Cleanup (#311)

chai_lab/data/parsing/fasta.py

@@ -5,7 +5,7 @@
 import logging
 from io import StringIO
 from pathlib import Path
-from typing import NamedTuple, Sequence
+from typing import NamedTuple, Sequence, TextIO
 
 from chai_lab.data.parsing.structure.entity_type import EntityType
 from chai_lab.data.residue_constants import restype_1to3_with_x
@@ -31,12 +31,15 @@ def fastas_to_str(fastas: Sequence[Fasta]) -> str:
     return "".join(f">{fasta.header}\n{fasta.sequence}\n" for fasta in fastas)
 
 
-def read_fasta(file_path: str | Path | StringIO) -> list[Fasta]:
+def read_fasta(file_path: str | Path) -> list[Fasta]:
+    with open(file_path) as source:
+        return read_fasta_content(source)
+
+
+def read_fasta_content(content: StringIO | TextIO) -> list[Fasta]:
     from Bio import SeqIO
 
-    fasta_sequences = SeqIO.parse(
-        open(file_path) if isinstance(file_path, (str, Path)) else file_path, "fasta"
-    )
+    fasta_sequences = SeqIO.parse(content, "fasta")
     return [Fasta(fasta.description, str(fasta.seq)) for fasta in fasta_sequences]
 
 

chai_lab/data/parsing/msas/a3m.py

@@ -19,7 +19,7 @@
 import numba
 import numpy as np
 
-from chai_lab.data.parsing.fasta import Fasta, read_fasta
+from chai_lab.data.parsing.fasta import Fasta, read_fasta_content
 from chai_lab.data.residue_constants import residue_types_with_nucleotides_order
 
 MAPPED_TOKEN_SKIP: Final[int] = -1
@@ -124,8 +124,10 @@ def read_colabfold_a3m(fname: Path) -> dict[str, list[Fasta]]:
         if not block:
             continue
         strio = StringIO(block)
-        hits = read_fasta(strio)
+        hits = read_fasta_content(strio)
         assert len(hits) > 0
-        assert re.match(r"^[0-9]{3}$", (query := hits[0].header))
+
+        query = hits[0].header
+        assert re.match(r"^[0-9]{3}$", query)
         retval[query] = hits
     return retval

scripts/stage_for_chai.py → scripts/stage_colabfold_outputs_for_chai.py

@@ -3,7 +3,8 @@
 # See the LICENSE file for details.
 """
 Given a output directory from a ColabFold run, traverses the directory structure and stage
-the same MSA and templates to run through Chai1.
+the same MSA and templates to run through Chai1. This is likely not applicable unless you are
+incorporating Chai1 into adhoc analyses based on ColabFold outputs.
 
 Some minimal example: