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

Added functionality to retry determining ST
Some checks reported errors
autoBIGS.engine/pipeline/head Something is wrong with the build of this commit
Details

Browse Source
This commit is contained in:
Harrison Deng 2025-03-13 16:13:14 +00:00
parent 8ffc7c7fb5
commit e088d1080b
1 changed files with 37 additions and 20 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

53
src/autobigs/engine/analysis/bigsdb.py
View File

@ -99,14 +99,17 @@ class RemoteBIGSdbMLSTProfiler(BIGSdbMLSTProfiler):
yield result_allele if isinstance(sequence_string, str) else (sequence_string.name, result_allele) yield result_allele if isinstance(sequence_string, str) else (sequence_string.name, result_allele)
else: else:
raise NoBIGSdbMatchesException(self._database_name, self._scheme_id, sequence_string.name if isinstance(sequence_string, NamedString) else None) raise NoBIGSdbMatchesException(self._database_name, self._scheme_id, sequence_string.name if isinstance(sequence_string, NamedString) else None)
except ConnectionResetError as e: except Exception as e:
last_error = e last_error = e
success = False success = False
await asyncio.sleep(5) # In case the connection issue is due to rate issues await asyncio.sleep(5) # In case the connection issue is due to rate issues
else: else:
success = True success = True
if not success and last_error is not None: if not success and last_error is not None:
raise last_error try:
raise last_error
except ConnectionResetError as e:
yield Allele("error", "error", None)
async def determine_mlst_st(self, alleles: Union[AsyncIterable[Union[Allele, tuple[str, Allele]]], Iterable[Union[Allele, tuple[str, Allele]]]]) -> Union[MLSTProfile, NamedMLSTProfile]: async def determine_mlst_st(self, alleles: Union[AsyncIterable[Union[Allele, tuple[str, Allele]]], Iterable[Union[Allele, tuple[str, Allele]]]]) -> Union[MLSTProfile, NamedMLSTProfile]:
uri_path = "designations" uri_path = "designations"
@ -130,22 +133,36 @@ class RemoteBIGSdbMLSTProfiler(BIGSdbMLSTProfiler):
"designations": allele_request_dict "designations": allele_request_dict
} }
async with self._http_client.post(uri_path, json=request_json) as response: attempts = 0
response_json: dict = await response.json() success = False
allele_set: Set[Allele] = set() last_error = None
response_json.setdefault("fields", dict()) while attempts < self._retry_limit and not success:
scheme_fields_returned: dict[str, str] = response_json["fields"] try:
scheme_fields_returned.setdefault("ST", "unknown") async with self._http_client.post(uri_path, json=request_json) as response:
scheme_fields_returned.setdefault("clonal_complex", "unknown") response_json: dict = await response.json()
scheme_exact_matches: dict = response_json["exact_matches"] allele_set: Set[Allele] = set()
for exact_match_locus, exact_match_alleles in scheme_exact_matches.items(): response_json.setdefault("fields", dict())
allele_set.add(Allele(exact_match_locus, exact_match_alleles[0]["allele_id"], None)) scheme_fields_returned: dict[str, str] = response_json["fields"]
if len(allele_set) == 0: scheme_fields_returned.setdefault("ST", "unknown")
raise ValueError("Passed in no alleles.") scheme_fields_returned.setdefault("clonal_complex", "unknown")
result_mlst_profile = MLSTProfile(allele_set, scheme_fields_returned["ST"], scheme_fields_returned["clonal_complex"]) scheme_exact_matches: dict = response_json["exact_matches"]
if len(names_list) > 0: for exact_match_locus, exact_match_alleles in scheme_exact_matches.items():
result_mlst_profile = NamedMLSTProfile(str(tuple(names_list)) if len(set(names_list)) > 1 else names_list[0], result_mlst_profile) allele_set.add(Allele(exact_match_locus, exact_match_alleles[0]["allele_id"], None))
return result_mlst_profile if len(allele_set) == 0:
raise ValueError("Passed in no alleles.")
result_mlst_profile = MLSTProfile(allele_set, scheme_fields_returned["ST"], scheme_fields_returned["clonal_complex"])
if len(names_list) > 0:
result_mlst_profile = NamedMLSTProfile(str(tuple(names_list)) if len(set(names_list)) > 1 else names_list[0], result_mlst_profile)
return result_mlst_profile
except Exception as e:
last_error = e
success = False
await asyncio.sleep(5)
if not success and last_error is not None:
try:
raise last_error
except ConnectionResetError as e:
result_mlst_profile = NamedMLSTProfile((str(tuple(names_list)) if len(set(names_list)) > 1 else names_list[0]) + ":Error", None)
async def profile_string(self, query_sequence_strings: Iterable[Union[NamedString, str]]) -> Union[NamedMLSTProfile, MLSTProfile]: async def profile_string(self, query_sequence_strings: Iterable[Union[NamedString, str]]) -> Union[NamedMLSTProfile, MLSTProfile]:
alleles = self.determine_mlst_allele_variants(query_sequence_strings) alleles = self.determine_mlst_allele_variants(query_sequence_strings)