Static typing for Bio.PDB (#4641)

mypy type checking had been disabled for Bio.PDB

.mypy.ini

@@ -17,8 +17,32 @@ warn_redundant_casts = True
 warn_unused_configs = True
 
 [mypy-Bio/PDB/*]
+disallow_incomplete_defs = False
+
+[mypy-Bio.PDB.ccealign]
+ignore_missing_imports = True
+
+[mypy-Bio.PDB.ic_rebuild]
 ignore_errors = True
 
+[mypy-Bio.PDB.internal_coords]
+ignore_errors = True
+
+[mypy-Bio.PDB.kdtrees]
+ignore_missing_imports = True
+
+[mypy-Bio.PDB.PDBMLParser]
+ignore_errors = True
+
+[mypy-Bio.PDB.PICIO]
+ignore_errors = True
+
+[mypy-igraph.*]
+ignore_missing_imports = True
+
+[mypy-mmtf.*]
+ignore_missing_imports = True
+
 [mypy-numpy.*]
 ignore_missing_imports = True
 

Bio/PDB/Atom.py

@@ -9,6 +9,7 @@
 
 import copy
 import sys
+from typing import Optional, TYPE_CHECKING
 import warnings
 
 import numpy as np
@@ -18,6 +19,9 @@ from Bio.PDB.PDBExceptions import PDBConstructionWarning
 from Bio.PDB.vectors import Vector
 from Bio.Data import IUPACData
 
+if TYPE_CHECKING:
+    from Bio.PDB.Residue import Residue
+
 
 class Atom:
     """Define Atom class.
@@ -33,16 +37,16 @@ class Atom:
 
     def __init__(
         self,
-        name,
-        coord,
-        bfactor,
-        occupancy,
-        altloc,
-        fullname,
+        name: str,
+        coord: np.ndarray,
+        bfactor: Optional[float],
+        occupancy: Optional[float],
+        altloc: str,
+        fullname: str,
         serial_number,
-        element=None,
-        pqr_charge=None,
-        radius=None,
+        element: Optional[str] = None,
+        pqr_charge: Optional[float] = None,
+        radius: Optional[float] = None,
     ):
         """Initialize Atom object.
 
@@ -76,7 +80,7 @@ class Atom:
         """
         self.level = "A"
         # Reference to the residue
-        self.parent = None
+        self.parent: Optional["Residue"] = None
         # the atomic data
         self.name = name  # eg. CA, spaces are removed from atom name
         self.fullname = fullname  # e.g. " CA ", spaces included
@@ -92,7 +96,7 @@ class Atom:
         self.sigatm_array = None
         self.serial_number = serial_number
         # Dictionary that keeps additional properties
-        self.xtra = {}
+        self.xtra: dict = {}
         assert not element or element == element.upper(), element
         self.element = self._assign_element(element)
         self.mass = self._assign_atom_mass()

Bio/PDB/Chain.py

@@ -10,10 +10,14 @@
 from Bio.PDB.Entity import Entity
 from Bio.PDB.internal_coords import IC_Chain
 
-from typing import Optional
+from typing import Optional, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from Bio.PDB.Model import Model
+    from Bio.PDB.Residue import Residue
 
 
-class Chain(Entity):
+class Chain(Entity["Model", "Residue"]):
     """Define Chain class.
 
     Chain is an object of type Entity, stores residues and includes a method to
@@ -208,7 +212,8 @@ class Chain(Entity):
                 verbose=verbose, start=start, fin=fin
             )
         else:
+            structure = None if self.parent is None else self.parent.parent
             raise Exception(
                 "Structure %s Chain %s does not have internal coordinates set"
-                % (self.parent.parent, self)
+                % (structure, self)
             )

Bio/PDB/Entity.py

@@ -13,20 +13,32 @@ import warnings
 
 from collections import deque
 from copy import copy
+from typing import TYPE_CHECKING, Any, Dict, Generic, List, Optional, TypeVar, Union
 
 import numpy as np
 
 from Bio import BiopythonWarning
 from Bio.PDB.PDBExceptions import PDBConstructionException
 
+if TYPE_CHECKING:
+    from Bio.PDB.Atom import Atom
 
-class Entity:
+_Child = TypeVar("_Child", bound=Union["Entity", "Atom"])
+_Parent = TypeVar("_Parent", bound=Optional["Entity"])
+
+
+class Entity(Generic[_Parent, _Child]):
     """Basic container object for PDB hierarchy.
 
     Structure, Model, Chain and Residue are subclasses of Entity.
     It deals with storage and lookup.
     """
 
+    parent: Optional[_Parent]
+    child_list: List[_Child]
+    child_dict: Dict[Any, _Child]
+    level: str
+
     def __init__(self, id):
         """Initialize the class."""
         self._id = id
@@ -174,7 +186,7 @@ class Entity:
         """
         if value == self._id:
             return
-        if self.parent:
+        if self.parent is not None:
             if value in self.parent.child_dict:
                 # See issue 1551 for details on the downgrade.
                 warnings.warn(
@@ -200,7 +212,7 @@ class Entity:
         """
         return self.level
 
-    def set_parent(self, entity):
+    def set_parent(self, entity: _Parent):
         """Set the parent Entity object."""
         self.parent = entity
         self._reset_full_id()
@@ -216,7 +228,7 @@ class Entity:
         del self.child_dict[id]
         self.child_list.remove(child)
 
-    def add(self, entity):
+    def add(self, entity: _Child):
         """Add a child to the Entity."""
         entity_id = entity.get_id()
         if self.has_id(entity_id):
@@ -225,7 +237,7 @@ class Entity:
         self.child_list.append(entity)
         self.child_dict[entity_id] = entity
 
-    def insert(self, pos, entity):
+    def insert(self, pos: int, entity: _Child):
         """Add a child to the Entity at a specified position."""
         entity_id = entity.get_id()
         if self.has_id(entity_id):

Bio/PDB/Model.py

@@ -8,8 +8,14 @@
 from Bio.PDB.Entity import Entity
 from Bio.PDB.internal_coords import IC_Chain
 
+from typing import TYPE_CHECKING
 
-class Model(Entity):
+if TYPE_CHECKING:
+    from Bio.PDB.Chain import Chain
+    from Bio.PDB.Structure import Structure
+
+
+class Model(Entity["Structure", "Chain"]):
     """The object representing a model in a structure.
 
     In a structure derived from an X-ray crystallography experiment,

Bio/PDB/NACCESS.py

@@ -200,7 +200,7 @@ class NACCESS_atomic(AbstractAtomPropertyMap):
 
 if __name__ == "__main__":
     import sys
-    from Bio.PDB import PDBParser
+    from Bio.PDB.PDBParser import PDBParser
 
     p = PDBParser()
     s = p.get_structure("X", sys.argv[1])

Bio/PDB/PDBList.py

@@ -484,14 +484,14 @@ class PDBList:
         data = json.dumps(body).encode("utf-8")
         headers = {
             "Content-Type": "application/json; charset=utf-8",
-            "Content-Length": len(data),
+            "Content-Length": str(len(data)),
         }
         request = Request("https://search.rcsb.org/rcsbsearch/v2/query", data, headers)
         with urlopen(request) as response:
             assemblies = json.loads(response.read().decode("utf-8"))["result_set"]
 
         # We transform the assemblies to match the format that they have historically been returned in.
-        def transform(assembly: dict) -> tuple[str, str]:
+        def transform(assembly: dict) -> Tuple[str, str]:
             split = assembly["identifier"].split("-")
             return split[0].lower(), split[-1]
 

Bio/PDB/Residue.py

@@ -10,6 +10,12 @@
 from Bio.PDB.PDBExceptions import PDBConstructionException
 from Bio.PDB.Entity import Entity, DisorderedEntityWrapper
 
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from Bio.PDB.Atom import Atom
+    from Bio.PDB.Chain import Chain
+
 
 _atom_name_dict = {}
 _atom_name_dict["N"] = 1
@@ -18,7 +24,7 @@ _atom_name_dict["C"] = 3
 _atom_name_dict["O"] = 4
 
 
-class Residue(Entity):
+class Residue(Entity["Chain", "Atom"]):
     """Represents a residue. A Residue object stores atoms."""
 
     def __init__(self, id, resname, segid):

Bio/PDB/SASA.py

@@ -18,6 +18,7 @@ Reference:
 
 import collections
 import math
+from typing import MutableMapping
 
 import numpy as np
 
@@ -41,7 +42,7 @@ _ENTITY_HIERARCHY = {
 # References:
 # A. Bondi (1964). "van der Waals Volumes and Radii".
 # M. Mantina, A.C. et al., J. Phys. Chem. 2009, 113, 5806.
-ATOMIC_RADII = collections.defaultdict(lambda: 2.0)
+ATOMIC_RADII: MutableMapping[str, float] = collections.defaultdict(lambda: 2.0)
 ATOMIC_RADII.update(
     {
         "H": 1.200,

Bio/PDB/Structure.py

@@ -9,8 +9,13 @@
 
 from Bio.PDB.Entity import Entity
 
+from typing import TYPE_CHECKING
 
-class Structure(Entity):
+if TYPE_CHECKING:
+    from Bio.PDB.Model import Model
+
+
+class Structure(Entity[None, "Model"]):
     """The Structure class contains a collection of Model instances."""
 
     def __init__(self, id):

Bio/PDB/StructureBuilder.py

@@ -8,6 +8,7 @@
 This is used by the PDBParser and MMCIFparser classes.
 """
 
+from typing import Optional
 import numpy as np
 import warnings
 
@@ -75,7 +76,7 @@ class StructureBuilder:
         """
         self.structure = Structure(structure_id)
 
-    def init_model(self, model_id: int, serial_num: int = None):
+    def init_model(self, model_id: int, serial_num: Optional[int] = None):
         """Create a new Model object with given id.
 
         Arguments:
@@ -185,15 +186,15 @@ class StructureBuilder:
     def init_atom(
         self,
         name: str,
-        coord: np.array,
+        coord: np.ndarray,
         b_factor: float,
         occupancy: float,
         altloc: str,
         fullname: str,
         serial_number=None,
-        element: str = None,
-        pqr_charge: float = None,
-        radius: float = None,
+        element: Optional[str] = None,
+        pqr_charge: Optional[float] = None,
+        radius: Optional[float] = None,
         is_pqr: bool = False,
     ):
         """Create a new Atom object.

Bio/PDB/vectors.py

@@ -388,7 +388,7 @@ Robert T. Miller 2019
 """
 
 
-def homog_rot_mtx(angle_rads: float, axis: str) -> np.array:
+def homog_rot_mtx(angle_rads: float, axis: str) -> np.ndarray:
     """Generate a 4x4 single-axis NumPy rotation matrix.
 
     :param float angle_rads: the desired rotation angle in radians
@@ -459,7 +459,7 @@ def set_X_homog_rot_mtx(angle_rads: float, mtx: np.ndarray):
     mtx[1][2] = -sinang
 
 
-def homog_trans_mtx(x: float, y: float, z: float) -> np.array:
+def homog_trans_mtx(x: float, y: float, z: float) -> np.ndarray:
     """Generate a 4x4 NumPy translation matrix.
 
     :param x, y, z: translation in each axis
@@ -477,7 +477,7 @@ def set_homog_trans_mtx(x: float, y: float, z: float, mtx: np.ndarray):
     mtx[2][3] = z
 
 
-def homog_scale_mtx(scale: float) -> np.array:
+def homog_scale_mtx(scale: float) -> np.ndarray:
     """Generate a 4x4 NumPy scaling matrix.
 
     :param float scale: scale multiplier
@@ -502,17 +502,17 @@ def _get_azimuth(x: float, y: float) -> float:
     )
 
 
-def get_spherical_coordinates(xyz: np.array) -> Tuple[float, float, float]:
+def get_spherical_coordinates(xyz: np.ndarray) -> Tuple[float, float, float]:
     """Compute spherical coordinates (r, azimuth, polar_angle) for X,Y,Z point.
 
     :param array xyz: column vector (3 row x 1 column NumPy array)
     :return: tuple of r, azimuth, polar_angle for input coordinate
     """
-    r = np.linalg.norm(xyz)
+    r = float(np.linalg.norm(xyz))
     if 0 == r:
         return (0, 0, 0)
     azimuth = _get_azimuth(xyz[0], xyz[1])
-    polar_angle = np.arccos(xyz[2] / r)
+    polar_angle: float = np.arccos(xyz[2] / r)
 
     return (r, azimuth, polar_angle)