Merge pull request #67 from BoothGroup/typing

Restructure package

.github/workflows/ci.yaml

@@ -15,10 +15,10 @@ jobs:
       fail-fast: false
       matrix:
         include:
-          - {python-version: "3.8", os: ubuntu-latest, documentation: True}
-          - {python-version: "3.9", os: ubuntu-latest, documentation: False}
+          - {python-version: "3.9", os: ubuntu-latest, documentation: True}
           - {python-version: "3.10", os: ubuntu-latest, documentation: False}
           - {python-version: "3.11", os: ubuntu-latest, documentation: False}
+          - {python-version: "3.12", os: ubuntu-latest, documentation: False}
 
     steps:
       - uses: actions/checkout@v2
@@ -38,6 +38,7 @@ jobs:
           python -m black ebcc/ --diff --check --verbose
           python -m isort ebcc/ --diff --check-only --verbose
           python -m flake8 ebcc/ --verbose
+          python -m mypy ebcc/ --verbose
       - name: Run unit tests
         run: |
           python -m pip install pytest pytest-cov

ebcc/__init__.py

@@ -34,35 +34,40 @@
 >>> ccsd.kernel()
 """
 
+from __future__ import annotations
+
+"""Version of the package."""
 __version__ = "1.4.5"
 
+"""List of supported ansatz types."""
+METHOD_TYPES = ["MP", "CC", "LCC", "QCI", "QCC", "DC"]
+
 import os
 import sys
-
-
-# --- Import NumPy here to allow drop-in replacements
+from typing import TYPE_CHECKING
 
 import numpy
 
+from ebcc.core.logging import NullLogger, default_log, init_logging
+from ebcc.cc import GEBCC, REBCC, UEBCC
+from ebcc.core import precision
+from ebcc.core.ansatz import Ansatz
+from ebcc.ham.space import Space
+from ebcc.opt import BruecknerGEBCC, BruecknerREBCC, BruecknerUEBCC
 
-# --- Logging:
+if TYPE_CHECKING:
+    from typing import Any, Callable
 
-from ebcc.logging import default_log, init_logging, NullLogger
+    from pyscf.scf.hf import SCF
 
+    from ebcc.cc.base import BaseEBCC
 
-# --- Types of ansatz supporting by the EBCC solvers:
 
-METHOD_TYPES = ["MP", "CC", "LCC", "QCI", "QCC", "DC"]
-
-
-# --- General constructor:
+sys.modules["ebcc.precision"] = precision  # Compatibility with older codegen versions
 
-from ebcc.gebcc import GEBCC
-from ebcc.rebcc import REBCC
-from ebcc.uebcc import UEBCC
 
-
-def EBCC(mf, *args, **kwargs):
+def EBCC(mf: SCF, *args: Any, **kwargs: Any) -> BaseEBCC:
+    """Construct an EBCC object for the given mean-field object."""
     from pyscf import scf
 
     if isinstance(mf, scf.uhf.UHF):
@@ -76,25 +81,20 @@ def EBCC(mf, *args, **kwargs):
 EBCC.__doc__ = REBCC.__doc__
 
 
-# --- Constructors for boson-free calculations:
-
-
-def _factory(ansatz):
-    def constructor(mf, *args, **kwargs):
-        from pyscf import scf
+def _factory(ansatz: str) -> Callable[[SCF, Any, Any], BaseEBCC]:
+    """Constructor for some specific ansatz."""
+    from pyscf import scf
 
+    def constructor(mf: SCF, *args: Any, **kwargs: Any) -> BaseEBCC:
+        """Construct an EBCC object for the given mean-field object."""
         kwargs["ansatz"] = ansatz
 
         if isinstance(mf, scf.uhf.UHF):
-            cc = UEBCC(mf, *args, **kwargs)
+            return UEBCC(mf, *args, **kwargs)
         elif isinstance(mf, scf.ghf.GHF):
-            cc = GEBCC(mf, *args, **kwargs)
+            return GEBCC(mf, *args, **kwargs)
         else:
-            cc = REBCC(mf, *args, **kwargs)
-
-        cc.__doc__ = REBCC.__doc__
-
-        return cc
+            return REBCC(mf, *args, **kwargs)
 
     return constructor
 
@@ -107,21 +107,10 @@ def constructor(mf, *args, **kwargs):
 del _factory
 
 
-# --- Other imports:
-
-from ebcc.ansatz import Ansatz
-from ebcc.brueckner import BruecknerGEBCC, BruecknerREBCC, BruecknerUEBCC
-from ebcc.space import Space
-
-
-# --- List available methods:
-
-
-def available_models(verbose=True):  # pragma: no cover
-    """List available coupled-cluster models for each of general (G),
-    restricted (R) and unrestricted (U) Hartree--Fock references.
-    """
-
+def available_models(
+    verbose: bool = True,
+) -> tuple[tuple[str, ...], tuple[str, ...], tuple[str, ...]]:  # pragma: no cover
+    """List available coupled-cluster models for each spin type."""
     cd = os.path.dirname(os.path.realpath(__file__))
     path = os.path.join(cd, "codegen")
     _, _, files = list(os.walk(path))[0]