39 lines
1.6 KiB
Python
39 lines
1.6 KiB
Python
from collections import defaultdict
|
|
import csv
|
|
from os import PathLike
|
|
from typing import AsyncIterable, Collection, Mapping, Sequence, Union
|
|
|
|
from autobigs.engine.structures.mlst import Allele, MLSTProfile
|
|
|
|
|
|
def alleles_to_map(alleles: Collection[Allele]) -> Mapping[str, Union[list[str], str]]:
|
|
result = defaultdict(list)
|
|
for allele in alleles:
|
|
result[allele.allele_locus].append(allele.allele_variant)
|
|
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]:
|
|
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:
|
|
if mlst_profile is None:
|
|
failed.append(name)
|
|
continue
|
|
allele_mapping = alleles_to_map(mlst_profile.alleles)
|
|
if writer is None:
|
|
header = ["id", "st", "clonal-complex", *sorted(allele_mapping.keys())]
|
|
writer = csv.DictWriter(filehandle, fieldnames=header)
|
|
writer.writeheader()
|
|
row_dictionary = {
|
|
"st": mlst_profile.sequence_type,
|
|
"clonal-complex": mlst_profile.clonal_complex,
|
|
"id": name,
|
|
**allele_mapping
|
|
}
|
|
writer.writerow(rowdict=row_dictionary)
|
|
return failed |