Clarify docstrings and type hints

src/quacc/__init__.py

@@ -2,25 +2,28 @@
 from __future__ import annotations
 
 from importlib.metadata import version
+from typing import TYPE_CHECKING
 
 from ase import Atoms
 from ase.io.jsonio import decode, encode
 
 from quacc.settings import QuaccSettings
 from quacc.wflow.decorators import flow, job, subflow
 
+if TYPE_CHECKING:
+    from typing import Any
 __all__ = ["flow", "job", "subflow"]
 
 
-def atoms_as_dict(s: Atoms) -> dict:
+def atoms_as_dict(s: Atoms) -> dict[str, Any]:
     # Uses Monty's MSONable spec
     # Normally, we would want to this to be a wrapper around atoms.todict() with @module and
     # @class key-value pairs inserted. However, atoms.todict()/atoms.fromdict() does not currently
     # work properly with constraints.
     return {"@module": "ase.atoms", "@class": "Atoms", "atoms_json": encode(s)}
 
 
-def atoms_from_dict(d: dict) -> Atoms:
+def atoms_from_dict(d: dict[str, Any]) -> Atoms:
     # Uses Monty's MSONable spec
     # Normally, we would want to have this be a wrapper around atoms.fromdict()
     # that just ignores the @module/@class key-value pairs. However, atoms.todict()/atoms.fromdict()

src/quacc/atoms/slabs.py

@@ -15,6 +15,8 @@
 from quacc.atoms.core import copy_atoms
 
 if TYPE_CHECKING:
+    from typing import Any
+
     from pymatgen.core import Structure
 
 logger = logging.getLogger(__name__)
@@ -204,8 +206,8 @@ def make_adsorbate_structures(
     modes: list[str] | None = None,
     allowed_surface_symbols: list[str] | None = None,
     allowed_surface_indices: list[int] | None = None,
-    ads_site_finder_kwargs: dict | None = None,
-    find_ads_sites_kwargs: dict | None = None,
+    ads_site_finder_kwargs: dict[str, Any] | None = None,
+    find_ads_sites_kwargs: dict[str, Any] | None = None,
 ) -> list[Atoms]:
     """
     Add a single adsorbate to a structure for every requested adsorption mode

src/quacc/calculators/vasp/custodian.py

@@ -5,6 +5,7 @@
 
 import os
 import shlex
+from typing import TYPE_CHECKING
 
 from custodian import Custodian
 from custodian.vasp.handlers import (
@@ -24,6 +25,9 @@
 from custodian.vasp.jobs import VaspJob
 from custodian.vasp.validators import VaspFilesValidator, VasprunXMLValidator
 
+if TYPE_CHECKING:
+    from typing import Any
+
 
 def run_custodian(
     vasp_parallel_cmd: str | None = None,
@@ -35,8 +39,8 @@ def run_custodian(
     vasp_custodian_handlers: list[str] | None = None,
     vasp_custodian_validators: list[str] | None = None,
     scratch_dir: str | None = None,
-    vasp_job_kwargs: dict | None = None,
-    custodian_kwargs: dict | None = None,
+    vasp_job_kwargs: dict[str, Any] | None = None,
+    custodian_kwargs: dict[str, Any] | None = None,
 ) -> None:
     """
     Function to run VASP Custodian

src/quacc/calculators/vasp/io.py

@@ -9,9 +9,10 @@
 
 if TYPE_CHECKING:
     from pathlib import Path
+    from typing import Any
 
 
-def load_vasp_yaml_calc(yaml_path: str | Path) -> dict:
+def load_vasp_yaml_calc(yaml_path: str | Path) -> dict[str, Any]:
     """
     Loads a YAML file containing calculator settings. Used for VASP calculations
     and can read quacc-formatted YAMLs that are of the following format:

src/quacc/calculators/vasp/params.py

@@ -15,16 +15,21 @@
 from quacc.atoms.core import check_is_metal
 
 if TYPE_CHECKING:
+    from typing import Any, Literal
+
     from ase import Atoms
 
 logger = logging.getLogger(__name__)
 
 
 def calc_swaps(
-    user_calc_params: dict, auto_kpts: dict, input_atoms: Atoms, force_copilot: bool
-) -> dict:
+    user_calc_params: dict[str, Any],
+    auto_kpts: dict[Literal["line_density", "kppvol", "kppa"], float],
+    input_atoms: Atoms,
+    force_copilot: bool,
+) -> dict[str, Any]:
     """
