Add type hints (#450)

mypy.ini

@@ -32,3 +32,4 @@ ignore_missing_imports = True
 ignore_missing_imports = True
 
 [mypy]
+show_error_codes = True

src/stk/serialization/json/serializers/molecule/molecule.py

@@ -4,135 +4,162 @@
 
 """
 
-from stk.molecular import InchiKey
+from __future__ import annotations
 
-from .utilities import atom_to_json, bond_to_json
+import typing
+
+from stk.molecular import InchiKey, Molecule, MoleculeKeyMaker
+from stk.utilities import OneOrMany
+
+from .utilities import AtomJson, BondJson, atom_to_json, bond_to_json
+
+
+class _MolecularGraphJson(typing.TypedDict):
+    a: tuple[AtomJson, ...]
+    b: tuple[BondJson, ...]
+
+
+class _PositionMatrixJson(typing.TypedDict):
+    m: list[list[float]]
+
+
+class _MoleculeJson(typing.TypedDict):
+    molecule: _MolecularGraphJson
+    matrix: _PositionMatrixJson
 
 
 class MoleculeJsonizer:
     """
     Abstract base class for creating JSONs of molecules.
 
-    See Also
-    --------
-    :class:`.ConstructedMoleculeJsonizer`
+    See Also:
 
-    Notes
-    -----
-    You might notice that the public methods of this abstract base
-    class are implemented. These are just default implementations,
-    which can be safely ignored or overridden, when implementing
-    subclasses. However, the default implementation can be used
-    directly, if it suits your needs.
+        :class:`.ConstructedMoleculeJsonizer`
 
-    Examples
-    --------
-    *Converting a Molecule to JSON*
+    Notes:
 
-    You want to create a JSON representation of a molecule
+        You might notice that the public methods of this abstract base
+        class are implemented. These are just default implementations,
+        which can be safely ignored or overridden, when implementing
+        subclasses. However, the default implementation can be used
+        directly, if it suits your needs.
 
-    .. testcode:: converting-a-molecule-to-json
+    Examples:
 
-        import stk
+        *Converting a Molecule to JSON*
 
-        jsonizer = stk.MoleculeJsonizer()
-        json = jsonizer.to_json(stk.BuildingBlock('NCCN'))
+        You want to create a JSON representation of a molecule
 
-    *Adding Additional Molecular Keys*
+        .. testcode:: converting-a-molecule-to-json
 
-    Apart from atoms, bonds and the position matrix, the JSON
-    representation holds additional fields, one for each
-    :class:`.MoleculeKeyMaker` provided to the initializer
+            import stk
 
-    .. testcode:: adding-additional-molecular-keys
+            jsonizer = stk.MoleculeJsonizer()
+            json = jsonizer.to_json(stk.BuildingBlock('NCCN'))
 
-        import stk
+        *Adding Additional Molecular Keys*
 
-        jsonizer = stk.MoleculeJsonizer(
-            key_makers=(
-                stk.Inchi(),
-                stk.InchiKey(),
-            ),
-        )
-        json = jsonizer.to_json(stk.BuildingBlock('NCCN'))
+        Apart from atoms, bonds and the position matrix, the JSON
+        representation holds additional fields, one for each
+        :class:`.MoleculeKeyMaker` provided to the initializer
 
-    In this case, ``json`` will look something like
+        .. testcode:: adding-additional-molecular-keys
 
-    .. code-block:: python
+            import stk
 
-        {
-            # A tuple of JSON atom representations.
-            'atoms': (...),
+            jsonizer = stk.MoleculeJsonizer(
+                key_makers=(
+                    stk.Inchi(),
+                    stk.InchiKey(),
+                ),
+            )
+            json = jsonizer.to_json(stk.BuildingBlock('NCCN'))
 
-            # A tuple of JSON bond representations.
-            'bonds': (...),
+        In this case, ``json`` will look something like
 
-            'InChI': 'The InChI of the molecule',
-            'InChIKey': 'The InChIKey of the molecule',
-        }
+        .. code-block:: python
 
-    For every :class:`.MoleculeKeyMaker` provided to `key_makers`,
-    a new key will be added to the JSON representation, with its name
-    given by :meth:`.MoleculeKeyMaker.get_key_name` and the value
-    given by :meth:`.MoleculeKeyMaker.get_key`.
+            {
+                # A tuple of JSON atom representations.
+                'atoms': (...),
+
+                # A tuple of JSON bond representations.
+                'bonds': (...),
+
+                'InChI': 'The InChI of the molecule',
+                'InChIKey': 'The InChIKey of the molecule',
+            }
+
+        For every :class:`.MoleculeKeyMaker` provided to `key_makers`,
+        a new key will be added to the JSON representation, with its
+        name given by :meth:`.MoleculeKeyMaker.get_key_name` and the
+        value given by :meth:`.MoleculeKeyMaker.get_key`.
 
     """
 
     def __init__(
         self,
-        key_makers=(InchiKey(), ),
-    ):
+        key_makers: OneOrMany[MoleculeKeyMaker] = (InchiKey(), ),
+    ) -> None:
         """
         Initialize a :class:`.MoleculeJsonizer` instance.
 
-        Parameters
-        ----------
-        key_makers : :class:`tuple` of \
-                :class:`.MoleculeKeyMaker`
-            Used to make the keys of molecules, which should be
-            included in their JSON representations. Keys allow
-            molecular data to reference itself when split across
-            multiple JSONs.
+        Parameters:
+
+            key_makers:
+                Used to make the keys of molecules, which should be
+                included in their JSON representations. Keys allow
+                molecular data to reference itself when split across
+                multiple JSONs.
 
         """
 
-        self._key_makers = key_makers
+        if isinstance(key_makers, MoleculeKeyMaker):
+            key_makers = (key_makers, )
 
-    def to_json(self, molecule):
+        self._key_makers = tuple(key_makers)
+
+    def to_json(
+        self,
+        molecule: Molecule,
+    ) -> _MoleculeJson:
         """
         Return a JSON representation of `molecule`.
 
-        Parameters
-        ----------
-        molecule : :class:`.Molecule`
-            The molecule to serialize.
+        Parameters:
+
+            molecule:
+                The molecule to serialize.
+
+        Returns:
 
-        Returns
-        -------
-        :class:`dict`
             A JSON representation of `molecule`.
 
         """
 
-        json = {
+        json: _MolecularGraphJson = {
             'a': tuple(map(atom_to_json, molecule.get_atoms())),
             'b': tuple(map(bond_to_json, molecule.get_bonds())),
         }
-        position_matrix = {
+        position_matrix: _PositionMatrixJson = {
             'm': molecule.get_position_matrix().tolist(),
         }
         for key_maker in self._key_makers:
             key_name = key_maker.get_key_name()
             key = key_maker.get_key(molecule)
-            json[key_name] = key
-            position_matrix[key_name] = key
+            # TypedDict does not allow keys which are not listed in the
+            # class definition. However, because we wish to create
+            # indices on keys created by KeyMakers we need to add
+            # these additional key-value pairs.
+            json[key_name] = key  # type: ignore
+            position_matrix[key_name] = key  # type: ignore
         return {
             'molecule': json,
             'matrix': position_matrix,
         }
 
-    def __str__(self):
+    def __str__(self) -> str:
         return repr(self)
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         return f'{self.__class__.__name__}({self._key_makers!r})'

src/stk/serialization/json/serializers/molecule/utilities.py

@@ -4,19 +4,40 @@
 
 """
 
+from __future__ import annotations
 
-def atom_to_json(atom):
+from .....molecular import Atom, Bond
+
+AtomicNumber = int
+AtomicCharge = int
+AtomJson = tuple[
+    AtomicNumber,
+    AtomicCharge,
+]
+AtomId = int
+BondOrder = int
+BondPeriodicity = tuple[int, int, int]
+BondJson = tuple[
+    AtomId,
+    AtomId,
+    BondOrder,
+    BondPeriodicity,
+]
+
+
+def atom_to_json(
+    atom: Atom,
+) -> AtomJson:
     """
     Return a JSON representation of `atom`.
 
-    Parameters
-    ----------
-    atom : :class:`.Atom`
-        The atom to serialize.
+    Parameters:
+
+        atom:
+            The atom to serialize.
+
+    Returns:
 
-    Returns
-    -------
-    :class:`dict`
         A JSON representation of `atom`.
 
     """
@@ -27,18 +48,19 @@ def atom_to_json(atom):
     )
 
 
-def bond_to_json(bond):
+def bond_to_json(
+    bond: Bond,
+) -> BondJson:
     """
     Return a JSON representation of `bond`.
 
-    Parameters
-    ----------
-    bond : :class:`.Bond`
-        The bond to serialize.
+    Parameters:
+
+        bond:
+            The bond to serialize.
+
+    Returns:
 
-    Returns
-    -------
-    :class:`dict`
         A JSON representation of `bond`.
 
     """