move all type-hint only imports behind if TYPE_CHECKING

src/atomate2/amset/files.py

@@ -3,7 +3,7 @@
 from __future__ import annotations
 
 import logging
-from pathlib import Path
+from typing import TYPE_CHECKING
 
 from monty.serialization import loadfn
 
@@ -12,6 +12,9 @@
 from atomate2.utils.file_client import FileClient, auto_fileclient
 from atomate2.utils.path import strip_hostname
 
+if TYPE_CHECKING:
+    from pathlib import Path
+
 __all__ = ["copy_amset_files"]
 
 

src/atomate2/common/flows/defect.py

@@ -5,12 +5,9 @@
 import logging
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
-from pathlib import Path
+from typing import TYPE_CHECKING
 
-import numpy.typing as npt
 from jobflow import Flow, Job, Maker, OutputReference
-from pymatgen.analysis.defects.core import Defect
-from pymatgen.core.structure import Structure
 
 from atomate2.common.jobs.defect import (
     bulk_supercell_calculation,
@@ -21,6 +18,13 @@
     spawn_energy_curve_calcs,
 )
 
+if TYPE_CHECKING:
+    from pathlib import Path
+
+    import numpy.typing as npt
+    from pymatgen.analysis.defects.core import Defect
+    from pymatgen.core.structure import Structure
+
 logger = logging.getLogger(__name__)
 
 DEFAULT_DISTORTIONS = (-1, -0.15, -0.1, -0.05, 0, 0.05, 0.1, 0.15, 1)

src/atomate2/common/jobs/defect.py

@@ -3,23 +3,27 @@
 from __future__ import annotations
 
 import logging
-from pathlib import Path
-from typing import Callable, Iterable
+from typing import TYPE_CHECKING, Callable, Iterable
 
 import numpy as np
-from emmet.core.tasks import TaskDoc
 from jobflow import Flow, Response, job
-from numpy.typing import NDArray
 from pydantic import BaseModel
-from pymatgen.analysis.defects.core import Defect
 from pymatgen.analysis.defects.supercells import (
     get_matched_structure_mapping,
     get_sc_fromstruct,
 )
 from pymatgen.core import Lattice, Structure
 
 from atomate2.common.schemas.defects import CCDDocument
-from atomate2.vasp.jobs.core import RelaxMaker, StaticMaker
+
+if TYPE_CHECKING:
+    from pathlib import Path
+
+    from emmet.core.tasks import TaskDoc
+    from numpy.typing import NDArray
+    from pymatgen.analysis.defects.core import Defect
+
+    from atomate2.vasp.jobs.core import RelaxMaker, StaticMaker
 
 logger = logging.getLogger(__name__)
 

src/atomate2/common/jobs/utils.py

@@ -2,12 +2,16 @@
 
 from __future__ import annotations
 
+from typing import TYPE_CHECKING
+
 from jobflow import Response, job
-from pymatgen.core import Structure
 from pymatgen.symmetry.analyzer import SpacegroupAnalyzer
 
 from atomate2 import SETTINGS
 
