Annotate most everything (openforcefield#1798)

* Annotate everything

* Shift more annotations

* Fix more typing errors

* Remove more (most?) annotations from docstrings

* Remove debug code

* Update

* Apply suggestions from code review

Co-authored-by: Josh A. Mitchell <[email protected]>

* Updates from code review

* Updates/fix

* Fix

* Update openff/toolkit/topology/topology.py

Co-authored-by: Josh A. Mitchell <[email protected]>

* Fix

* Fix

* Avoid name clash

* Update release history

---------

Co-authored-by: Josh A. Mitchell <[email protected]>

.github/workflows/CI.yml

@@ -143,7 +143,8 @@ jobs:
         run: pytest -r fE --tb=short openff/toolkit/_tests/test_links.py
 
       - name: Run mypy
-        if: ${{ matrix.rdkit == true && matrix.openeye == true }}
+        # subtle differences in Python/mypy versions, just keep it passing on one
+        if: ${{ matrix.rdkit == true && matrix.openeye && matrix.python-version == 3.11 }}
         run: mypy -p "openff.toolkit"
 
       - name: Run unit tests

.pre-commit-config.yaml

@@ -3,7 +3,7 @@ ci:
 files: ^openff|(^examples/((?!deprecated).)*$)|^docs
 repos:
 - repo: https://github.com/psf/black
-  rev: 24.1.1
+  rev: 24.2.0
   hooks:
   - id: black
   - id: black-jupyter

devtools/conda-envs/test_env.yaml

@@ -44,7 +44,7 @@ dependencies:
   - nbval
   # No idea why this is necessary, see https://github.com/openforcefield/openff-toolkit/pull/1821
   - nomkl
-  - mypy
+  - mypy =1.8
   - typing_extensions
   - pydantic =1
   - pip:

docs/releasehistory.md

@@ -8,6 +8,8 @@ Releases follow the `major.minor.micro` scheme recommended by [PEP440](https://w
 
 ## Current development
 
+* #1798 Adds type annotations to most of the codebase.
+
 ### API-breaking changes
 
 ### Behavior changes

openff/toolkit/topology/_mm_molecule.py

@@ -10,7 +10,7 @@
 
 """
 
-from typing import TYPE_CHECKING, Generator, NoReturn, Optional, Union
+from typing import TYPE_CHECKING, Generator, Iterable, NoReturn, Optional, Union
 
 from openff.units.elements import MASSES, SYMBOLS
 
@@ -37,7 +37,7 @@ class _SimpleMolecule:
     def __init__(self):
         self.atoms = []
         self.bonds = []
-        self.conformers = None
+        self._conformers = None
         self._hierarchy_schemes = dict()
 
     def add_atom(self, atomic_number: int, **kwargs):
@@ -61,10 +61,14 @@ def add_bond(self, atom1, atom2, **kwargs):
         bond = _SimpleBond(atom1_atom, atom2_atom, **kwargs)
         self.bonds.append(bond)
 
+    @property
+    def conformers(self) -> Optional[list["Quantity"]]:
+        return self._conformers
+
     def add_conformer(self, conformer):
-        if self.conformers is None:
-            self.conformers = list()
-        self.conformers.append(conformer)
+        if self._conformers is None:
+            self._conformers = list()
+        self._conformers.append(conformer)
 
     @property
     def n_atoms(self) -> int:
@@ -76,9 +80,7 @@ def n_bonds(self) -> int:
 
     @property
     def n_conformers(self) -> int:
-        if self.conformers is None:
-            return 0
-        return len(self.conformers)
+        return 0 if self._conformers is None else len(self._conformers)
 
     def atom(self, index):
         return self.atoms[index]
@@ -305,6 +307,15 @@ def _from_subgraph(cls, subgraph: "nx.Graph"):
 
         return molecule
 
+    # TODO: Implement me - shouldn't need to be too different than Molecule.add_hierarchy_scheme
+    #       since the extra chemical information is unrelated to hierarchy/metadata
+    def add_hierarchy_scheme(
+        self,
+        uniqueness_criteria: Iterable[str],
+        iterator_name: str,
+    ) -> NoReturn:
+        raise NotImplementedError()
+
     def to_dict(self) -> dict:
         molecule_dict = dict()
         special_serialization_logic = [
@@ -325,14 +336,14 @@ def to_dict(self) -> dict:
 
         molecule_dict["bonds"] = [bond.to_dict() for bond in self.bonds]
 
-        if self.conformers is None:
+        if self._conformers is None:
             molecule_dict["conformers"] = None
         else:
             molecule_dict["conformers"] = []
             molecule_dict["conformers_unit"] = (
                 "angstrom"  # Have this defined as a class variable?
             )
-            for conf in self.conformers:
+            for conf in self._conformers:
                 conf_unitless = conf.m_as(unit.angstrom)
                 conf_serialized, conf_shape = serialize_numpy((conf_unitless))
                 molecule_dict["conformers"].append(conf_serialized)
@@ -363,15 +374,15 @@ def from_dict(cls, molecule_dict):
 
         conformers = molecule_dict.pop("conformers")
         if conformers is None:
-            molecule.conformers = None
+            molecule._conformers = None
         else:
             conformers_unit = molecule_dict.pop("conformers_unit")
-            molecule.conformers = list()
+            molecule._conformers = list()
             for ser_conf in conformers:
                 conformers_shape = (molecule.n_atoms, 3)
                 conformer_unitless = deserialize_numpy(ser_conf, conformers_shape)
                 conformer = unit.Quantity(conformer_unitless, conformers_unit)
-                molecule.conformers.append(conformer)
+                molecule._conformers.append(conformer)
 
         hier_scheme_dicts = molecule_dict.pop("hierarchy_schemes")
         for iterator_name, hierarchy_scheme_dict in hier_scheme_dicts.items():
@@ -387,7 +398,8 @@ def from_dict(cls, molecule_dict):
                     atom_indices=element_dict["atom_indices"],
                 )
 
-        {setattr(molecule, key, val) for key, val in molecule_dict.items()}
+        for key, val in molecule_dict.items():
+            setattr(molecule, key, val)
 
         return molecule
 
@@ -407,11 +419,13 @@ def from_molecule(cls, molecule: Molecule):
                 atom2=mm_molecule.atom(bond.atom2_index),
             )
 
-        mm_molecule.conformers = molecule.conformers
+        if molecule.conformers is not None:
+            for conformer in molecule.conformers:
+                mm_molecule.add_conformer(conformer)
 
         for name, hierarchy_scheme in molecule.hierarchy_schemes.items():
             assert name == hierarchy_scheme.iterator_name
-            mm_molecule.add_hierarchy_scheme(  # type: ignore[operator]
+            mm_molecule.add_hierarchy_scheme(
                 uniqueness_criteria=hierarchy_scheme.uniqueness_criteria,
                 iterator_name=hierarchy_scheme.iterator_name,
             )
@@ -503,7 +517,7 @@ def generate_unique_atom_names(self):
         """Generate unique atom names. See `Molecule.generate_unique_atom_names`."""
         from collections import defaultdict
 
-        element_counts = defaultdict(int)
+        element_counts: defaultdict[str, int] = defaultdict(int)
         for atom in self.atoms:
             symbol = SYMBOLS[atom.atomic_number]
             element_counts[symbol] += 1
@@ -528,7 +542,7 @@ def __deepcopy__(self, memo):
         return self.__class__.from_dict(self.to_dict())
 
 
-_SimpleMolecule.add_hierarchy_scheme = Molecule.add_hierarchy_scheme  # type: ignore[attr-defined]
+_SimpleMolecule.add_hierarchy_scheme = Molecule.add_hierarchy_scheme  # type: ignore[assignment]
 _SimpleMolecule.update_hierarchy_schemes = Molecule.update_hierarchy_schemes  # type: ignore[attr-defined]
 
 
@@ -550,7 +564,7 @@ def __init__(
         self._name = name
         self._atomic_number = atomic_number
         self._molecule = molecule
-        self._bonds: list[Optional[_SimpleBond]] = list()
+        self._bonds: list[_SimpleBond] = list()
         for key, val in kwargs.items():
             setattr(self, key, val)