ydeng/autoBIGS.engine
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Adapted latest merged reading codebase to current codebase

Browse Source
This commit is contained in:
Harrison Deng 2025-02-11 19:13:29 +00:00
parent 4b9eb8674d
commit e568e9fb2c
1 changed files with 13 additions and 9 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

22
src/autobigs/engine/writing.py
View File

@ -1,16 +1,19 @@
from collections import defaultdict
import csv import csv
from os import PathLike from os import PathLike
from typing import AsyncIterable, Mapping, Sequence, Union from typing import AsyncIterable, Collection, Mapping, Sequence, Union
from autobigs.engine.structures.mlst import Allele, MLSTProfile from autobigs.engine.structures.mlst import Allele, MLSTProfile
def dict_loci_alleles_variants_from_loci(alleles_map: Mapping[str, Allele]): def alleles_to_map(alleles: Collection[Allele]) -> Mapping[str, Union[list[str], str]]:
result_dict: dict[str, Union[list[str], str]] = {} result = defaultdict(list)
for loci, alleles in alleles_map.items(): for allele in alleles:
result_dict[loci] = alleles.allele_variant result[allele.allele_locus].append(allele.allele_variant)
return result_dict for locus in result.keys():
if len(result[locus]) == 1:
result[locus] = result[locus][0] # Take the only one
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[tuple[str, Union[MLSTProfile, None]]], handle: Union[str, bytes, PathLike[str], PathLike[bytes]]) -> Sequence[str]:
failed = list() failed = list()
@ -21,15 +24,16 @@ async def write_mlst_profiles_as_csv(mlst_profiles_iterable: AsyncIterable[tuple
if mlst_profile is None: if mlst_profile is None:
failed.append(name) failed.append(name)
continue continue
allele_mapping = alleles_to_map(mlst_profile.alleles)
if writer is None: if writer is None:
header = ["id", "st", "clonal-complex", *sorted(mlst_profile.alleles.keys())] header = ["id", "st", "clonal-complex", *sorted(allele_mapping.keys())]
writer = csv.DictWriter(filehandle, fieldnames=header) writer = csv.DictWriter(filehandle, fieldnames=header)
writer.writeheader() writer.writeheader()
row_dictionary = { row_dictionary = {
"st": mlst_profile.sequence_type, "st": mlst_profile.sequence_type,
"clonal-complex": mlst_profile.clonal_complex, "clonal-complex": mlst_profile.clonal_complex,
"id": name, "id": name,
**mlst_profile.alleles **allele_mapping
} }
writer.writerow(rowdict=row_dictionary) writer.writerow(rowdict=row_dictionary)
return failed return failed