+if TYPE_CHECKING:
+    from pymatgen.core import Structure
+
 __all__ = [
     "structure_to_primitive",
     "structure_to_conventional",

src/atomate2/common/utils.py

@@ -4,11 +4,13 @@
 
 import re
 from importlib import import_module
-from pathlib import Path
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 from monty.serialization import loadfn
 
+if TYPE_CHECKING:
+    from pathlib import Path
+
 
 def get_transformations(
     transformations: tuple[str, ...], params: tuple[dict, ...] | None

src/atomate2/cp2k/files.py

@@ -5,17 +5,20 @@
 import logging
 import re
 from pathlib import Path
-from typing import Sequence
+from typing import TYPE_CHECKING, Sequence
 
-from pymatgen.core import Structure
 from pymatgen.io.cp2k.outputs import Cp2kOutput
 
 from atomate2 import SETTINGS
 from atomate2.common.files import copy_files, get_zfile, gunzip_files, rename_files
-from atomate2.cp2k.sets.base import Cp2kInputGenerator
 from atomate2.utils.file_client import FileClient, auto_fileclient
 from atomate2.utils.path import strip_hostname
 
+if TYPE_CHECKING:
+    from pymatgen.core import Structure
+
+    from atomate2.cp2k.sets.base import Cp2kInputGenerator
+
 __all__ = ["copy_cp2k_outputs", "write_cp2k_input_set"]
 
 

src/atomate2/cp2k/flows/core.py

@@ -4,12 +4,10 @@
 
 from copy import deepcopy
 from dataclasses import dataclass, field
-from pathlib import Path
+from typing import TYPE_CHECKING
 
 from jobflow import Flow, Job, Maker
-from pymatgen.core.structure import Structure
 
-from atomate2.cp2k.jobs.base import BaseCp2kMaker
 from atomate2.cp2k.jobs.core import (
     HybridCellOptMaker,
     HybridRelaxMaker,
@@ -21,6 +19,13 @@
 from atomate2.cp2k.schemas.calculation import Cp2kObject
 from atomate2.cp2k.sets.base import recursive_update
 
+if TYPE_CHECKING:
+    from pathlib import Path
+
+    from pymatgen.core.structure import Structure
+
+    from atomate2.cp2k.jobs.base import BaseCp2kMaker
+
 __all__ = [
     "DoubleRelaxMaker",
     "BandStructureMaker",

src/atomate2/cp2k/jobs/base.py

@@ -4,14 +4,13 @@
 
 from dataclasses import dataclass, field
 from pathlib import Path
-from typing import Callable
+from typing import TYPE_CHECKING, Callable
 
 from jobflow import Maker, Response, job
 from monty.serialization import dumpfn
 from monty.shutil import gzip_dir
 from pymatgen.alchemy.materials import TransformedStructure
 from pymatgen.alchemy.transmuters import StandardTransmuter
-from pymatgen.core import Structure
 from pymatgen.core.trajectory import Trajectory
 from pymatgen.electronic_structure.bandstructure import (
     BandStructure,
@@ -30,6 +29,9 @@
 from atomate2.cp2k.schemas.task import TaskDocument
 from atomate2.cp2k.sets.base import Cp2kInputGenerator
 
+if TYPE_CHECKING:
+    from pymatgen.core import Structure
+
 __all__ = ["BaseCp2kMaker", "cp2k_job"]
 
 

src/atomate2/cp2k/jobs/core.py

@@ -4,7 +4,7 @@
 
 import logging
 from dataclasses import dataclass, field
-from pathlib import Path
+from typing import TYPE_CHECKING
 
 from custodian.cp2k.handlers import (
     AbortHandler,
@@ -15,11 +15,9 @@
 )
 from pymatgen.alchemy.materials import TransformedStructure
 from pymatgen.alchemy.transmuters import StandardTransmuter
-from pymatgen.core.structure import Structure
 
 from atomate2.common.utils import get_transformations
 from atomate2.cp2k.jobs.base import BaseCp2kMaker, cp2k_job
-from atomate2.cp2k.sets.base import Cp2kInputGenerator
 from atomate2.cp2k.sets.core import (
     CellOptSetGenerator,
     HybridCellOptSetGenerator,
@@ -31,6 +29,14 @@
     StaticSetGenerator,
 )
 
+if TYPE_CHECKING:
+    from pathlib import Path
+
+    from pymatgen.core.structure import Structure
+
+    from atomate2.cp2k.sets.base import Cp2kInputGenerator
+
+
 logger = logging.getLogger(__name__)
 
 __all__ = [

src/atomate2/cp2k/run.py

@@ -6,7 +6,7 @@
 import shlex
 import subprocess
 from os.path import expandvars
-from typing import Any, Sequence
+from typing import TYPE_CHECKING, Any, Sequence
 
 from custodian import Custodian
 from custodian.cp2k.handlers import (
@@ -21,11 +21,14 @@
 )
 from custodian.cp2k.jobs import Cp2kJob
 from custodian.cp2k.validators import Cp2kOutputValidator
-from custodian.custodian import ErrorHandler, Validator
 from jobflow.utils import ValueEnum
 
 from atomate2 import SETTINGS
-from atomate2.cp2k.schemas.task import TaskDocument
+
+if TYPE_CHECKING:
+    from custodian.custodian import ErrorHandler, Validator
+
+    from atomate2.cp2k.schemas.task import TaskDocument
 
 __all__ = [
     "JobType",

src/atomate2/cp2k/sets/core.py

@@ -4,14 +4,17 @@
 
 import logging
 from dataclasses import dataclass
+from typing import TYPE_CHECKING
 
-from pymatgen.core import Structure
-from pymatgen.io.cp2k.inputs import Cp2kInput
-from pymatgen.io.cp2k.outputs import Cp2kOutput
 from pymatgen.io.cp2k.utils import get_truncated_coulomb_cutoff
 
 from atomate2.cp2k.sets.base import Cp2kInputGenerator
 
+if TYPE_CHECKING:
+    from pymatgen.core import Structure
+    from pymatgen.io.cp2k.inputs import Cp2kInput
+    from pymatgen.io.cp2k.outputs import Cp2kOutput
+
 logger = logging.getLogger(__name__)
 
 

src/atomate2/forcefields/flows.py

@@ -3,14 +3,18 @@
 from __future__ import annotations
 
 from dataclasses import dataclass, field
+from typing import TYPE_CHECKING
 
 from jobflow import Flow, Maker
-from pymatgen.core.structure import Structure
 
 from atomate2.forcefields.jobs import CHGNetRelaxMaker
-from atomate2.vasp.jobs.base import BaseVaspMaker
 from atomate2.vasp.jobs.core import RelaxMaker
 
+if TYPE_CHECKING:
+    from pymatgen.core.structure import Structure
+
+    from atomate2.vasp.jobs.base import BaseVaspMaker
+
 
 @dataclass
 class CHGNetVaspRelaxMaker(Maker):

src/atomate2/forcefields/jobs.py

@@ -4,12 +4,15 @@
 
 import logging
 from dataclasses import dataclass, field
+from typing import TYPE_CHECKING
 
 from jobflow import Maker, job
-from pymatgen.core.structure import Structure
 
 from atomate2.forcefields.schemas import ForceFieldTaskDocument
 
+if TYPE_CHECKING:
+    from pymatgen.core.structure import Structure
+
 logger = logging.getLogger(__name__)
 
 

src/atomate2/lobster/files.py

@@ -3,12 +3,15 @@
 from __future__ import annotations
 
 import logging
-from pathlib import Path
+from typing import TYPE_CHECKING
 
 from atomate2.common.files import copy_files, get_zfile, gunzip_files
 from atomate2.utils.file_client import FileClient, auto_fileclient
 from atomate2.utils.path import strip_hostname
 
+if TYPE_CHECKING:
+    from pathlib import Path
+
 LOBSTEROUTPUT_FILES = [
     "lobsterout",
     "CHARGE.lobster",

src/atomate2/lobster/run.py

@@ -6,16 +6,18 @@
 import shlex
 import subprocess
 from os.path import expandvars
-from typing import Any, Sequence
+from typing import TYPE_CHECKING, Any, Sequence
 
 from custodian import Custodian
-from custodian.custodian import Validator
 from custodian.lobster.handlers import EnoughBandsValidator, LobsterFilesValidator
 from custodian.lobster.jobs import LobsterJob
 from jobflow.utils import ValueEnum
 
 from atomate2 import SETTINGS
 
+if TYPE_CHECKING:
+    from custodian.custodian import Validator
+
 __all__ = ["JobType", "run_lobster"]
 
 _DEFAULT_VALIDATORS = (LobsterFilesValidator(), EnoughBandsValidator())

src/atomate2/vasp/builders/elastic.py

@@ -3,16 +3,19 @@
 from __future__ import annotations
 
 from itertools import chain
+from typing import TYPE_CHECKING
 
 import numpy as np
 from maggma.builders import Builder
-from maggma.core import Store
 from pydash import get
 from pymatgen.analysis.elasticity import Deformation, Stress
 
 from atomate2 import SETTINGS
 from atomate2.common.schemas.elastic import ElasticDocument
 
+if TYPE_CHECKING:
+    from maggma.core import Store
+
 
 class ElasticBuilder(Builder):
     """

src/atomate2/vasp/files.py

@@ -5,15 +5,17 @@
 import logging
 import re
 from pathlib import Path
-from typing import Sequence
-
-from pymatgen.core import Structure
+from typing import TYPE_CHECKING, Sequence
 
 from atomate2 import SETTINGS
 from atomate2.common.files import copy_files, get_zfile, gunzip_files, rename_files
 from atomate2.utils.file_client import FileClient, auto_fileclient
 from atomate2.utils.path import strip_hostname
-from atomate2.vasp.sets.base import VaspInputGenerator
+
+if TYPE_CHECKING:
+    from pymatgen.core import Structure
+
+    from atomate2.vasp.sets.base import VaspInputGenerator
 
 __all__ = ["copy_vasp_outputs", "get_largest_relax_extension"]