-    Swaps out bad INCAR flags.
+    Swaps Fout bad INCAR flags.
 
     Parameters
     ----------
@@ -230,7 +235,7 @@ def calc_swaps(
     return calc.parameters if force_copilot else calc.parameters | user_calc_params
 
 
-def remove_unused_flags(user_calc_params: dict) -> dict:
+def remove_unused_flags(user_calc_params: dict[str, Any]) -> dict[str, Any]:
     """
     Removes unused flags in the INCAR, like EDIFFG if you are doing NSW = 0.
 
@@ -273,7 +278,9 @@ def remove_unused_flags(user_calc_params: dict) -> dict:
     return user_calc_params
 
 
-def set_auto_dipole(user_calc_params: dict, input_atoms: Atoms) -> dict:
+def set_auto_dipole(
+    user_calc_params: dict[str, Any], input_atoms: Atoms
+) -> dict[str, Any]:
     """
     Sets flags related to the auto_dipole kwarg.
 
@@ -302,10 +309,10 @@ def set_auto_dipole(user_calc_params: dict, input_atoms: Atoms) -> dict:
 
 
 def convert_auto_kpts(
-    user_calc_params: dict,
-    auto_kpts: dict,
+    user_calc_params: dict[str, Any],
+    auto_kpts: dict[Literal["line_density", "kppvol", "kppa"], float],
     input_atoms: Atoms,
-) -> dict:
+) -> dict[str, Any]:
     """
     Shortcuts for pymatgen k-point generation schemes.
 

src/quacc/calculators/vasp/vasp.py

@@ -47,7 +47,7 @@ def __init__(
         copy_magmoms: bool | None = None,
         preset_mag_default: float | None = None,
         mag_cutoff: None | float = None,
-        elemental_magmoms: dict | None = None,
+        elemental_magmoms: dict[str, float] | None = None,
         auto_kpts: dict[Literal["line_density", "kppvol", "kppa"], float]
         | dict[Literal["length_densities"], list[float]]
         | None = None,

src/quacc/recipes/dftb/core.py

@@ -12,7 +12,7 @@
 from quacc.utils.files import check_logfile
 
 if TYPE_CHECKING:
-    from typing import Literal
+    from typing import Any, Literal
 
     from ase import Atoms
 
@@ -27,7 +27,7 @@ def static_job(
     atoms: Atoms,
     method: Literal["GFN1-xTB", "GFN2-xTB", "DFTB"] = "GFN2-xTB",
     kpts: tuple | list[tuple] | dict | None = None,
-    calc_swaps: dict | None = None,
+    calc_swaps: dict[str, Any] | None = None,
     copy_files: list[str] | None = None,
 ) -> RunSchema:
     """
@@ -88,7 +88,7 @@ def relax_job(
     method: Literal["GFN1-xTB", "GFN2-xTB", "DFTB"] = "GFN2-xTB",
     kpts: tuple | list[tuple] | dict | None = None,
     relax_cell: bool = False,
-    calc_swaps: dict | None = None,
+    calc_swaps: dict[str, Any] | None = None,
     copy_files: list[str] | None = None,
 ) -> RunSchema:
     """
@@ -155,9 +155,9 @@ def relax_job(
 
 def _base_job(
     atoms: Atoms,
-    defaults: dict | None = None,
-    calc_swaps: dict | None = None,
-    additional_fields: dict | None = None,
+    defaults: dict[str, Any] | None = None,
+    calc_swaps: dict[str, Any] | None = None,
+    additional_fields: dict[str, Any] | None = None,
     copy_files: list[str] | None = None,
 ) -> RunSchema:
     """

src/quacc/recipes/emt/core.py

@@ -16,6 +16,8 @@
 from quacc.utils.dicts import merge_dicts
 
 if TYPE_CHECKING:
+    from typing import Any
+
     from ase import Atoms
 
     from quacc.schemas.ase import OptSchema, RunSchema
