fixup! Add type annotations to code base

Author: padix-key

src/biotite/application/application.py

@@ -2,6 +2,8 @@
 # under the 3-Clause BSD License. Please see 'LICENSE.rst' for further
 # information.
 
+from __future__ import annotations
+
 __name__ = "biotite.application"
 __author__ = "Patrick Kunzmann"
 __all__ = [
@@ -15,8 +17,13 @@
 
 import abc
 import time
+from collections.abc import Callable
 from enum import Flag, auto
 from functools import wraps
+from typing import ParamSpec, TypeVar
+
+_P = ParamSpec("_P")
+_R = TypeVar("_R")
 
 
 class AppState(Flag):
@@ -31,7 +38,9 @@ class AppState(Flag):
     CANCELLED = auto()
 
 
-def requires_state(app_state):
+def requires_state(
+    app_state: AppState,
+) -> Callable[[Callable[_P, _R]], Callable[_P, _R]]:
     """
     A decorator for methods of :class:`Application` subclasses that
     raises an :class:`AppStateError` in case the method is called, when
@@ -53,14 +62,16 @@ def requires_state(app_state):
     ...     pass
     """
 
-    def decorator(func):
+    def decorator(func: Callable[_P, _R]) -> Callable[_P, _R]:
         @wraps(func)
-        def wrapper(*args, **kwargs):
+        def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> _R:
             # First parameter of method is always 'self'
             try:
                 instance = args[0]
             except IndexError:
                 raise TypeError("This method must be called from a class instance")
+            if not isinstance(instance, Application):
+                raise TypeError("This method must be an 'Application' method")
             if not instance._state & app_state:
                 raise AppStateError(
                     f"The application is in {instance.get_app_state()} state, "
@@ -115,21 +126,21 @@ class Application(metaclass=abc.ABCMeta):
     executed, while the application runs in the background.
     """
 
-    def __init__(self):
-        self._state = AppState.CREATED
+    def __init__(self) -> None:
+        self._state: AppState = AppState.CREATED
 
     @requires_state(AppState.CREATED)
-    def start(self):
+    def start(self) -> None:
         """
         Start the application run and set its state to *RUNNING*.
         This can only be done from the *CREATED* state.
         """
         self.run()
-        self._start_time = time.time()
+        self._start_time: float = time.time()
         self._state = AppState.RUNNING
 
     @requires_state(AppState.RUNNING | AppState.FINISHED)
-    def join(self, timeout=None):
+    def join(self, timeout: float | None = None) -> None:
         """
         Conclude the application run and set its state to *JOINED*.
         This can only be done from the *RUNNING* or *FINISHED* state.
@@ -174,14 +185,14 @@ def join(self, timeout=None):
         self.clean_up()
 
     @requires_state(AppState.RUNNING | AppState.FINISHED)
-    def cancel(self):
+    def cancel(self) -> None:
         """
         Cancel the application when in *RUNNING* or *FINISHED* state.
         """
         self._state = AppState.CANCELLED
         self.clean_up()
 
-    def get_app_state(self):
+    def get_app_state(self) -> AppState:
         """
         Get the current app state.
 
@@ -196,7 +207,7 @@ def get_app_state(self):
         return self._state
 
     @abc.abstractmethod
-    def run(self):
+    def run(self) -> None:
         """
         Commence the application run. Called in :func:`start()`.
 
@@ -205,7 +216,7 @@ def run(self):
         pass
 
     @abc.abstractmethod
-    def is_finished(self):
+    def is_finished(self) -> bool:
         """
         Check if the application has finished.
 
@@ -219,7 +230,7 @@ def is_finished(self):
         pass
 
     @abc.abstractmethod
-    def wait_interval(self):
+    def wait_interval(self) -> float:
         """
         The time interval of :func:`is_finished()` calls in the joining
         process.
@@ -235,15 +246,15 @@ def wait_interval(self):
         pass
 
     @abc.abstractmethod
-    def evaluate(self):
+    def evaluate(self) -> None:
         """
         Evaluate application results. Called in :func:`join()`.
 
         PROTECTED: Override when inheriting.
         """
         pass
 
-    def clean_up(self):
+    def clean_up(self) -> None:
         """
         Do clean up work after the application terminates.
 

src/biotite/application/autodock/app.py

@@ -2,19 +2,25 @@
 # under the 3-Clause BSD License. Please see 'LICENSE.rst' for further
 # information.
 
+from __future__ import annotations
+
 __name__ = "biotite.application.autodock"
 __author__ = "Patrick Kunzmann"
 __all__ = ["VinaApp"]
 
 import copy
+from os import PathLike
 from tempfile import NamedTemporaryFile
+from typing import Any
 import numpy as np
 from biotite.application.application import AppState, requires_state
 from biotite.application.localapp import LocalApp, cleanup_tempfile
+from biotite.structure.atoms import AtomArray, AtomArrayStack
 from biotite.structure.connect import find_connected
 from biotite.structure.error import BadStructureError
 from biotite.structure.io.pdbqt import PDBQTFile
 from biotite.structure.residues import get_residue_masks, get_residue_starts_for
+from biotite.typing import XYZ, K, M, N, NDArray1, NDArray3
 
 
 class VinaApp(LocalApp):
@@ -24,7 +30,7 @@ class VinaApp(LocalApp):
     Parameters
     ----------
     ligand : AtomArray
-        The structure of the receptor molecule.
+        The structure of the ligand molecule.
         Must have an associated :class:`BondList`.
         An associated ``charge`` annotation is recommended for proper
         calculation of partial charges.
@@ -63,25 +69,33 @@ class VinaApp(LocalApp):
     ... )
     """
 
-    def __init__(self, ligand, receptor, center, size, flexible=None, bin_path="vina"):
+    def __init__(
+        self,
+        ligand: AtomArray[Any],
+        receptor: AtomArray[N],
+        center: NDArray1[XYZ, np.floating],
+        size: NDArray1[XYZ, np.floating],
+        flexible: NDArray1[N, np.bool_] | None = None,
+        bin_path: PathLike[str] | str = "vina",
+    ) -> None:
         super().__init__(bin_path)
 
         if ligand.bonds is None:
             raise ValueError("The ligand has no associated BondList")
         if receptor.bonds is None:
             raise ValueError("The receptor has no associated BondList")
 
-        self._ligand = ligand.copy()
-        self._receptor = receptor.copy()
-        self._center = copy.deepcopy(center)
-        self._size = copy.deepcopy(size)
-        self._is_flexible = flexible is not None
-        self._seed = None
-        self._exhaustiveness = None
-        self._number = None
-        self._energy_range = None
-
-        if self._is_flexible:
+        self._ligand: AtomArray[Any] = ligand.copy()
+        self._receptor: AtomArray[Any] = receptor.copy()
+        self._center: NDArray1[Any, np.floating] = copy.deepcopy(center)
+        self._size: NDArray1[Any, np.floating] = copy.deepcopy(size)
+        self._is_flexible: bool = flexible is not None
+        self._seed: int | None = None
+        self._exhaustiveness: int | None = None
+        self._number: int | None = None
+        self._energy_range: float | None = None
+
+        if flexible is not None:
             flexible_indices = np.where(flexible)[0]
             self._flex_res_starts = np.unique(
                 get_residue_starts_for(receptor, flexible_indices)
@@ -95,7 +109,7 @@ def __init__(self, ligand, receptor, center, size, flexible=None, bin_path="vina
         self._out_file = NamedTemporaryFile("r", suffix=".pdbqt", delete=False)
 
     @requires_state(AppState.CREATED)
-    def set_seed(self, seed):
+    def set_seed(self, seed: int) -> None:
         """
         Fix the seed for the random number generator to get
         reproducible results.
@@ -110,7 +124,7 @@ def set_seed(self, seed):
         self._seed = seed
 
     @requires_state(AppState.CREATED)
-    def set_exhaustiveness(self, exhaustiveness):
+    def set_exhaustiveness(self, exhaustiveness: int) -> None:
         """
         Set the *exhaustiveness* parameter for *Vina*.
 
@@ -127,7 +141,7 @@ def set_exhaustiveness(self, exhaustiveness):
         self._exhaustiveness = exhaustiveness
 
     @requires_state(AppState.CREATED)
-    def set_max_number_of_models(self, number):
+    def set_max_number_of_models(self, number: int) -> None:
         """
         Set the maximum number of binding modes to generate.
 
@@ -143,7 +157,7 @@ def set_max_number_of_models(self, number):
         self._number = number
 
     @requires_state(AppState.CREATED)
-    def set_energy_range(self, energy_range):
+    def set_energy_range(self, energy_range: float) -> None:
         """
         Set the maximum energy range of the generated models.
 
@@ -158,7 +172,7 @@ def set_energy_range(self, energy_range):
         """
         self._energy_range = energy_range
 
-    def run(self):
+    def run(self) -> None:
         # Use different atom ID ranges for atoms in ligand and receptor
         # for unambiguous assignment, if the receptor contains flexible
         # residues
@@ -258,7 +272,7 @@ def run(self):
         self.set_arguments(arguments)
         super().run()
 
-    def evaluate(self):
+    def evaluate(self) -> None:
         super().evaluate()
         out_file = PDBQTFile.read(self._out_file)
 
@@ -276,15 +290,15 @@ def evaluate(self):
             [float(remark[12:].split()[0]) for remark in remarks]
         )
 
-    def clean_up(self):
+    def clean_up(self) -> None:
         super().clean_up()
         cleanup_tempfile(self._ligand_file)
         cleanup_tempfile(self._receptor_file)
         cleanup_tempfile(self._receptor_flex_file)
         cleanup_tempfile(self._out_file)
 
     @requires_state(AppState.JOINED)
-    def get_energies(self):
+    def get_energies(self) -> NDArray1[K, np.floating]:
         """
         Get the predicted binding energy for each generated binding
         mode.
@@ -298,7 +312,7 @@ def get_energies(self):
         return self._energies
 
     @requires_state(AppState.JOINED)
-    def get_ligand_models(self):
+    def get_ligand_models(self) -> AtomArrayStack[Any, Any]:
         """
         Get the ligand structure with the conformations for each
         generated binding mode.
@@ -321,7 +335,7 @@ def get_ligand_models(self):
         return self._ligand_models
 
     @requires_state(AppState.JOINED)
-    def get_ligand_coord(self):
+    def get_ligand_coord(self) -> NDArray3[M, N, XYZ, np.floating]:
         """
         Get the ligand coordinates for each generated binding mode.
 
@@ -342,7 +356,7 @@ def get_ligand_coord(self):
         return coord
 
     @requires_state(AppState.JOINED)
-    def get_flexible_residue_models(self):
+    def get_flexible_residue_models(self) -> AtomArrayStack[Any, Any]:
         """
         Get the structure for the flexible side chains with the
         conformations for each generated binding mode.
@@ -368,7 +382,7 @@ def get_flexible_residue_models(self):
         return self._flex_models
 
     @requires_state(AppState.JOINED)
-    def get_receptor_coord(self):
+    def get_receptor_coord(self) -> NDArray3[M, N, XYZ, np.floating]:
         """
         Get the get_receptor_coord coordinates for each generated
         binding mode.
@@ -401,7 +415,9 @@ def get_receptor_coord(self):
             coord[:, self._receptor_mask] = self._flex_models.coord
         return coord
 
-    def _get_flexible_residue(self, residue_start):
+    def _get_flexible_residue(
+        self, residue_start: int
+    ) -> tuple[np.ndarray, np.ndarray, int]:
         residue_indices = np.where(
             get_residue_masks(self._receptor, [residue_start])[0]
         )[0]
@@ -417,6 +433,8 @@ def _get_flexible_residue(self, residue_start):
 
         # Find the index of the atom connected to root on the flexible
         # side chain (CB)
+        if self._receptor.bonds is None:
+            raise ValueError("The receptor has no associated BondList")
         root_connect_indices, _ = self._receptor.bonds.get_bonds(root_index)
         connected_index = None
         try:
@@ -450,7 +468,14 @@ def _get_flexible_residue(self, residue_start):
         return flex_mask, rigid_mask, root_index
 
     @staticmethod
-    def dock(ligand, receptor, center, size, flexible=None, bin_path="vina"):
+    def dock(
+        ligand: AtomArray[Any],
+        receptor: AtomArray[N],
+        center: NDArray1[XYZ, np.floating],
+        size: NDArray1[XYZ, np.floating],
+        flexible: NDArray1[N, np.bool_] | None = None,
+        bin_path: PathLike[str] | str = "vina",
+    ) -> tuple[NDArray3[M, N, XYZ, np.floating], NDArray1[K, np.floating]]:
         """
         Dock a ligand to a receptor molecule using *AutoDock Vina*.
 
@@ -460,7 +485,7 @@ def dock(ligand, receptor, center, size, flexible=None, bin_path="vina"):
         Parameters
         ----------
         ligand : AtomArray
-            The structure of the receptor molecule.
+            The structure of the ligand molecule.
             Must have an associated :class:`BondList`.
             An associated ``charge`` annotation is recommended for proper
             calculation of partial charges.