From 7235fbf91eb072e21c385d8100ab48c81a0e9eaa Mon Sep 17 00:00:00 2001 From: Andrew Freiburger Date: Wed, 26 Aug 2026 14:49:02 -0500 Subject: [PATCH] Sync modelseedpy/core/rast_client.py from freiburgermsu fork File-level sync to the state of this file in freiburgermsu/ModelSEEDpy@971df5d, submitted as an individual PR so each module's divergence can be reviewed independently. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_017rcBA87xipNeaukuFW8N2G --- modelseedpy/core/rast_client.py | 64 +++++++++++++++++++++++++-------- 1 file changed, 50 insertions(+), 14 deletions(-) diff --git a/modelseedpy/core/rast_client.py b/modelseedpy/core/rast_client.py index cc36b7bb..84b59b4a 100644 --- a/modelseedpy/core/rast_client.py +++ b/modelseedpy/core/rast_client.py @@ -45,6 +45,13 @@ def aux_template(template): return rxn_roles +def split_annotation(annotation) -> list[str]: + if annotation: + functions = re.split("; | / | @", annotation) + return functions + return [] + + class RastClient: def __init__(self): self.rpc_client = RPCClient( @@ -67,18 +74,18 @@ def annotate_genome(self, genome, split_terms=True): p_features.append({"id": f.id, "protein_translation": f.seq}) res = self.f(p_features) - for o in res[0]["features"]: - feature = genome.features.get_by_id(o["id"]) - if "function" in o: - rast_function = o["function"] - if split_terms: - functions = re.split("; | / | @", rast_function) - for function in functions: - feature.add_ontology_term("RAST", function) - else: - feature.add_ontology_term("RAST", rast_function) + if res: + for o in res[0]["features"]: + feature = genome.features.get_by_id(o["id"]) + if "function" in o: + rast_function = o["function"] + if split_terms: + for function in split_annotation(rast_function): + feature.add_ontology_term("RAST", function) + else: + feature.add_ontology_term("RAST", rast_function) - return res[0]["analysis_events"] + return res[0]["analysis_events"] if res else None def annotate_genome_from_fasta(self, filepath, split="|"): genome = MSGenome.from_fasta(filepath, split) @@ -90,9 +97,38 @@ def annotate_protein_sequence(self, protein_id: str, protein_seq: str): p_features = [{"id": protein_id, "protein_translation": protein_seq}] return self.f(p_features) - def annotate_protein_sequences(self, protein_seqs: dict): - p_features = [{"id": protein_id, "protein_translation": protein_seq}] - return self.f(p_features) + def annotate_protein_batch(self, protein_batch: dict): + protein_annotations = {} + features = [ + {"id": _id, "protein_translation": seq} + for _id, seq in protein_batch.items() + ] + params = [{"features": features}, {"stages": self.stages}] + result = self.rpc_client.call("GenomeAnnotation.run_pipeline", params) + if result and len(result) > 0 and "features" in result[0]: + for feature in result[0]["features"]: + if "function" in feature: + protein_annotations[feature["id"]] = feature["function"] + + return protein_annotations + + def annotate_protein_sequences(self, protein_seqs: dict, chunk_size: int = 5000): + protein_annotations = {} + chunk = {} + + for protein_id, protein_seq in protein_seqs.items(): + chunk[protein_id] = protein_seq + + # Process chunk when it reaches the specified size + if len(chunk) >= chunk_size: + protein_annotations.update(self.annotate_protein_batch(chunk)) + chunk = {} + + # Process remaining proteins in the last chunk + if chunk: + protein_annotations.update(self.annotate_protein_batch(chunk)) + + return protein_annotations def f1(self, protein_id, protein_seq): p_features = [{"id": protein_id, "protein_translation": protein_seq}]