Merge pull request #172 from speleo3/mypy

Add type annotations and mypy CI job

.github/workflows/python-package.yml

@@ -56,3 +56,13 @@ jobs:
       with:
         name: coverage-html
         path: htmlcov/*
+
+  static_type_check:
+    runs-on: ubuntu-latest
+    steps:
+    - uses: actions/checkout@v2
+    - uses: actions/setup-python@v2
+      with:
+        python-version: "3.12"
+    - run: python -m pip install mypy types-setuptools
+    - run: mypy

.gitignore

@@ -13,3 +13,5 @@
 docs/build
 docs/source/api/*.rst
 build/
+dist/
+.coverage

propka/atom.py

@@ -7,9 +7,15 @@ The :class:`Atom` class contains all atom information found in the PDB file.
 """
 
 import string
+from typing import cast, List, NoReturn, Optional, TYPE_CHECKING
+
 from propka.lib import make_tidy_atom_label
 from . import hybrid36
 
+if TYPE_CHECKING:
+    from propka.group import Group
+    from propka.molecular_container import MolecularContainer
+    from propka.conformation_container import ConformationContainer
 
 # Format strings that get used in multiple places (or are very complex)
 PDB_LINE_FMT1 = (
@@ -37,37 +43,24 @@ class Atom:
        removed as reading/writing PROPKA input is no longer supported.
     """
 
-    def __init__(self, line=None):
+    def __init__(self, line: Optional[str] = None):
         """Initialize Atom object.
 
         Args:
             line:  Line from a PDB file to set properties of atom.
         """
-        self.occ = None
+        self.number_of_bonded_elements: NoReturn = cast(NoReturn, {})  # FIXME unused?
-        self.numb = None
+        self.group: Optional[Group] = None
-        self.res_name = None
+        self.group_type: Optional[str] = None
-        self.type = None
+        self.cysteine_bridge: bool = False
-        self.chain_id = None
+        self.bonded_atoms: List[Atom] = []
-        self.beta = None
-        self.icode = None
-        self.res_num = None
-        self.name = None
-        self.element = None
-        self.x = None
-        self.y = None
-        self.z = None
-        self.group = None
-        self.group_type = None
-        self.number_of_bonded_elements = {}
-        self.cysteine_bridge = False
-        self.bonded_atoms = []
         self.residue = None
-        self.conformation_container = None
+        self.conformation_container: Optional[ConformationContainer] = None
-        self.molecular_container = None
+        self.molecular_container: Optional[MolecularContainer] = None
         self.is_protonated = False
         self.steric_num_lone_pairs_set = False
-        self.terminal = None
+        self.terminal: Optional[str] = None
-        self.charge = 0
+        self.charge = 0.0
         self.charge_set = False
         self.steric_number = 0
         self.number_of_lone_pairs = 0
@@ -84,7 +77,7 @@ class Atom:
         self.sybyl_assigned = False
         self.marvin_pka = False
 
-    def set_properties(self, line):
+    def set_properties(self, line: Optional[str]):
         """Line from PDB file to set properties of atom.
 
         Args:
@@ -112,10 +105,8 @@ class Atom:
             self.z = float(line[46:54].strip())
             self.res_num = int(line[22:26].strip())
             self.res_name = "{0:<3s}".format(line[17:20].strip())
-            self.chain_id = line[21]
             # Set chain id to "_" if it is just white space.
-            if not self.chain_id.strip():
+            self.chain_id = line[21].strip() or '_'
-                self.chain_id = '_'
             self.type = line[:6].strip().lower()
 
             # TODO - define nucleic acid residue names elsewhere
@@ -134,7 +125,7 @@ class Atom:
                 self.element = '{0:1s}{1:1s}'.format(
                     self.element[0], self.element[1].lower())
 
-    def set_group_type(self, type_):
+    def set_group_type(self, type_: str):
         """Set group type of atom.
 
         Args:

propka/bonds.py

@@ -10,6 +10,10 @@ import math
 import json
 import pkg_resources
 import propka.calculations
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from propka.molecular_container import MolecularContainer
 
 
 _LOGGER = logging.getLogger(__name__)
@@ -329,7 +333,7 @@ class BondMaker:
                 return True
         return False
 
-    def find_bonds_for_molecules_using_boxes(self, molecules):
+    def find_bonds_for_molecules_using_boxes(self, molecules: "MolecularContainer"):
         """ Finds all bonds for a molecular container.
 
         Args:

propka/calculations.py

@@ -6,13 +6,17 @@ Mathematical helper functions.
 """
 
 import math
+from typing import Iterable, Optional, Tuple, TypeVar
+from .vector_algebra import _XYZ
 
+_BoundXYZ_1 = TypeVar("_BoundXYZ_1", bound=_XYZ)
+_BoundXYZ_2 = TypeVar("_BoundXYZ_2", bound=_XYZ)
 
 #: Maximum distance used to bound calculations of smallest distance
 MAX_DISTANCE = 1e6
 
 
-def squared_distance(atom1, atom2):
+def squared_distance(atom1: _XYZ, atom2: _XYZ) -> float:
     """Calculate the squared distance between two atoms.
 
     Args:
@@ -28,7 +32,7 @@ def squared_distance(atom1, atom2):
     return res
 
 
-def distance(atom1, atom2):
+def distance(atom1: _XYZ, atom2: _XYZ) -> float:
     """Calculate the distance between two atoms.
 
     Args:
@@ -40,7 +44,10 @@ def distance(atom1, atom2):
     return math.sqrt(squared_distance(atom1, atom2))
 
 
-def get_smallest_distance(atoms1, atoms2):
+def get_smallest_distance(
+    atoms1: Iterable[_BoundXYZ_1],
+    atoms2: Iterable[_BoundXYZ_2],
+) -> Tuple[Optional[_BoundXYZ_1], float, Optional[_BoundXYZ_2]]:
     """Calculate the smallest distance between two groups of atoms.
 
     Args:
@@ -59,4 +66,4 @@ def get_smallest_distance(atoms1, atoms2):
                 res_dist = dist
                 res_atom1 = atom1
                 res_atom2 = atom2
-    return [res_atom1, math.sqrt(res_dist), res_atom2]
+    return (res_atom1, math.sqrt(res_dist), res_atom2)

propka/conformation_container.py

@@ -6,6 +6,12 @@ Container data structure for molecular conformations.
 """
 import logging
 import functools
+from typing import Iterable, List, NoReturn, Optional, TYPE_CHECKING, Set
+
+if TYPE_CHECKING:
+    from propka.atom import Atom
+    from propka.molecular_container import MolecularContainer
+
 import propka.ligand
 from propka.output import make_interaction_map
 from propka.determinant import Determinant
@@ -13,7 +19,6 @@ from propka.coupled_groups import NCCG
 from propka.determinants import set_backbone_determinants, set_ion_determinants
 from propka.determinants import set_determinants
 from propka.group import Group, is_group
-from typing import Iterable
 
 
 _LOGGER = logging.getLogger(__name__)
@@ -38,7 +43,10 @@ class ConformationContainer:
        PROPKA inputs is no longer supported.
     """
 
-    def __init__(self, name='', parameters=None, molecular_container=None):
+    def __init__(self,
+                 name: str = '',
+                 parameters=None,
+                 molecular_container: Optional["MolecularContainer"] = None):
         """Initialize conformation container.
 
         Args:
@@ -49,9 +57,9 @@ class ConformationContainer:
         self.molecular_container = molecular_container
         self.name = name
         self.parameters = parameters
-        self.atoms = []
+        self.atoms: List["Atom"] = []
-        self.groups = []
+        self.groups: List[Group] = []
-        self.chains = []
+        self.chains: List[str] = []
         self.current_iter_item = 0
         self.marvin_pkas_calculated = False
         self.non_covalently_coupled_groups = False
@@ -126,7 +134,8 @@ class ConformationContainer:
                            self.get_titratable_groups()))) > 0:
             self.non_covalently_coupled_groups = True
 
-    def find_bonded_titratable_groups(self, atom, num_bonds, original_atom):
+    def find_bonded_titratable_groups(self, atom: "Atom", num_bonds: int,
+                                      original_atom: "Atom"):
         """Find bonded titrable groups.
 
         Args:
@@ -136,7 +145,7 @@ class ConformationContainer:
         Returns:
             a set of bonded atom groups
         """
-        res = set()
+        res: Set[Group] = set()
         for bond_atom in atom.bonded_atoms:
             # skip the original atom
             if bond_atom == original_atom:
@@ -152,7 +161,7 @@ class ConformationContainer:
                     bond_atom, num_bonds+1, original_atom)
         return res
 
-    def setup_and_add_group(self, group):
+    def setup_and_add_group(self, group: Optional[Group]):
         """Check if we want to include this group in the calculations.
 
         Args:
@@ -166,7 +175,7 @@ class ConformationContainer:
         self.init_group(group)
         self.groups.append(group)
 
-    def init_group(self, group):
+    def init_group(self, group: Group):
         """Initialize the given Group object.
 
         Args:
@@ -178,10 +187,11 @@ class ConformationContainer:
 
         # If --titrate_only option is set, make non-specified residues
         # un-titratable:
+        assert self.molecular_container is not None
         titrate_only = self.molecular_container.options.titrate_only
         if titrate_only is not None:
             atom = group.atom
-            if not (atom.chain_id, atom.res_num, atom.icode) in titrate_only:
+            if (atom.chain_id, atom.res_num, atom.icode) not in titrate_only:
                 group.titratable = False
                 if group.residue_type == 'CYS':
                     group.exclude_cys_from_results = True
@@ -475,7 +485,7 @@ class ConformationContainer:
             group for group in self.groups
             if group.residue_type in self.parameters.ions.keys()]
 
-    def get_group_names(self, group_list):
+    def get_group_names(self, group_list: NoReturn) -> NoReturn:  # FIXME unused?
         """Get names of groups in list.
 
         Args:
@@ -483,9 +493,11 @@ class ConformationContainer:
         Returns:
             list of groups
         """
+        if TYPE_CHECKING:
+            assert False
         return [group for group in self.groups if group.type in group_list]
 
-    def get_ligand_atoms(self):
+    def get_ligand_atoms(self) -> List["Atom"]:
         """Get atoms associated with ligands.
 
         Returns:
@@ -493,7 +505,7 @@ class ConformationContainer:
         """
         return [atom for atom in self.atoms if atom.type == 'hetatm']
 
-    def get_heavy_ligand_atoms(self):
+    def get_heavy_ligand_atoms(self) -> List["Atom"]:
         """Get heavy atoms associated with ligands.
 
         Returns:
@@ -503,7 +515,7 @@ class ConformationContainer:
             atom for atom in self.atoms
             if atom.type == 'hetatm' and atom.element != 'H']
 
-    def get_chain(self, chain):
+    def get_chain(self, chain: str) -> List["Atom"]:
         """Get atoms associated with a specific chain.
 
         Args:
@@ -513,7 +525,7 @@ class ConformationContainer:
         """
         return [atom for atom in self.atoms if atom.chain_id != chain]
 
-    def add_atom(self, atom):
+    def add_atom(self, atom: "Atom"):
         """Add atom to container.
 
         Args:
@@ -556,7 +568,7 @@ class ConformationContainer:
         """
         self.top_up_from_atoms(other.atoms)
 
-    def top_up_from_atoms(self, other_atoms: Iterable["propka.atom.Atom"]):
+    def top_up_from_atoms(self, other_atoms: Iterable["Atom"]):
         """Adds atoms which are missing from this container.
 
         Args:
@@ -613,7 +625,7 @@ class ConformationContainer:
             self.atoms[i].numb = i+1
 
     @staticmethod
-    def sort_atoms_key(atom):
+    def sort_atoms_key(atom: "Atom") -> float:
         """Generate key for atom sorting.
 
         Args:

propka/energy.py

@@ -7,6 +7,12 @@ Energy calculations.
 """
 import math
 import logging
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from propka.conformation_container import ConformationContainer
+    from propka.group import Group
+
 from propka.calculations import squared_distance, get_smallest_distance
 
 
@@ -27,13 +33,14 @@ COMBINED_NUM_BURIED_MAX = 900
 SEPARATE_NUM_BURIED_MAX = 400
 
 
-def radial_volume_desolvation(parameters, group):
+def radial_volume_desolvation(parameters, group: "Group") -> None:
     """Calculate desolvation terms for group.
 
     Args:
         parameters:  parameters for desolvation calculation
         group:  group of atoms for calculation
     """
+    assert group.atom.conformation_container is not None
     all_atoms = group.atom.conformation_container.get_non_hydrogen_atoms()
     volume = 0.0
     group.num_volume = 0
@@ -66,7 +73,7 @@ def radial_volume_desolvation(parameters, group):
         * volume_after_allowance * scale_factor)
 
 
-def calculate_scale_factor(parameters, weight):
+def calculate_scale_factor(parameters, weight: float) -> float:
     """Calculate desolvation scaling factor.
 
     Args:
@@ -82,7 +89,7 @@ def calculate_scale_factor(parameters, weight):
     return scale_factor
 
 
-def calculate_weight(parameters, num_volume):
+def calculate_weight(parameters, num_volume: int) -> float:
     """Calculate the atom-based desolvation weight.
 
     TODO - figure out why a similar function exists in version.py
@@ -102,7 +109,7 @@ def calculate_weight(parameters, num_volume):
     return weight
 
 
-def calculate_pair_weight(parameters, num_volume1, num_volume2):
+def calculate_pair_weight(parameters, num_volume1: int, num_volume2: int) -> float:
     """Calculate the atom-pair based desolvation weight.
 
     Args:
@@ -120,7 +127,7 @@ def calculate_pair_weight(parameters, num_volume1, num_volume2):
     return weight
 
 
-def hydrogen_bond_energy(dist, dpka_max, cutoffs, f_angle=1.0):
+def hydrogen_bond_energy(dist, dpka_max: float, cutoffs, f_angle=1.0) -> float:
     """Calculate hydrogen-bond interaction pKa shift.
 
     Args:
@@ -319,7 +326,7 @@ def check_coulomb_pair(parameters, group1, group2, dist):
     return do_coulomb
 
 
-def coulomb_energy(dist, weight, parameters):
+def coulomb_energy(dist: float, weight: float, parameters) -> float:
     """Calculates the Coulomb interaction pKa shift based on Coulomb's law.
 
     Args:
@@ -340,7 +347,7 @@ def coulomb_energy(dist, weight, parameters):
     return abs(dpka)
 
 
-def backbone_reorganization(_, conformation):
+def backbone_reorganization(_, conformation: "ConformationContainer") -> None:
     """Perform calculations related to backbone reorganizations.
 
     NOTE - this was described in the code as "adding test stuff"

propka/group.py

@@ -11,8 +11,11 @@ Routines and classes for storing groups important to PROPKA calculations.
 """
 import logging
 import math
+from typing import cast, Dict, Iterable, List, NoReturn, Optional
+
 import propka.ligand
 import propka.protonate
+from propka.atom import Atom
 from propka.ligand_pka_values import LigandPkaValues
 from propka.determinant import Determinant
 
@@ -57,7 +60,7 @@ class Group:
        longer supported.
     """
 
-    def __init__(self, atom):
+    def __init__(self, atom: Atom):
         """Initialize with an atom.
 
         Args:
@@ -67,7 +70,11 @@ class Group:
         self.type = ''
         atom.group = self
         # set up data structures
-        self.determinants = {'sidechain': [], 'backbone': [], 'coulomb': []}
+        self.determinants: Dict[str, List[Determinant]] = {
+            'sidechain': [],
+            'backbone': [],
+            'coulomb': [],
+        }
         self.pka_value = 0.0
         self.model_pka = 0.0
         # Energy associated with volume interactions
@@ -84,16 +91,16 @@ class Group:
         self.z = 0.0
         self.charge = 0
         self.parameters = None
-        self.exclude_cys_from_results = None
+        self.exclude_cys_from_results = False
-        self.interaction_atoms_for_acids = []
+        self.interaction_atoms_for_acids: List[Atom] = []
-        self.interaction_atoms_for_bases = []
+        self.interaction_atoms_for_bases: List[Atom] = []
         self.model_pka_set = False
         self.intrinsic_pka = None
-        self.titratable = None
+        self.titratable = False
         # information on covalent and non-covalent coupling
-        self.non_covalently_coupled_groups = []
+        self.non_covalently_coupled_groups: List["Group"] = []
-        self.covalently_coupled_groups = []
+        self.covalently_coupled_groups: List["Group"] = []
-        self.coupled_titrating_group = None
+        self.coupled_titrating_group: Optional["Group"] = None
         self.common_charge_centre = False
         self.residue_type = self.atom.res_name
         if self.atom.terminal:
@@ -112,9 +119,9 @@ class Group:
             self.label = fmt.format(
                 type=self.residue_type, name=atom.name, chain=atom.chain_id)
         # container for squared distances
-        self.squared_distances = {}
+        self.squared_distances: NoReturn = cast(NoReturn, {})  # FIXME unused?
 
-    def couple_covalently(self, other):
+    def couple_covalently(self, other: "Group") -> None:
         """Couple this group with another group.
 
         Args:
@@ -126,7 +133,7 @@ class Group:
         if self not in other.covalently_coupled_groups:
             other.covalently_coupled_groups.append(self)
 
-    def couple_non_covalently(self, other):
+    def couple_non_covalently(self, other: "Group") -> None:
         """Non-covalenthly couple this group with another group.
 
         Args:
@@ -154,7 +161,7 @@ class Group:
         """
         return self.non_covalently_coupled_groups
 
-    def share_determinants(self, others):
+    def share_determinants(self, others: Iterable["Group"]) -> None:
         """Share determinants between this group and others.
 
         Args:
@@ -172,7 +179,7 @@ class Group:
         self.calculate_total_pka()
         the_other.calculate_total_pka()
 
-    def share_determinant(self, new_determinant, type_):
+    def share_determinant(self, new_determinant: Determinant, type_: str) -> None:
         """Add determinant to this group's list of determinants.
 
         Args:
@@ -230,7 +237,7 @@ class Group:
                 self.add_determinant(determinant, type_)
         return self
 
-    def add_determinant(self, new_determinant, type_):
+    def add_determinant(self, new_determinant: Determinant, type_: str) -> None:
         """Add to current and creates non-present determinants.
 
         Args:
@@ -247,7 +254,7 @@ class Group:
         self.determinants[type_].append(Determinant(new_determinant.group,
                                                     new_determinant.value))
 
-    def set_determinant(self, new_determinant, type_):
+    def set_determinant(self, new_determinant: Determinant, type_: str) -> None:
         """Overwrite current and create non-present determinants.
 
         Args:
@@ -345,8 +352,8 @@ class Group:
         # set the main atom as interaction atom
         self.set_interaction_atoms([self.atom], [self.atom])
 
-    def set_interaction_atoms(self, interaction_atoms_for_acids,
+    def set_interaction_atoms(self, interaction_atoms_for_acids: List[Atom],
-                              interaction_atoms_for_bases):
+                              interaction_atoms_for_bases: List[Atom]):
         """Set interacting atoms and group types.
 
         Args:
@@ -359,10 +366,10 @@ class Group:
         self.interaction_atoms_for_bases = interaction_atoms_for_bases
         # check if all atoms have been identified
         ok = True
-        for [expect, found, _] in [[EXPECTED_ATOMS_ACID_INTERACTIONS,
+        for (expect, found) in [
-                                    self.interaction_atoms_for_acids, 'acid'],
+            (EXPECTED_ATOMS_ACID_INTERACTIONS, self.interaction_atoms_for_acids),
-                                   [EXPECTED_ATOMS_BASE_INTERACTIONS,
+            (EXPECTED_ATOMS_BASE_INTERACTIONS, self.interaction_atoms_for_bases),
-                                    self.interaction_atoms_for_bases, 'base']]:
+        ]:
             if self.type in expect.keys():
                 for elem in expect[self.type].keys():
                     if (len([a for a in found if a.element == elem])
@@ -395,7 +402,7 @@ class Group:
                     '             {0:s}'.format(
                         str(self.interaction_atoms_for_bases[i])))
 
-    def get_interaction_atoms(self, interacting_group):
+    def get_interaction_atoms(self, interacting_group) -> List[Atom]:
         """Get atoms involved in interaction with other group.
 
         Args:
@@ -403,6 +410,7 @@ class Group:
         Returns:
             list of atoms
         """
+        assert self.parameters is not None
         if interacting_group.residue_type in self.parameters.base_list:
             return self.interaction_atoms_for_bases
         else:
@@ -518,7 +526,7 @@ class Group:
             self.model_pka + self.energy_volume + self.energy_local
             + back_bone + side_chain)
 
-    def get_summary_string(self, remove_penalised_group=False):
+    def get_summary_string(self, remove_penalised_group: bool = False) -> str:
         """Create summary string for this group.
 
         Args:
@@ -1210,7 +1218,7 @@ class TitratableLigandGroup(Group):
         self.model_pka_set = True
 
 
-def is_group(parameters, atom):
+def is_group(parameters, atom: Atom) -> Optional[Group]:
     """Identify whether the atom belongs to a group.
 
     Args:
@@ -1244,7 +1252,7 @@ def is_group(parameters, atom):
     return None
 
 
-def is_protein_group(parameters, atom):
+def is_protein_group(parameters, atom: Atom) -> Optional[Group]:
     """Identify whether the atom belongs to a protein group.
 
     Args:
@@ -1278,7 +1286,7 @@ def is_protein_group(parameters, atom):
     return None
 
 
-def is_ligand_group_by_groups(_, atom):
+def is_ligand_group_by_groups(_, atom: Atom) -> Optional[Group]:
     """Identify whether the atom belongs to a ligand group by checking groups.
 
     Args:
@@ -1360,7 +1368,7 @@ def is_ligand_group_by_groups(_, atom):
     return None
 
 
-def is_ligand_group_by_marvin_pkas(parameters, atom):
+def is_ligand_group_by_marvin_pkas(parameters, atom: Atom) -> Optional[Group]:
     """Identify whether the atom belongs to a ligand group by calculating
     'Marvin pKas'.
 
@@ -1375,6 +1383,7 @@ def is_ligand_group_by_marvin_pkas(parameters, atom):
     # calculate Marvin ligand pkas for this conformation container
     # if not already done
     # TODO - double-check testing coverage of these functions.
+    assert atom.conformation_container is not None
     if not atom.conformation_container.marvin_pkas_calculated:
         lpka = LigandPkaValues(parameters)
         lpka.get_marvin_pkas_for_molecular_container(
@@ -1396,7 +1405,7 @@ def is_ligand_group_by_marvin_pkas(parameters, atom):
     return None
 
 
-def is_ion_group(parameters, atom):
+def is_ion_group(parameters, atom: Atom) -> Optional[Group]:
     """Identify whether the atom belongs to an ion group.
 
     Args:

propka/hydrogens.py

@@ -7,15 +7,19 @@ Calculations related to hydrogen placement.
 """
 import math
 import logging
+from typing import List, Optional, Tuple, TYPE_CHECKING
+
 from propka.protonate import Protonate
 from propka.bonds import BondMaker
 from propka.atom import Atom
 
+if TYPE_CHECKING:
+    from propka.molecular_container import MolecularContainer
 
 _LOGGER = logging.getLogger(__name__)
 
 
-def setup_bonding_and_protonation(molecular_container):
+def setup_bonding_and_protonation(molecular_container: "MolecularContainer") -> None:
     """Set up bonding and protonation for a molecule.
 
     Args:
@@ -34,7 +38,7 @@ def setup_bonding_and_protonation(molecular_container):
         protonator.protonate(molecular_container)
 
 
-def setup_bonding(molecular_container):
+def setup_bonding(molecular_container: "MolecularContainer") -> BondMaker:
     """Set up bonding for a molecular container.
 
     Args:
@@ -47,7 +51,7 @@ def setup_bonding(molecular_container):
     return my_bond_maker
 
 
-def setup_bonding_and_protonation_30_style(molecular_container):
+def setup_bonding_and_protonation_30_style(molecular_container: "MolecularContainer") -> BondMaker:
     """Set up bonding for a molecular container.
 
     Args:
@@ -63,7 +67,7 @@ def setup_bonding_and_protonation_30_style(molecular_container):
     return bond_maker
 
 
-def protonate_30_style(molecular_container):
+def protonate_30_style(molecular_container: "MolecularContainer") -> None:
     """Protonate the molecule.
 
     Args:
@@ -73,9 +77,9 @@ def protonate_30_style(molecular_container):
         _LOGGER.info('Now protonating %s', name)
         # split atom into residues
         curres = -1000000
-        residue = []
+        residue: List[Atom] = []
-        o_atom = None
+        o_atom: Optional[Atom] = None
-        c_atom = None
+        c_atom: Optional[Atom] = None
         for atom in molecular_container.conformations[name].atoms:
             if atom.res_num != curres:
                 curres = atom.res_num
@@ -100,7 +104,7 @@ def protonate_30_style(molecular_container):
                 residue.append(atom)
 
 
-def set_ligand_atom_names(molecular_container):
+def set_ligand_atom_names(molecular_container: "MolecularContainer") -> None:
     """Set names for ligands in molecular container.
 
     Args:
@@ -110,7 +114,7 @@ def set_ligand_atom_names(molecular_container):
         molecular_container.conformations[name].set_ligand_atom_names()
 
 
-def add_arg_hydrogen(residue):
+def add_arg_hydrogen(residue: List[Atom]) -> List[Atom]:
     """Adds Arg hydrogen atoms to residues according to the 'old way'.
 
     Args:
@@ -142,7 +146,7 @@ def add_arg_hydrogen(residue):
     return [h1_atom, h2_atom, h3_atom, h4_atom, h5_atom]
 
 
-def add_his_hydrogen(residue):
+def add_his_hydrogen(residue: List[Atom]) -> None:
     """Adds His hydrogen atoms to residues according to the 'old way'.
 
     Args:
@@ -165,7 +169,7 @@ def add_his_hydrogen(residue):
     he_atom.name = "HNE"
 
 
-def add_trp_hydrogen(residue):
+def add_trp_hydrogen(residue: List[Atom]) -> None:
     """Adds Trp hydrogen atoms to residues according to the 'old way'.
 
     Args:
@@ -188,7 +192,7 @@ def add_trp_hydrogen(residue):
     he_atom.name = "HNE"
 
 
-def add_amd_hydrogen(residue):
+def add_amd_hydrogen(residue: List[Atom]) -> None:
     """Adds Gln & Asn hydrogen atoms to residues according to the 'old way'.
 
     Args:
@@ -217,7 +221,9 @@ def add_amd_hydrogen(residue):
     h2_atom.name = "HN2"
 
 
-def add_backbone_hydrogen(residue, o_atom, c_atom):
+def add_backbone_hydrogen(residue: List[Atom],
+                          o_atom: Optional[Atom],
+                          c_atom: Optional[Atom]) -> Tuple[Optional[Atom], Optional[Atom]]:
     """Adds hydrogen backbone atoms to residues according to the old way.
 
     dR is wrong for the N-terminus (i.e. first residue) but it doesn't affect
@@ -240,18 +246,18 @@ def add_backbone_hydrogen(residue, o_atom, c_atom):
             new_c_atom = atom
         if atom.name == "O":
             new_o_atom = atom
-    if None in [c_atom, o_atom, n_atom]:
+    if c_atom is None or o_atom is None or n_atom is None:
-        return [new_o_atom, new_c_atom]
+        return (new_o_atom, new_c_atom)
     if n_atom.res_name == "PRO":
         # PRO doesn't have an H-atom; do nothing
         pass
     else:
         h_atom = protonate_direction(n_atom, o_atom, c_atom)
         h_atom.name = "H"
-    return [new_o_atom, new_c_atom]
+    return (new_o_atom, new_c_atom)
 
 
-def protonate_direction(x1_atom, x2_atom, x3_atom):
+def protonate_direction(x1_atom: Atom, x2_atom: Atom, x3_atom: Atom) -> Atom:
     """Protonates an atom, x1_atom, given a direction.
 
     New direction for x1_atom proton is (x2_atom -> x3_atom).
@@ -275,7 +281,7 @@ def protonate_direction(x1_atom, x2_atom, x3_atom):
     return h_atom
 
 
-def protonate_average_direction(x1_atom, x2_atom, x3_atom):
+def protonate_average_direction(x1_atom: Atom, x2_atom: Atom, x3_atom: Atom) -> Atom:
     """Protonates an atom, x1_atom, given a direction.
 
     New direction for x1_atom is (x1_atom/x2_atom -> x3_atom).
@@ -301,7 +307,7 @@ def protonate_average_direction(x1_atom, x2_atom, x3_atom):
     return h_atom
 
 
-def protonate_sp2(x1_atom, x2_atom, x3_atom):
+def protonate_sp2(x1_atom: Atom, x2_atom: Atom, x3_atom: Atom) -> Atom:
     """Protonates a SP2 atom, given a list of atoms
 
     Args:
@@ -323,7 +329,7 @@ def protonate_sp2(x1_atom, x2_atom, x3_atom):
     return h_atom
 
 
-def make_new_h(atom, x, y, z):
+def make_new_h(atom: Atom, x: float, y: float, z: float) -> Atom:
     """Add a new hydrogen to an atom at the specified position.
 
     Args:
@@ -347,5 +353,6 @@ def make_new_h(atom, x, y, z):
     new_h.number_of_protons_to_add = 0
     new_h.num_pi_elec_2_3_bonds = 0
     atom.bonded_atoms.append(new_h)
+    assert atom.conformation_container is not None
     atom.conformation_container.add_atom(new_h)
     return new_h

propka/input.py

@@ -10,12 +10,15 @@ Input routines.
    :func:`get_atom_lines_from_input`) have been removed.
 """
 import typing
+from typing import Iterator, Tuple
 import contextlib
 from pathlib import Path
 from pkg_resources import resource_filename
 from propka.lib import protein_precheck
 from propka.atom import Atom
 from propka.conformation_container import ConformationContainer
+from propka.molecular_container import MolecularContainer
+from propka.parameters import Parameters
 
 
 def open_file_for_reading(
@@ -27,18 +30,14 @@ def open_file_for_reading(
         input_file: path to file or file-like object. If file-like object,
         then will attempt seek(0).
     """
-    try:
+    if not isinstance(input_file, (str, Path)):
         input_file.seek(0)
-    except AttributeError:
+        return contextlib.nullcontext(input_file)
-        pass
-    else:
-        # TODO use contextlib.nullcontext when dropping Python 3.6 support
-        return contextlib.contextmanager(lambda: (yield input_file))()
 
     return contextlib.closing(open(input_file, 'rt'))
 
 
-def read_molecule_file(filename: str, mol_container, stream=None):
+def read_molecule_file(filename: str, mol_container: MolecularContainer, stream=None) -> MolecularContainer:
     """Read input file or stream (PDB or PROPKA) for a molecular container
 
     Args:
@@ -123,7 +122,7 @@ def read_molecule_file(filename: str, mol_container, stream=None):
     return mol_container
 
 
-def read_parameter_file(input_file, parameters):
+def read_parameter_file(input_file, parameters: Parameters) -> Parameters:
     """Read a parameter file.
 
     Args:
@@ -144,7 +143,7 @@ def read_parameter_file(input_file, parameters):
     return parameters
 
 
-def conformation_sorter(conf):
+def conformation_sorter(conf: str) -> int:
     """TODO - figure out what this function does."""
     model = int(conf[:-1])
     altloc = conf[-1:]
@@ -152,7 +151,7 @@ def conformation_sorter(conf):
 
 
 def get_atom_lines_from_pdb(pdb_file, ignore_residues=[], keep_protons=False,
-                            tags=['ATOM  ', 'HETATM'], chains=None):
+                            tags=['ATOM  ', 'HETATM'], chains=None) -> Iterator[Tuple[str, Atom]]:
     """Get atom lines from PDB file.
 
     Args:
@@ -237,7 +236,7 @@ def read_pdb(pdb_file, parameters, molecule):
         keep_protons=molecule.options.keep_protons,
         chains=molecule.options.chains)
     for (name, atom) in lines:
-        if not name in conformations.keys():
+        if name not in conformations.keys():
             conformations[name] = ConformationContainer(
                 name=name, parameters=parameters, molecular_container=molecule)
         conformations[name].add_atom(atom)

propka/molecular_container.py

@@ -6,6 +6,8 @@ Molecular container for storing all contents of PDB files.
 """
 import logging
 import os
+from typing import Dict, List, Optional, Tuple
+
 import propka.version
 from propka.output import write_pka, print_header, print_result
 from propka.conformation_container import ConformationContainer
@@ -28,7 +30,12 @@ class MolecularContainer:
        PROPKA input files is no longer supported.
     """
 
-    def __init__(self, parameters, options=None):
+    conformation_names: List[str]
+    conformations: Dict[str, ConformationContainer]
+    name: Optional[str]
+    version: propka.version.Version
+
+    def __init__(self, parameters, options=None) -> None:
         """Initialize molecular container.
 
         Args:
@@ -50,7 +57,7 @@ class MolecularContainer:
                 parameters.version)
             raise Exception(errstr)
 
-    def top_up_conformations(self):
+    def top_up_conformations(self) -> None:
         """Makes sure that all atoms are present in all conformations."""
         ref_atoms = {
             atom.residue_label: atom
@@ -60,24 +67,24 @@ class MolecularContainer:
         for conf in self.conformations.values():
             conf.top_up_from_atoms(ref_atoms.values())
 
-    def find_covalently_coupled_groups(self):
+    def find_covalently_coupled_groups(self) -> None:
         """Find covalently coupled groups."""
         for name in self.conformation_names:
             self.conformations[name].find_covalently_coupled_groups()
 
-    def find_non_covalently_coupled_groups(self):
+    def find_non_covalently_coupled_groups(self) -> None:
         """Find non-covalently coupled groups."""
         verbose = self.options.display_coupled_residues
         for name in self.conformation_names:
             self.conformations[name].find_non_covalently_coupled_groups(
                 verbose=verbose)
 
-    def extract_groups(self):
+    def extract_groups(self) -> None:
         """Identify the groups needed for pKa calculation."""
         for name in self.conformation_names:
             self.conformations[name].extract_groups()
 
-    def calculate_pka(self):
+    def calculate_pka(self) -> None:
         """Calculate pKa values."""
         # calculate for each conformation
         for name in self.conformation_names:
@@ -90,7 +97,7 @@ class MolecularContainer:
         # print out the conformation-average results
         print_result(self, 'AVR', self.version.parameters)
 
-    def average_of_conformations(self):
+    def average_of_conformations(self) -> None:
         """Generate an average of conformations."""
         parameters = self.conformations[self.conformation_names[0]].parameters
         # make a new configuration to hold the average values
@@ -124,7 +131,7 @@ class MolecularContainer:
         self.conformations['AVR'] = avr_conformation
 
     def write_pka(self, filename=None, reference="neutral",
-                  direction="folding", options=None):
+                  direction="folding", options=None) -> None:
         """Write pKa information to a file.
 
         Args:
@@ -187,7 +194,7 @@ class MolecularContainer:
             stability_range = [min(stable_values), max(stable_values)]
         return profile, opt, range_80pct, stability_range
 
-    def get_charge_profile(self, conformation='AVR', grid=[0., 14., .1]):
+    def get_charge_profile(self, conformation: str = 'AVR', grid=[0., 14., .1]):
         """Get charge profile for conformation as function of pH.
 
         Args:
@@ -196,7 +203,7 @@ class MolecularContainer:
         Returns:
             list of charge state values
         """
-        charge_profile = []
+        charge_profile: List[List[float]] = []
         for ph in make_grid(*grid):
             conf = self.conformations[conformation]
             q_unfolded, q_folded = conf.calculate_charge(
@@ -204,8 +211,8 @@ class MolecularContainer:
             charge_profile.append([ph, q_unfolded, q_folded])
         return charge_profile
 
-    def get_pi(self, conformation='AVR', grid=[0., 14., 1], *,
+    def get_pi(self, conformation: str = 'AVR', grid=[0., 14., 1], *,
-            precision: float = 1e-4):
+               precision: float = 1e-4) -> Tuple[float, float]:
         """Get the isoelectric points for folded and unfolded states.
 
         Args:

propka/vector_algebra.py

@@ -6,16 +6,35 @@ Vector algebra for PROPKA.
 """
 import logging
 import math
+from typing import Optional, Protocol, Union
 from propka.lib import get_sorted_configurations
 
 
 _LOGGER = logging.getLogger(__name__)
 
 
+class _XYZ(Protocol):
+    """
+    Protocol for types which have x/y/z attributes, like Vector or Atom.
+    """
+    x: float
+    y: float
+    z: float
+
+
 class Vector:
     """Vector"""
 
-    def __init__(self, xi=0.0, yi=0.0, zi=0.0, atom1=None, atom2=None):
+    x: float
+    y: float
+    z: float
+
+    def __init__(self,
+                 xi: float = 0.0,
+                 yi: float = 0.0,
+                 zi: float = 0.0,
+                 atom1: Optional[_XYZ] = None,
+                 atom2: Optional[_XYZ] = None):
         """Initialize vector.
 
         Args:
@@ -41,17 +60,17 @@ class Vector:
                 self.y = atom2.y - self.y
                 self.z = atom2.z - self.z
 
-    def __add__(self, other):
+    def __add__(self, other: _XYZ):
         return Vector(self.x + other.x,
                       self.y + other.y,
                       self.z + other.z)
 
-    def __sub__(self, other):
+    def __sub__(self, other: _XYZ):
         return Vector(self.x - other.x,
                       self.y - other.y,
                       self.z - other.z)
 
-    def __mul__(self, other):
+    def __mul__(self, other: Union["Vector", "Matrix4x4", float]):
         """Dot product, scalar and matrix multiplication."""
         if isinstance(other, Vector):
             return self.x * other.x + self.y * other.y + self.z * other.z
@@ -66,14 +85,12 @@ class Vector:
                 )
         elif type(other) in [int, float]:
             return Vector(self.x * other, self.y * other, self.z * other)
-        else:
+        raise TypeError(f'{type(other)} not supported')
-            _LOGGER.info('{0:s} not supported'.format(type(other)))
-            raise TypeError
 
     def __rmul__(self, other):
         return self.__mul__(other)
 
-    def __pow__(self, other):
+    def __pow__(self, other: _XYZ):
         """Cross product."""
         return Vector(self.y * other.z - self.z * other.y,
                       self.z * other.x - self.x * other.z,
@@ -89,7 +106,7 @@ class Vector:
         """Return vector squared-length"""
         return self.x * self.x + self.y * self.y + self.z * self.z
 
-    def length(self):
+    def length(self) -> float:
         """Return vector length."""
         return math.sqrt(self.sq_length())
 
@@ -107,7 +124,7 @@ class Vector:
             res = Vector(self.z, 0, -self.x)
         return res
 
-    def rescale(self, new_length):
+    def rescale(self, new_length: float):
         """ Rescale vector to new length while preserving direction """
         frac = new_length/(self.length())
         res = Vector(xi=self.x*frac, yi=self.y*frac, zi=self.z*frac)
@@ -145,7 +162,7 @@ class Matrix4x4:
         self.a44 = a44i
 
 
-def angle(avec, bvec):
+def angle(avec: Vector, bvec: Vector) -> float:
     """Get the angle between two vectors.
 
     Args:
@@ -158,7 +175,7 @@ def angle(avec, bvec):
     return math.acos(dot / (avec.length() * bvec.length()))
 
 
-def angle_degrees(avec, bvec):
+def angle_degrees(avec: Vector, bvec: Vector) -> float:
     """Get the angle between two vectors in degrees.
 
     Args:
@@ -170,7 +187,7 @@ def angle_degrees(avec, bvec):
     return math.degrees(angle(avec, bvec))
 
 
-def signed_angle_around_axis(avec, bvec, axis):
+def signed_angle_around_axis(avec: Vector, bvec: Vector, axis: Vector) -> float:
     """Get signed angle of two vectors around axis in radians.
 
     Args:
@@ -189,7 +206,7 @@ def signed_angle_around_axis(avec, bvec, axis):
     return ang
 
 
-def rotate_vector_around_an_axis(theta, axis, vec):
+def rotate_vector_around_an_axis(theta: float, axis: Vector, vec: Vector) -> Vector:
     """Rotate vector around an axis.
 
     Args:
@@ -225,7 +242,7 @@ def rotate_vector_around_an_axis(theta, axis, vec):
     return vec
 
 
-def rotate_atoms_around_z_axis(theta):
+def rotate_atoms_around_z_axis(theta: float) -> Matrix4x4:
     """Get rotation matrix for z-axis.
 
     Args:
@@ -253,7 +270,7 @@ def rotate_atoms_around_z_axis(theta):
         )
 
 
-def rotate_atoms_around_y_axis(theta):
+def rotate_atoms_around_y_axis(theta: float) -> Matrix4x4:
     """Get rotation matrix for y-axis.
 
     Args:

setup.cfg

@@ -19,3 +19,15 @@ omit =
 exclude_lines =
     pragma: no cover
     
+[yapf]
+column_limit = 88
+based_on_style = pep8
+allow_split_before_dict_value = False
+
+[mypy]
+files = propka,tests
+exclude = (?x)(
+    /_version\.py$
+    )
+explicit_package_bases = True
+ignore_missing_imports = True

tests/test_basic_regression.py

@@ -10,6 +10,7 @@ from propka.parameters import Parameters
 from propka.molecular_container import MolecularContainer
 from propka.input import read_parameter_file, read_molecule_file
 from propka.lib import loadOptions
+from typing import List
 
 
 _LOGGER = logging.getLogger(__name__)
@@ -32,8 +33,6 @@ RESULTS_DIR = Path("tests/results")
 if not RESULTS_DIR.is_dir():
     _LOGGER.warning("Switching to sub-directory")
     RESULTS_DIR = Path("results")
-# Arguments to add to all tests
-DEFAULT_ARGS = []
 
 
 def get_test_dirs():
@@ -87,8 +86,8 @@ def run_propka(options, pdb_path, tmp_path):
 def parse_pka(pka_path: Path) -> dict:
     """Parse testable data from a .pka file into a dictionary.
     """
-    pka_list = []
+    pka_list: List[float] = []
-    data = {"pKa": pka_list}
+    data: dict = {"pKa": pka_list}
 
     with open(pka_path, "rt") as pka_file:
         at_pka = False
@@ -98,6 +97,7 @@ def parse_pka(pka_path: Path) -> dict:
                     at_pka = False
                 else:
                     m = re.search(r'\d+\.\d+', line[13:])
+                    assert m is not None
                     pka_list.append(float(m.group()))
             elif "model-pKa" in line:
                 at_pka = True