chore: improve type annotations

.github/workflows/pyright.yml

@@ -0,0 +1,19 @@
+on:
+  - push
+  - pull_request
+
+name: Type checker
+jobs:
+  pyright:
+    name: pyright
+    runs-on: ubuntu-latest
+    steps:
+    - uses: actions/checkout@master
+    - uses: actions/setup-python@v5
+      with:
+        python-version: '3.12'
+    - run: pip install uv
+    - run: uv pip install --system -e .[amber,ase,pymatgen] rdkit openbabel-wheel
+    - uses: jakebailey/pyright-action@v2
+      with:
+        version: 1.1.363

benchmark/test_import.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import subprocess
 import sys
 

docs/conf.py

@@ -11,6 +11,8 @@
 # add these directories to sys.path here. If the directory is relative to the
 # documentation root, use os.path.abspath to make it absolute, like shown here.
 #
+from __future__ import annotations
+
 import os
 import subprocess as sp
 import sys

docs/make_format.py

@@ -1,8 +1,15 @@
+from __future__ import annotations
+
 import csv
 import os
+import sys
 from collections import defaultdict
 from inspect import Parameter, Signature, cleandoc, signature
-from typing import Literal
+
+if sys.version_info >= (3, 8):
+    from typing import Literal
+else:
+    from typing_extensions import Literal
 
 from numpydoc.docscrape import Parameter as numpydoc_Parameter
 from numpydoc.docscrape_sphinx import SphinxDocString

docs/nb/try_dpdata.ipynb

@@ -13,6 +13,8 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "from __future__ import annotations\n",
+    "\n",
     "import dpdata"
    ]
   },

dpdata/__about__.py

@@ -1 +1,3 @@
+from __future__ import annotations
+
 __version__ = "unknown"

dpdata/__init__.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from . import lammps, md, vasp
 from .bond_order_system import BondOrderSystem
 from .system import LabeledSystem, MultiSystems, System

dpdata/__main__.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from dpdata.cli import dpdata_cli
 
 if __name__ == "__main__":

dpdata/abacus/md.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import os
 import warnings
 

dpdata/abacus/relax.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import os
 
 import numpy as np

dpdata/abacus/scf.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import os
 import re
 import warnings

dpdata/amber/mask.py

@@ -1,5 +1,7 @@
 """Amber mask."""
 
+from __future__ import annotations
+
 try:
     import parmed
 except ImportError:

dpdata/amber/md.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import os
 import re
 

dpdata/amber/sqm.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import numpy as np
 
 from dpdata.periodic_table import ELEMENTS

dpdata/ase_calculator.py

