Merge pull request #259 from SCM-NV/annotation

TYP: Update type annotations

src/qmflows/cp2k_utils.py

@@ -27,12 +27,11 @@
 
 import copy
 import textwrap
-from io import TextIOBase
 from functools import singledispatch
 from itertools import repeat, islice
 from collections import abc
 from typing import (Union, Optional, List, Dict, Tuple, MutableMapping, NoReturn,
-                    Sequence, Any, Iterable, Iterator, overload, cast)
+                    Sequence, Any, Iterable, Iterator, overload, cast, IO)
 
 import numpy as np
 import pandas as pd
@@ -90,13 +89,13 @@ class LengthError(ValueError):
 
 @to_runtime_error
 def _map_psf_atoms(settings: None, key: str,
-                   value: Union[PathLike, TextIOBase],
+                   value: Union[PathLike, IO[Any]],
                    mol: None, **kwargs: Any) -> Dict[str, str]:
     """A small wrapper around :func:`map_psf_atoms`."""
     return map_psf_atoms(value, **kwargs)
 
 
-def map_psf_atoms(file: Union[PathLike, TextIOBase], **kwargs: Any) -> Dict[str, str]:
+def map_psf_atoms(file: Union[PathLike, IO[Any]], **kwargs: Any) -> Dict[str, str]:
     r"""Take a .psf file and construct a :class:`dict` mapping atom types to atom names.
 
     Examples

src/qmflows/fileFunctions.py

@@ -2,7 +2,7 @@
 
 __all__ = ['yaml2Settings']
 
-from typing import Union, Callable
+from typing import Union, Callable, overload, Dict, Any
 
 import yaml
 
@@ -11,10 +11,12 @@
 from .yaml_utils import UniqueSafeLoader
 
 
-def yaml2Settings(
-    xs: Union[str, bytes],
-    mapping_type: Callable[[dict], T] = Settings,
-) -> T:
+@overload
+def yaml2Settings(xs: Union[str, bytes]) -> Settings: ...
+@overload
+def yaml2Settings(xs: Union[str, bytes], mapping_type: Callable[[Dict[str, Any]], T]) -> T: ...
+
+def yaml2Settings(xs, mapping_type=Settings):
     """Transform a string containing some data in .yaml format to a Settings object."""
     if isinstance(xs, bytes):
         xs = xs.decode()

src/qmflows/packages/SCM.py

@@ -102,35 +102,38 @@ class ADF_Result(Result):
     def __init__(self, settings: Optional[Settings],
                  molecule: Optional[plams.Molecule],
                  job_name: str,
-                 dill_path: Union[None, str, os.PathLike] = None,
-                 plams_dir: Union[None, str, os.PathLike] = None,
-                 work_dir: Union[None, str, os.PathLike] = None,
+                 dill_path: "None | str | os.PathLike[str]" = None,
+                 plams_dir: "None | str | os.PathLike[str]" = None,
+                 work_dir: "None | str | os.PathLike[str]" = None,
                  status: str = 'done',
                  warnings: Optional[WarnMap] = None) -> None:
         # Load available property parser from yaml file.
         super().__init__(settings, molecule, job_name, dill_path,
                          plams_dir=plams_dir, status=status, warnings=warnings)
 
         # Create a KF reader instance
-        if work_dir is not None:
+        if work_dir is not None and plams_dir is not None:
             # The t21 path has to be absolute: use workdir instead of plams_dir
             name_t21 = basename(normpath(work_dir))
             path_t21 = join(plams_dir, f'{name_t21}.t21')
-            self.kf = plams.KFFile(path_t21)
+            self.kf: "None | plams.KFFile" = plams.KFFile(path_t21)
         else:
             self.kf = None
 
-    def get_property_kf(self, prop: str, section: Optional[str] = None) -> Any:
+    def get_property_kf(self, prop: str, section: Optional[str] = None) -> Optional[Any]:
         """Interface for :meth:`plams.KFFile.read()<scm.plams.tools.kftools.KFFile.read>`."""
-        return self.kf.read(section, prop)
+        if self.kf is None:
+            return None
+        else:
+            return self.kf.read(section, prop)
 
     @property
     def molecule(self) -> Optional[plams.Molecule]:
         """WARNING: Cheap copy from PLAMS, do not keep this!!!."""
-        try:
+        if self._molecule is None or self.kf is None:
+            return None
+        else:
             m = self._molecule.copy()
-        except AttributeError:
-            return None  # self._molecule can be None
         # Find out correct location
         coords = self.kf.read('Geometry', 'xyz InputOrder')
         coords = np.array([coords[i: i + 3] for i in range(0, len(coords), 3)])
@@ -151,9 +154,9 @@ class DFTB_Result(Result):
     def __init__(self, settings: Optional[Settings],
                  molecule: Optional[plams.Molecule],
                  job_name: str,
-                 dill_path: Union[None, str, os.PathLike] = None,
-                 plams_dir: Union[None, str, os.PathLike] = None,
-                 work_dir: Union[None, str, os.PathLike] = None,
+                 dill_path: "None | str | os.PathLike[str]" = None,
+                 plams_dir: "None | str | os.PathLike[str]" = None,
+                 work_dir: "None | str | os.PathLike[str]" = None,
                  status: str = 'done',
                  warnings: Optional[WarnMap] = None) -> None:
         # Read available propiety parsers from a yaml file
@@ -163,17 +166,17 @@ def __init__(self, settings: Optional[Settings],
         if plams_dir is not None:
             kf_filename = join(plams_dir, 'dftb.rkf')
             # create a kf reader instance
-            self.kf = plams.KFFile(kf_filename)
+            self.kf: "None | plams.KFFile" = plams.KFFile(kf_filename)
         else:
             self.kf = None
 
     @property
-    def molecule(self) -> plams.Molecule:
+    def molecule(self) -> Optional[plams.Molecule]:
         """Read molecule from output."""
-        try:
+        if self._molecule is None or self.kf is None:
+            return None
+        else:
             m = self._molecule.copy()
-        except AttributeError:
-            return None  # self._molecule can be None
         coords = self.kf.read('Molecule', 'Coords')
         coords = np.array([coords[i: i + 3] for i in range(0, len(coords), 3)])
         m.from_array(coords)
@@ -197,7 +200,7 @@ class ADF(Package):
     """
 
     generic_mapping: ClassVar[_Settings] = load_properties('ADF', prefix='generic2')
