From 7384895578acc0a74e68b896978b70751ea2f827 Mon Sep 17 00:00:00 2001 From: Harrison Deng Date: Tue, 18 Feb 2025 16:03:17 +0000 Subject: [PATCH] Writing now uses named MLST profile --- src/autobigs/engine/writing.py | 8 +++++--- tests/autobigs/engine/test_writing.py | 14 +++++++------- 2 files changed, 12 insertions(+), 10 deletions(-) diff --git a/src/autobigs/engine/writing.py b/src/autobigs/engine/writing.py index 1f17ccd..1d19ec6 100644 --- a/src/autobigs/engine/writing.py +++ b/src/autobigs/engine/writing.py @@ -3,7 +3,7 @@ import csv from os import PathLike from typing import AsyncIterable, Collection, Mapping, Sequence, Union -from autobigs.engine.structures.mlst import Allele, MLSTProfile +from autobigs.engine.structures.mlst import Allele, MLSTProfile, NamedMLSTProfile def alleles_to_text_map(alleles: Collection[Allele]) -> Mapping[str, Union[Sequence[str], str]]: @@ -17,12 +17,14 @@ def alleles_to_text_map(alleles: Collection[Allele]) -> Mapping[str, Union[Seque result[locus] = tuple(result[locus]) # type: ignore return dict(result) -async def write_mlst_profiles_as_csv(mlst_profiles_iterable: AsyncIterable[tuple[str, Union[MLSTProfile, None]]], handle: Union[str, bytes, PathLike[str], PathLike[bytes]]) -> Sequence[str]: +async def write_mlst_profiles_as_csv(mlst_profiles_iterable: AsyncIterable[NamedMLSTProfile], handle: Union[str, bytes, PathLike[str], PathLike[bytes]]) -> Sequence[str]: failed = list() with open(handle, "w", newline='') as filehandle: header = None writer: Union[csv.DictWriter, None] = None - async for name, mlst_profile in mlst_profiles_iterable: + async for named_mlst_profile in mlst_profiles_iterable: + name = named_mlst_profile.name + mlst_profile = named_mlst_profile.mlst_profile if mlst_profile is None: failed.append(name) continue diff --git a/tests/autobigs/engine/test_writing.py b/tests/autobigs/engine/test_writing.py index 8c19620..cbd120b 100644 --- a/tests/autobigs/engine/test_writing.py +++ b/tests/autobigs/engine/test_writing.py @@ -3,7 +3,7 @@ from typing import AsyncIterable, Iterable import pytest from autobigs.engine.structures.alignment import AlignmentStats from autobigs.engine.writing import alleles_to_text_map, write_mlst_profiles_as_csv -from autobigs.engine.structures.mlst import Allele, MLSTProfile +from autobigs.engine.structures.mlst import Allele, MLSTProfile, NamedMLSTProfile import tempfile from csv import reader from os import path @@ -11,20 +11,20 @@ from os import path @pytest.fixture def dummy_alphabet_mlst_profile(): - return MLSTProfile(( + return NamedMLSTProfile("name", MLSTProfile(( Allele("A", "1", None), Allele("D", "1", None), Allele("B", "1", None), Allele("C", "1", None), Allele("C", "2", AlignmentStats(90, 10, 0, 90)) - ), "mysterious", "very mysterious") + ), "mysterious", "very mysterious")) async def iterable_to_asynciterable(iterable: Iterable): for iterated in iterable: yield iterated async def test_column_order_is_same_as_expected_file(dummy_alphabet_mlst_profile: MLSTProfile): - dummy_profiles = [("test_1", dummy_alphabet_mlst_profile)] + dummy_profiles = [dummy_alphabet_mlst_profile] with tempfile.TemporaryDirectory() as temp_dir: output_path = path.join(temp_dir, "out.csv") await write_mlst_profiles_as_csv(iterable_to_asynciterable(dummy_profiles), output_path) @@ -34,8 +34,8 @@ async def test_column_order_is_same_as_expected_file(dummy_alphabet_mlst_profile target_columns = lines[4:] assert target_columns == sorted(target_columns) -async def test_alleles_to_text_map_mapping_is_correct(dummy_alphabet_mlst_profile: MLSTProfile): - mapping = alleles_to_text_map(dummy_alphabet_mlst_profile.alleles) +async def test_alleles_to_text_map_mapping_is_correct(dummy_alphabet_mlst_profile: NamedMLSTProfile): + mapping = alleles_to_text_map(dummy_alphabet_mlst_profile.mlst_profile.alleles) # type: ignore expected_mapping = { "A": "1", "B": "1", @@ -44,4 +44,4 @@ async def test_alleles_to_text_map_mapping_is_correct(dummy_alphabet_mlst_profil } for allele_name, allele_ids in mapping.items(): assert allele_name in expected_mapping - assert allele_ids == expected_mapping[allele_name] \ No newline at end of file + assert allele_ids == expected_mapping[allele_name]