refactor(models,io): Simplify helpers and make models/io mypy-clean

pyproject.toml

@@ -236,17 +236,4 @@ module = [
   "q2mm.backends.mm.jax_md_engine",
   "q2mm.backends.mm.openmm",
   "q2mm.backends.mm.tinker",
-  "q2mm.io._helpers",
-  "q2mm.io.amber",
-  "q2mm.io.cmap",
-  "q2mm.io.gaussian",
-  "q2mm.io.jaguar",
-  "q2mm.io.macromodel",
-  "q2mm.io.mm3",
-  "q2mm.io.openmm",
-  "q2mm.io.tinker",
-  "q2mm.models.forcefield",
-  "q2mm.models.hessian",
-  "q2mm.models.molecule",
-  "q2mm.models.seminario",
 ]

q2mm/io/_helpers.py

@@ -5,8 +5,10 @@
 
 from __future__ import annotations
 
+import os
 from collections.abc import Callable, Sequence
-from typing import TYPE_CHECKING, Any
+from pathlib import Path
+from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
     from q2mm.models.forcefield import (
@@ -23,8 +25,36 @@
     canonicalize_bond_env_id,
 )
 
+
+class _TextFileBase:
+    """Shared path and lazy-line-reading base for legacy text-file parsers."""
+
+    __slots__ = ("_lines", "path", "directory", "filename")
+
+    def __init__(self, path: str) -> None:
+        self._lines: list[str] | None = None
+        self.path = os.path.abspath(path)
+        self.directory = os.path.dirname(self.path)
+        self.filename = os.path.basename(self.path)
+
+    @property
+    def lines(self) -> list[str]:
+        """Return the lines of the file."""
+        if self._lines is None:
+            with open(self.path) as f:
+                self._lines = f.readlines()
+        return self._lines
+
+    def write(self, path: str, lines: list[str] | None = None) -> None:
+        """Write lines to file at *path* (defaults to ``self.lines``)."""
+        if lines is None:
+            lines = self.lines
+        with open(path, "w") as f:
+            f.writelines(lines)
+
+
 # ---------------------------------------------------------------------------
-# Format compatibility
+# Format compatibility & template resolution
 # ---------------------------------------------------------------------------
 
 _FORMAT_COMPATIBLE_FORMS: dict[str, set[str]] = {
@@ -46,8 +76,21 @@ def _validate_form_for_format(ff: ForceField, target_format: str) -> None:
         )
 
 
+def _resolve_template_path(
+    ff: ForceField,
+    template_path: str | Path | None,
+    source_format: str,
+) -> Path | None:
+    """Resolve an explicit *template_path* or fall back to ``ff.source_path`` when formats match."""
+    if template_path is not None:
+        return Path(template_path)
+    if ff.source_format == source_format and ff.source_path is not None:
+        return ff.source_path
+    return None
+
+
 # ---------------------------------------------------------------------------
-# Env-id splitting
+# Env-id splitting & atom-type cleaning
 # ---------------------------------------------------------------------------
 
 
@@ -58,12 +101,7 @@ def _split_env_id(env_id: str, expected_len: int) -> list[str]:
     return []
 
 
-# ---------------------------------------------------------------------------
-# Atom-type cleaning
-# ---------------------------------------------------------------------------
-
-
-def _clean_atom_types(atom_types: list[str] | tuple[str, ...] | None, expected_len: int) -> list[str]:
+def _clean_atom_types(atom_types: Sequence[str] | None, expected_len: int) -> list[str]:
     if atom_types is None:
         return []
     if isinstance(atom_types, str):