-    result_type: ClassVar[Type[Result]] = ADF_Result
+    result_type: ClassVar[Type[ADF_Result]] = ADF_Result
 
     def __init__(self, pkg_name: str = "adf") -> None:
         super().__init__(pkg_name)
@@ -264,7 +267,7 @@ class DFTB(Package):
     """:class:`~qmflows.packages.packages.Package` subclass for DFTB."""
 
     generic_mapping: ClassVar[_Settings] = load_properties('DFTB', prefix='generic2')
-    result_type: ClassVar[Type[Result]] = DFTB_Result
+    result_type: ClassVar[Type[DFTB_Result]] = DFTB_Result
 
     def __init__(self, pkg_name: str = "dftb") -> None:
         super().__init__(pkg_name)
@@ -292,7 +295,7 @@ def run_job(cls, settings: Settings, mol: plams.Molecule, job_name: str,
         try:
             result = job.run()
             name = result.job.name
-            path = result.job.path
+            path: "None | str" = result.job.path
         except struct.error:
             job.status = 'failed'
             name = job_name

src/qmflows/packages/cp2k_mm.py

@@ -14,12 +14,12 @@
 
 import os
 from os.path import join, abspath
-from typing import Union, Any, ClassVar, Dict, Type
+from typing import Union, Any, ClassVar, Dict, Type, TYPE_CHECKING
 
 import numpy as np
 from scm import plams
 
-from .packages import Result, parse_output_warnings, load_properties
+from .packages import parse_output_warnings, load_properties
 from .cp2k_package import CP2K, CP2K_Result
 from ..cp2k_utils import set_prm, _map_psf_atoms, CP2K_KEYS_ALIAS
 from ..parsers.cp2KParser import parse_cp2k_warnings
@@ -47,7 +47,7 @@ class CP2KMM(CP2K):
     """  # noqa: E501
 
     generic_mapping: ClassVar[_Settings] = load_properties('CP2KMM', prefix='generic2')
-    result_type: ClassVar[Type[Result]] = CP2KMM_Result
+    result_type: ClassVar[Type[CP2KMM_Result]] = CP2KMM_Result
 
     def __init__(self, pkg_name: str = "cp2k") -> None:
         super().__init__(pkg_name)
@@ -90,6 +90,18 @@ def prerun(self, settings: Settings, mol: plams.Molecule, **kwargs: Any) -> None
             settings.prm = str(prm_name)
         """
 
+    if TYPE_CHECKING:
+        @classmethod
+        def run_job(
+            cls,
+            settings: Settings,
+            mol: plams.Molecule,
+            job_name: str = 'cp2k_job',
+            work_dir: "None | str | os.PathLike[str]" = ...,
+            validate_output: bool = True,
+            **kwargs: Any,
+        ) -> CP2KMM_Result: ...
+
     @classmethod
     def handle_special_keywords(cls, settings: Settings, key: str,
                                 value: Any, mol: plams.Molecule) -> None:
@@ -119,7 +131,7 @@ def handle_special_keywords(cls, settings: Settings, key: str,
 
     @staticmethod
     def _parse_psf(settings: Settings, key: str,
-                   value: Any, mol: plams.Molecule) -> None:
+                   value: "None", mol: plams.Molecule) -> None:
         """Assign a .psf file."""
         subsys = settings.specific.cp2k.force_eval.subsys
         if value is None:

src/qmflows/packages/cp2k_package.py

@@ -24,7 +24,7 @@ class CP2K_Result(Result):
     prop_mapping: ClassVar[_Settings] = load_properties('CP2K', prefix='properties')
 
     @property
-    def molecule(self) -> plams.Molecule:
+    def molecule(self) -> "None | plams.Molecule":
         """Return the current geometry.
 
         If the job is an optimization, try to read the ` *-pos-1.xyz` file.
@@ -47,15 +47,15 @@ class CP2K(Package):
     """
 
     generic_mapping: ClassVar[_Settings] = load_properties('CP2K', prefix='generic2')
-    result_type: ClassVar[Type[Result]] = CP2K_Result
+    result_type: ClassVar[Type[CP2K_Result]] = CP2K_Result
 
     def __init__(self, pkg_name: str = "cp2k") -> None:
         super().__init__(pkg_name)
 
     @classmethod
     def run_job(cls, settings: Settings, mol: plams.Molecule,
                 job_name: str = 'cp2k_job',
-                work_dir: Union[None, str, os.PathLike] = None,
+                work_dir: "None | str | os.PathLike[str]" = None,
                 validate_output: bool = True,
                 **kwargs: Any) -> CP2K_Result:
         """Call the Cp2K binary using plams interface.

src/qmflows/packages/orca.py

@@ -31,11 +31,10 @@ def molecule(self) -> Optional[plams.Molecule]:
             return None
 
         plams_dir = self.archive["plams_dir"]
-        try:
-            file_name = join(plams_dir, f'{self.job_name}.out')
-        except TypeError:  # plams_dir can be None
+        if plams_dir is None:
             return None
         else:
+            file_name = join(plams_dir, f'{self.job_name}.out')
             return parse_molecule(file_name, self._molecule)
 
 
@@ -49,15 +48,15 @@ class ORCA(Package):
     """
 
     generic_mapping: ClassVar[_Settings] = load_properties('ORCA', prefix='generic2')
-    result_type: ClassVar[Type[Result]] = ORCA_Result
+    result_type: ClassVar[Type[ORCA_Result]] = ORCA_Result
 
     def __init__(self, pkg_name: str = "orca") -> None:
         super().__init__(pkg_name)
 
     @classmethod
     def run_job(cls, settings: Settings, mol: plams.Molecule,
                 job_name: str = "ORCAjob",
-                work_dir: Union[None, str, os.PathLike] = None,
+                work_dir: "None | str | os.PathLike[str]" = None,
                 validate_output: bool = True,
                 **kwargs: Any) -> ORCA_Result:
 
@@ -92,10 +91,10 @@ def inithess(value: Any) -> Settings:
             value = value if isinstance(value, np.ndarray) else np.array(value)
 
             def format_atom(atom: plams.Atom) -> str:
-                symbol, mass, coords = atom.symbol, atom._getmass(), atom.coords
+                symbol, mass, coords = atom.symbol, atom.mass, atom.coords
                 return '{:2s}{:12.4f}{:14.6f}{:14.6f}{:14.6f}\n'.format(symbol, mass, *coords)
 
-            def format_hessian(dim, hess: Union[List[List[float]], np.array]) -> str:
+            def format_hessian(dim, hess: Union[List[List[float]], np.ndarray]) -> str:
                 """Format numpy array to Orca matrix format."""
                 ret = ''
                 for i in range((dim - 1) // 6 + 1):

src/qmflows/packages/package_wrapper.py

@@ -184,7 +184,7 @@ class PackageWrapper(Package, Generic[JT]):
     """  # noqa
 
     generic_mapping: ClassVar[_Settings] = load_properties('PackageWrapper', prefix='generic2')
-    result_type: ClassVar[Type[Result]] = ResultWrapper
+    result_type: ClassVar[Type[ResultWrapper]] = ResultWrapper
     job_type: Type[JT]
 
     def __init__(self, job_type: Type[JT], name: Optional[str] = None) -> None:

src/qmflows/packages/packages.py

@@ -64,9 +64,9 @@ class Result:
     def __init__(self, settings: Optional[Settings],
                  molecule: Optional[plams.Molecule],
                  job_name: str,
-                 dill_path: Union[None, str, os.PathLike] = None,
-                 plams_dir: Union[None, str, os.PathLike] = None,
-                 work_dir: Union[None, str, os.PathLike] = None,
+                 dill_path: "None | str | os.PathLike[str]" = None,
+                 plams_dir: "None | str | os.PathLike[str]" = None,
+                 work_dir: "None | str | os.PathLike[str]" = None,
                  status: str = 'done',
                  warnings: Optional[WarnMap] = None) -> None:
         """Initialize a :class:`Result` instance.
@@ -681,15 +681,17 @@ def import_parser(ds: Mapping[str, str], module_root: str = "qmflows.parsers") -
     return importlib.import_module(module_name)
 
 
-def find_file_pattern(path: Union[str, os.PathLike],
-                      folder: Union[None, str, os.PathLike] = None) -> Iterator[str]:
+def find_file_pattern(
+    path: "str | os.PathLike[str]",
+    folder: "None | str | os.PathLike[str]" = None,
+) -> Iterator[str]:
     if folder is not None and os.path.exists(folder):
         return map(lambda x: join(folder, x), fnmatch.filter(os.listdir(folder), str(path)))
     else:
         return iter([])
 
 
-def ignore_unused_kwargs(fun: Callable, *args: Any, **kwargs: Any) -> Any:
+def ignore_unused_kwargs(fun: Callable[..., Any], *args: Any, **kwargs: Any) -> Any:
     """Inspect the signature of function `fun` and filter the keyword arguments.
 
     Searches for the keyword arguments that are present in both `**kwargs`