@@ -1,4 +1,6 @@
-from typing import TYPE_CHECKING, List, Optional
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
 
 from ase.calculators.calculator import (  # noqa: TID253
     Calculator,
@@ -23,7 +25,10 @@ class DPDataCalculator(Calculator):
         dpdata driver
     """
 
-    name = "dpdata"
+    @property
+    def name(self) -> str:
+        return "dpdata"
+
     implemented_properties = ["energy", "free_energy", "forces", "virial", "stress"]
 
     def __init__(self, driver: Driver, **kwargs) -> None:
@@ -32,9 +37,9 @@ def __init__(self, driver: Driver, **kwargs) -> None:
 
     def calculate(
         self,
-        atoms: Optional["Atoms"] = None,
-        properties: List[str] = ["energy", "forces"],
-        system_changes: List[str] = all_changes,
+        atoms: Atoms | None = None,
+        properties: list[str] = ["energy", "forces"],
+        system_changes: list[str] = all_changes,
     ):
         """Run calculation with a driver.
 
@@ -48,10 +53,10 @@ def calculate(
         system_changes : List[str], optional
             unused, only for function signature compatibility, by default all_changes
         """
-        if atoms is not None:
-            self.atoms = atoms.copy()
+        assert atoms is not None
+        atoms = atoms.copy()
 
-        system = dpdata.System(self.atoms, fmt="ase/structure")
+        system = dpdata.System(atoms, fmt="ase/structure")
         data = system.predict(driver=self.driver).data
 
         self.results["energy"] = data["energies"][0]

dpdata/bond_order_system.py

@@ -1,5 +1,7 @@
 # %%
 # Bond Order System
+from __future__ import annotations
+
 from copy import deepcopy
 
 import numpy as np
@@ -96,13 +98,14 @@
         mol = fmtobj.from_bond_order_system(file_name, **kwargs)
         self.from_rdkit_mol(mol)
         if hasattr(fmtobj.from_bond_order_system, "post_func"):
-            for post_f in fmtobj.from_bond_order_system.post_func:
+            for post_f in fmtobj.from_bond_order_system.post_func:  # type: ignore
                 self.post_funcs.get_plugin(post_f)(self)
         return self
 
     def to_fmt_obj(self, fmtobj, *args, **kwargs):
         from rdkit.Chem import Conformer
 
+        assert self.rdkit_mol is not None
         self.rdkit_mol.RemoveAllConformers()
         for ii in range(self.get_nframes()):
             conf = Conformer()
@@ -145,9 +148,9 @@
         """Return the formal charges on each atom."""
         return self.data["formal_charges"]
 
-    def copy(self):
+    def copy(self):  # type: ignore
         new_mol = deepcopy(self.rdkit_mol)
-        self.__class__(data=deepcopy(self.data), rdkit_mol=new_mol)
+        return self.__class__(data=deepcopy(self.data), rdkit_mol=new_mol)
 
     def __add__(self, other):
         raise NotImplementedError(

dpdata/cli.py

@@ -1,7 +1,8 @@
 """Command line interface for dpdata."""
 
+from __future__ import annotations
+
 import argparse
-from typing import Optional
 
 from . import __version__
 from .system import LabeledSystem, MultiSystems, System
@@ -59,11 +60,11 @@
     *,
     from_file: str,
     from_format: str = "auto",
-    to_file: Optional[str] = None,
-    to_format: Optional[str] = None,
+    to_file: str | None = None,
+    to_format: str | None = None,
     no_labeled: bool = False,
     multi: bool = False,
-    type_map: Optional[list] = None,
+    type_map: list | None = None,
     **kwargs,
 ):
     """Convert files from one format to another one.

dpdata/cp2k/cell.py

@@ -1,4 +1,5 @@
 # %%
+from __future__ import annotations
 
 import numpy as np
 

dpdata/cp2k/output.py

@@ -1,4 +1,6 @@
 # %%
+from __future__ import annotations
+
 import math
 import re
 from collections import OrderedDict

dpdata/data_type.py

@@ -1,5 +1,7 @@
+from __future__ import annotations
+
 from enum import Enum, unique
-from typing import TYPE_CHECKING, Tuple
+from typing import TYPE_CHECKING
 
 import numpy as np
 
@@ -50,16 +52,17 @@ def __init__(
         self,
         name: str,
         dtype: type,
-        shape: Tuple[int, Axis] = None,
+        shape: tuple[int | Axis, ...] | None = None,
         required: bool = True,
     ) -> None:
         self.name = name
         self.dtype = dtype
         self.shape = shape
         self.required = required
 
-    def real_shape(self, system: "System") -> Tuple[int]:
+    def real_shape(self, system: System) -> tuple[int]:
         """Returns expected real shape of a system."""
+        assert self.shape is not None
         shape = []
         for ii in self.shape:
             if ii is Axis.NFRAMES:
@@ -70,7 +73,7 @@ def real_shape(self, system: "System") -> Tuple[int]:
                 shape.append(system.get_natoms())
             elif ii is Axis.NBONDS:
                 # BondOrderSystem
-                shape.append(system.get_nbonds())
+                shape.append(system.get_nbonds())  # type: ignore
             elif ii == -1:
                 shape.append(AnyInt(-1))
             elif isinstance(ii, int):
@@ -79,7 +82,7 @@ def real_shape(self, system: "System") -> Tuple[int]:
                 raise RuntimeError("Shape is not an int!")
         return tuple(shape)
 
-    def check(self, system: "System"):
+    def check(self, system: System):
         """Check if a system has correct data of this type.
 
         Parameters

dpdata/deepmd/comp.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import glob
 import os
 import shutil

dpdata/deepmd/mixed.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import glob
 import os
 import shutil

dpdata/deepmd/raw.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import os
 import warnings
 

dpdata/dftbplus/output.py

@@ -1,9 +1,9 @@
-from typing import Tuple
+from __future__ import annotations
 
 import numpy as np
 
 
-def read_dftb_plus(fn_1: str, fn_2: str) -> Tuple[str, np.ndarray, float, np.ndarray]:
+def read_dftb_plus(fn_1: str, fn_2: str) -> tuple[str, np.ndarray, float, np.ndarray]:
     """Read from DFTB+ input and output.
 
     Parameters

dpdata/driver.py

@@ -1,12 +1,14 @@
 """Driver plugin system."""
 
+from __future__ import annotations
+
 from abc import ABC, abstractmethod
-from typing import TYPE_CHECKING, Callable, List, Union
+from typing import TYPE_CHECKING, Callable
 
 from .plugin import Plugin
 
 if TYPE_CHECKING:
-    import ase
+    import ase.calculators.calculator
 
 
 class Driver(ABC):
@@ -43,7 +45,7 @@
         return Driver.__DriverPlugin.register(key)
 
     @staticmethod
-    def get_driver(key: str) -> "Driver":
+    def get_driver(key: str) -> type[Driver]:
         """Get a driver plugin.
 
         Parameters
@@ -97,7 +99,7 @@
         return NotImplemented
 
     @property
-    def ase_calculator(self) -> "ase.calculators.calculator.Calculator":
+    def ase_calculator(self) -> ase.calculators.calculator.Calculator:
         """Returns an ase calculator based on this driver."""
         from .ase_calculator import DPDataCalculator
 
@@ -130,7 +132,7 @@
     This driver is the hybrid of SQM and DP.
     """
 
-    def __init__(self, drivers: List[Union[dict, Driver]]) -> None:
+    def __init__(self, drivers: list[dict | Driver]) -> None:
         self.drivers = []
         for driver in drivers:
             if isinstance(driver, Driver):
@@ -157,6 +159,7 @@
         dict
             labeled data with energies and forces
         """
+        labeled_data = {}
         for ii, driver in enumerate(self.drivers):
             lb_data = driver.label(data.copy())
             if ii == 0:
@@ -199,7 +202,7 @@
         return Minimizer.__MinimizerPlugin.register(key)
 
     @staticmethod
-    def get_minimizer(key: str) -> "Minimizer":
+    def get_minimizer(key: str) -> type[Minimizer]:
         """Get a minimizer plugin.
 
         Parameters

dpdata/fhi_aims/output.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import re
 import warnings
 

dpdata/format.py

@@ -1,5 +1,7 @@
 """Implement the format plugin system."""
 
+from __future__ import annotations
+
 import os
 from abc import ABC
 
@@ -163,7 +165,7 @@
             if not isinstance(func_name, (list, tuple, set)):
                 object.post_func = (func_name,)
             else:
-                object.post_func = func_name
+                object.post_func = tuple(func_name)
             return object
 
         return decorator