@@ -79,28 +117,30 @@ def _clean_atom_types(atom_types: list[str] | tuple[str, ...] | None, expected_l
 # ---------------------------------------------------------------------------
 
 
-def _build_param_maps(params: list, secondary_key: str) -> tuple[dict, dict]:
-    """Build ff_row and secondary-key lookup dicts for a list of parameters."""
+def _build_param_maps[P: (BondParam, AngleParam, StretchBendParam, VdwParam)](
+    params: Sequence[P], secondary_key: str
+) -> tuple[dict[int, P], dict[str, P]]:
+    """Build ff_row and secondary-key lookup dicts for a sequence of parameters."""
     by_row = {p.ff_row: p for p in params if p.ff_row is not None}
     by_key = {getattr(p, secondary_key): p for p in params if getattr(p, secondary_key, None)}
     return by_row, by_key
 
 
-def _build_bond_maps(bonds: list[BondParam]) -> tuple[dict[int, BondParam], dict[str, BondParam]]:
+def _build_bond_maps(bonds: Sequence[BondParam]) -> tuple[dict[int, BondParam], dict[str, BondParam]]:
     return _build_param_maps(bonds, "env_id")
 
 
-def _build_angle_maps(angles: list[AngleParam]) -> tuple[dict[int, AngleParam], dict[str, AngleParam]]:
+def _build_angle_maps(angles: Sequence[AngleParam]) -> tuple[dict[int, AngleParam], dict[str, AngleParam]]:
     return _build_param_maps(angles, "env_id")
 
 
 def _build_sb_maps(
-    stretch_bends: list[StretchBendParam],
+    stretch_bends: Sequence[StretchBendParam],
 ) -> tuple[dict[int, StretchBendParam], dict[str, StretchBendParam]]:
     return _build_param_maps(stretch_bends, "env_id")
 
 
-def _build_vdw_maps(vdws: list[VdwParam]) -> tuple[dict[int, VdwParam], dict[str, VdwParam]]:
+def _build_vdw_maps(vdws: Sequence[VdwParam]) -> tuple[dict[int, VdwParam], dict[str, VdwParam]]:
     return _build_param_maps(vdws, "atom_type")
 
 
@@ -109,14 +149,14 @@ def _build_vdw_maps(vdws: list[VdwParam]) -> tuple[dict[int, VdwParam], dict[str
 # ---------------------------------------------------------------------------
 
 
-def _match_for_export(
+def _match_for_export[P: (BondParam, AngleParam, StretchBendParam)](
     ff_row: int | None,
-    atom_types: list[str] | tuple[str, ...] | None,
-    by_row: dict,
-    by_env: dict,
+    atom_types: Sequence[str] | None,
+    by_row: dict[int, P],
+    by_env: dict[str, P],
     expected_len: int,
-    canonicalize_fn: Callable,
-) -> Any:
+    canonicalize_fn: Callable[[Sequence[str]], str],
+) -> P | None:
     """Match a parsed file row to an internal param by ``ff_row`` or ``env_id``.
 
     Pure: takes the row's primitive identity data (*ff_row*, *atom_types*)
@@ -133,7 +173,7 @@ def _match_for_export(
 
 def _match_bond_for_export(
     ff_row: int | None,
-    atom_types: list[str] | tuple[str, ...] | None,
+    atom_types: Sequence[str] | None,
     bond_by_row: dict[int, BondParam],
     bond_by_env: dict[str, BondParam],
 ) -> BondParam | None:
@@ -142,7 +182,7 @@ def _match_bond_for_export(
 
 def _match_angle_for_export(
     ff_row: int | None,
-    atom_types: list[str] | tuple[str, ...] | None,
+    atom_types: Sequence[str] | None,
     angle_by_row: dict[int, AngleParam],
     angle_by_env: dict[str, AngleParam],
 ) -> AngleParam | None:
@@ -151,7 +191,7 @@ def _match_angle_for_export(
 
 def _match_sb_for_export(
     ff_row: int | None,
-    atom_types: list[str] | tuple[str, ...] | None,
+    atom_types: Sequence[str] | None,
     sb_by_row: dict[int, StretchBendParam],
     sb_by_env: dict[str, StretchBendParam],
 ) -> StretchBendParam | None:

q2mm/io/amber.py

@@ -9,7 +9,7 @@
 from decimal import Decimal
 from pathlib import Path
 
-from q2mm.io._helpers import _validate_form_for_format
+from q2mm.io._helpers import _resolve_template_path, _validate_form_for_format
 from q2mm.models.forcefield import (
     AngleParam,
     BondParam,
@@ -308,10 +308,10 @@ def load_amber_frcmod(path: str | Path) -> ForceField:
             types, rest = _parse_amber_types(line, 2)
             vals = _parse_floats(rest)
             if len(types) == 2 and all(types) and len(vals) >= 2:
-                elems = tuple(_amber_type_to_element(t, mass_map) for t in types)
+                t0, t1 = types
                 bonds.append(
                     BondParam(
-                        elements=elems,
+                        elements=(_amber_type_to_element(t0, mass_map), _amber_type_to_element(t1, mass_map)),
                         equilibrium=vals[1],
                         force_constant=vals[0],
                         env_id=canonicalize_bond_env_id(types),
@@ -324,10 +324,14 @@ def load_amber_frcmod(path: str | Path) -> ForceField:
             types, rest = _parse_amber_types(line, 3)
             vals = _parse_floats(rest)
             if len(types) == 3 and all(types) and len(vals) >= 2:
-                elems = tuple(_amber_type_to_element(t, mass_map) for t in types)
+                t0, t1, t2 = types
                 angles.append(
                     AngleParam(
-                        elements=elems,
+                        elements=(
+                            _amber_type_to_element(t0, mass_map),
+                            _amber_type_to_element(t1, mass_map),
+                            _amber_type_to_element(t2, mass_map),
+                        ),
                         equilibrium=vals[1],
                         force_constant=vals[0],
                         env_id=canonicalize_angle_env_id(types),
@@ -338,10 +342,15 @@ def load_amber_frcmod(path: str | Path) -> ForceField:
 
         elif section == "DIHE":
             record = dihe_rows[row]
-            elems = tuple(_amber_type_to_element(t, mass_map) for t in record.atom_types)
+            t0, t1, t2, t3 = record.atom_types
             torsions.append(
                 TorsionParam(
-                    elements=elems,
+                    elements=(
+                        _amber_type_to_element(t0, mass_map),
+                        _amber_type_to_element(t1, mass_map),
+                        _amber_type_to_element(t2, mass_map),
+                        _amber_type_to_element(t3, mass_map),
+                    ),
                     periodicity=abs(record.pn),
                     force_constant=record.barrier / record.effective_idivf,
                     phase=record.phase,
@@ -356,15 +365,20 @@ def load_amber_frcmod(path: str | Path) -> ForceField:
             vals = _parse_floats(rest)
             # vals: barrier, phase, periodicity (no IDIVF)
             if len(types) == 4 and all(types) and len(vals) >= 3:
-                elems = tuple(_amber_type_to_element(t, mass_map) for t in types)
+                t0, t1, t2, t3 = types
                 pn = vals[2]
                 # IMPROPER has neither DIHE's continuation sign nor a zero-to-one convention.
                 if not math.isfinite(pn) or pn <= 0 or pn != int(pn):
                     raise ValueError(f"AMBER IMPROPER row {row}: raw PN must be a finite positive integer")
                 periodicity = int(pn)
                 torsions.append(
                     TorsionParam(
-                        elements=elems,
+                        elements=(
+                            _amber_type_to_element(t0, mass_map),
+                            _amber_type_to_element(t1, mass_map),
+                            _amber_type_to_element(t2, mass_map),
+                            _amber_type_to_element(t3, mass_map),
+                        ),
                         periodicity=periodicity,
                         force_constant=vals[0],
                         phase=vals[1],
@@ -397,10 +411,10 @@ def load_amber_frcmod(path: str | Path) -> ForceField:
 
     return ForceField(
         name=f"AMBER from {path.name}",
-        bonds=bonds,
-        angles=angles,
-        torsions=torsions,
-        vdws=vdws,
+        bonds=tuple(bonds),
+        angles=tuple(angles),
+        torsions=tuple(torsions),
+        vdws=tuple(vdws),
         source_path=path,
         source_format="amber_frcmod",
         functional_form=FunctionalForm.HARMONIC,
@@ -453,9 +467,7 @@ def save_amber_frcmod(
     """
     _validate_form_for_format(ff, "amber_frcmod")
     output_path = Path(path)
-    template = Path(template_path) if template_path is not None else None
-    if template is None and ff.source_format == "amber_frcmod" and ff.source_path is not None:
-        template = ff.source_path
+    template = _resolve_template_path(ff, template_path, "amber_frcmod")
 
     if template is not None:
         return _save_amber_frcmod_template(ff, output_path, template)