@@ -24,18 +26,12 @@
 @job
 def static_job(
     atoms: Atoms,
-    calc_swaps: dict | None = None,
+    calc_swaps: dict[str, Any] | None = None,
     copy_files: list[str] | None = None,
 ) -> RunSchema:
     """
     Carry out a static calculation.
 
-    !!! Info "Calculator defaults, which can be overriden by `calc_swaps`"
-
-        ```python
-        {}
-        ```
-
     Parameters
     ----------
     atoms
@@ -44,6 +40,12 @@ def static_job(
         Dictionary of custom kwargs for the EMT calculator. Set a value to
         `None` to remove a pre-existing key entirely. For a list of available
         keys, refer to the `ase.calculators.emt.EMT` calculator.
+
+        !!! Info "Calculator defaults"
+
+            ```python
+            {}
+            ```
     copy_files
         Files to copy to the runtime directory.
 
@@ -68,25 +70,13 @@ def static_job(
 def relax_job(
     atoms: Atoms,
     relax_cell: bool = False,
-    calc_swaps: dict | None = None,
-    opt_swaps: dict | None = None,
+    calc_swaps: dict[str, Any] | None = None,
+    opt_swaps: dict[str, Any] | None = None,
     copy_files: list[str] | None = None,
 ) -> OptSchema:
     """
     Carry out a geometry optimization.
 
-    !!! Info "Calculator defaults, which can be overriden by `calc_swaps`"
-
-        ```python
-        {}
-        ```
-
-    !!! Info "Optimizer defaults, which can be overriden by `opt_swaps`"
-
-        ```python
-        {"fmax": 0.01, "max_steps": 1000, "optimizer": FIRE}
-        ```
-
     Parameters
     ----------
     atoms
@@ -97,10 +87,22 @@ def relax_job(
         Dictionary of custom kwargs for the EMT calculator. Set a value to
         `None` to remove a pre-existing key entirely. For a list of available
         keys, refer to the `ase.calculators.emt.EMT` calculator.
+
+        !!! Info "Calculator defaults"
+
+            ```python
+            {}
+            ```
     opt_swaps
         Dictionary of custom kwargs for the optimization process. Set a value
         to `None` to remove a pre-existing key entirely. For a list of available
         keys, refer to [quacc.runners.calc.run_ase_opt][].
+
+        !!! Info "Optimizer defaults"
+
+            ```python
+            {"fmax": 0.01, "max_steps": 1000, "optimizer": FIRE}
+            ```
     copy_files
         Files to copy to the runtime directory.
 

src/quacc/recipes/emt/defects.py

@@ -10,6 +10,8 @@
 from quacc.recipes.emt.core import relax_job, static_job
 
 if TYPE_CHECKING:
+    from typing import Any
+
     from ase import Atoms
     from pymatgen.analysis.defects.generators import (
         AntiSiteGenerator,
@@ -34,10 +36,10 @@ def bulk_to_defects_flow(
         | VoronoiInterstitialGenerator
     ) = VacancyGenerator,
     defect_charge: int = 0,
-    make_defects_kwargs: dict | None = None,
+    make_defects_kwargs: dict[str, Any] | None = None,
     run_static: bool = True,
-    defect_relax_kwargs: dict | None = None,
-    defect_static_kwargs: dict | None = None,
+    defect_relax_kwargs: dict[str, Any] | None = None,
+    defect_static_kwargs: dict[str, Any] | None = None,
 ) -> list[RunSchema | OptSchema]:
     """
     Workflow consisting of:

src/quacc/recipes/emt/slabs.py

@@ -8,6 +8,8 @@
 from quacc.recipes.emt.core import relax_job, static_job
 
 if TYPE_CHECKING:
+    from typing import Any
+
     from ase import Atoms
 
     from quacc.schemas.ase import OptSchema, RunSchema
@@ -16,10 +18,10 @@
 @flow
 def bulk_to_slabs_flow(
     atoms: Atoms,
-    make_slabs_kwargs: dict | None = None,
+    make_slabs_kwargs: dict[str, Any] | None = None,
     run_static: bool = True,
-    slab_relax_kwargs: dict | None = None,
-    slab_static_kwargs: dict | None = None,
+    slab_relax_kwargs: dict[str, Any] | None = None,
+    slab_static_kwargs: dict[str, Any] | None = None,
 ) -> list[RunSchema | OptSchema]:
     """
     Workflow consisting of: