Guard TYPE_CHECKING only imports (#3827)

* flake8 --select=TYP001

* fix type import

* format docstring to google style

* fix more unguarded typecheck imports

* format more docs to Google format

* trigger CI (better var names)

---------

Co-authored-by: Janosh Riebesell <janosh.riebesell@gmail.com>

pymatgen/alchemy/materials.py

@@ -8,7 +8,7 @@
 import datetime
 import json
 import re
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING
 from warnings import warn
 
 from monty.json import MSONable, jsanitize
@@ -22,6 +22,7 @@
 
 if TYPE_CHECKING:
     from collections.abc import Sequence
+    from typing import Any
 
     from typing_extensions import Self
 

pymatgen/alchemy/transmuters.py

@@ -12,13 +12,14 @@
 import os
 import re
 from multiprocessing import Pool
-from typing import TYPE_CHECKING, Callable
+from typing import TYPE_CHECKING
 
 from pymatgen.alchemy.materials import TransformedStructure
 from pymatgen.io.vasp.sets import MPRelaxSet, VaspInputSet
 
 if TYPE_CHECKING:
     from collections.abc import Sequence
+    from typing import Callable
 
     from typing_extensions import Self
 

pymatgen/analysis/chemenv/coordination_environments/coordination_geometries.py

@@ -13,13 +13,15 @@
 import itertools
 import json
 import os
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING
 
 import numpy as np
 from monty.json import MontyDecoder, MSONable
 from scipy.special import factorial
 
 if TYPE_CHECKING:
+    from typing import Any
+
     from typing_extensions import Self
 
 __author__ = "David Waroquiers"

pymatgen/analysis/chemenv/utils/coordination_geometry_utils.py

@@ -3,7 +3,7 @@
 from __future__ import annotations
 
 import math
-from typing import TYPE_CHECKING, Callable
+from typing import TYPE_CHECKING
 
 import numpy as np
 from numpy.linalg import norm
@@ -14,6 +14,8 @@
 from pymatgen.analysis.chemenv.utils.chemenv_errors import SolidAngleError
 
 if TYPE_CHECKING:
+    from typing import Callable
+
     from numpy.typing import ArrayLike
     from typing_extensions import Self
 

pymatgen/analysis/elasticity/elastic.py

@@ -9,7 +9,7 @@
 import itertools
 import math
 import warnings
-from typing import TYPE_CHECKING, Literal
+from typing import TYPE_CHECKING
 
 import numpy as np
 import sympy as sp
@@ -25,6 +25,7 @@
 
 if TYPE_CHECKING:
     from collections.abc import Sequence
+    from typing import Literal
 
     from numpy.typing import ArrayLike
     from typing_extensions import Self

pymatgen/analysis/elasticity/strain.py

@@ -8,7 +8,7 @@
 
 import collections
 import itertools
-from typing import TYPE_CHECKING, Literal
+from typing import TYPE_CHECKING
 
 import numpy as np
 import scipy
@@ -18,6 +18,7 @@
 
 if TYPE_CHECKING:
     from collections.abc import Sequence
+    from typing import Literal
 
     from numpy.typing import ArrayLike
     from typing_extensions import Self

pymatgen/analysis/ewald.py

@@ -3,10 +3,10 @@
 from __future__ import annotations
 
 import bisect
+import math
 from copy import copy, deepcopy
 from datetime import datetime
-from math import log, pi, sqrt
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING
 from warnings import warn
 
 import numpy as np
@@ -18,6 +18,8 @@
 from pymatgen.util.due import Doi, due
 
 if TYPE_CHECKING:
+    from typing import Any
+
     from typing_extensions import Self
 
 __author__ = "Shyue Ping Ong, William Davidson Richard"
@@ -55,7 +57,7 @@ class EwaldSummation(MSONable):
     """
 
     # Converts unit of q*q/r into eV
-    CONV_FACT = 1e10 * constants.e / (4 * pi * constants.epsilon_0)
+    CONV_FACT = 1e10 * constants.e / (4 * math.pi * constants.epsilon_0)
 
     def __init__(
         self,
@@ -102,12 +104,12 @@ def __init__(
 
         self._acc_factor = acc_factor
         # set screening length
-        self._eta = eta or (len(structure) * w / (self._vol**2)) ** (1 / 3) * pi
-        self._sqrt_eta = sqrt(self._eta)
+        self._eta = eta or (len(structure) * w / (self._vol**2)) ** (1 / 3) * math.pi
+        self._sqrt_eta = math.sqrt(self._eta)
 
         # acc factor used to automatically determine the optimal real and
         # reciprocal space cutoff radii
-        self._accf = sqrt(log(10**acc_factor))
+        self._accf = math.sqrt(math.log(10**acc_factor))
 
         self._rmax = real_space_cut or self._accf / self._sqrt_eta
         self._gmax = recip_space_cut or 2 * self._sqrt_eta * self._accf
@@ -307,7 +309,7 @@ def _calc_recip(self):
         This method is heavily vectorized to utilize numpy's C backend for speed.
         """
         n_sites = len(self._struct)
-        prefactor = 2 * pi / self._vol
+        prefactor = 2 * math.pi / self._vol
         e_recip = np.zeros((n_sites, n_sites), dtype=np.float64)
         forces = np.zeros((n_sites, 3), dtype=np.float64)
         coords = self._coords
@@ -332,7 +334,7 @@ def _calc_recip(self):
 
         for g, g2, gr, exp_val, s_real, s_imag in zip(gs, g2s, grs, exp_vals, s_reals, s_imags):
             # Uses the identity sin(x)+cos(x) = 2**0.5 sin(x + pi/4)
-            m = np.sin((gr[None, :] + pi / 4) - gr[:, None])
+            m = np.sin((gr[None, :] + math.pi / 4) - gr[:, None])
             m *= exp_val / g2
 
             e_recip += m
@@ -350,7 +352,7 @@ def _calc_recip(self):
     def _calc_real_and_point(self):
         """Determine the self energy -(eta/pi)**(1/2) * sum_{i=1}^{N} q_i**2."""
         frac_coords = self._struct.frac_coords
-        force_pf = 2 * self._sqrt_eta / sqrt(pi)
+        force_pf = 2 * self._sqrt_eta / math.sqrt(math.pi)
         coords = self._coords
         n_sites = len(self._struct)
         e_real = np.empty((n_sites, n_sites), dtype=np.float64)
@@ -359,7 +361,7 @@ def _calc_real_and_point(self):
 
         qs = np.array(self._oxi_states)
 
-        e_point = -(qs**2) * sqrt(self._eta / pi)
+        e_point = -(qs**2) * math.sqrt(self._eta / math.pi)
 
         for idx in range(n_sites):
             nf_coords, rij, js, _ = self._struct.lattice.get_points_in_sphere(

pymatgen/analysis/interface_reactions.py

@@ -8,7 +8,7 @@
 import json
 import os
 import warnings
-from typing import Literal
+from typing import TYPE_CHECKING
 
 import matplotlib.pyplot as plt
 import numpy as np
@@ -23,6 +23,9 @@
 from pymatgen.util.plotting import pretty_plot
 from pymatgen.util.string import htmlify, latexify
 
+if TYPE_CHECKING:
+    from typing import Literal
+
 __author__ = "Yihan Xiao, Matthew McDermott"
 __maintainer__ = "Matthew McDermott"
 __email__ = "mcdermott@lbl.gov"

pymatgen/analysis/magnetism/jahnteller.py

@@ -4,7 +4,7 @@
 
 import os
 import warnings
-from typing import TYPE_CHECKING, Any, Literal, cast
+from typing import TYPE_CHECKING, Literal, cast
 
 import numpy as np
 
@@ -14,6 +14,8 @@
 from pymatgen.symmetry.analyzer import SpacegroupAnalyzer
 
 if TYPE_CHECKING:
+    from typing import Any
+
     from pymatgen.core import Structure
 
 MODULE_DIR = os.path.dirname(os.path.abspath(__file__))

pymatgen/analysis/phase_diagram.py

@@ -11,7 +11,7 @@
 import warnings
 from collections import defaultdict
 from functools import lru_cache
-from typing import TYPE_CHECKING, Any, Literal, no_type_check
+from typing import TYPE_CHECKING, no_type_check
 
 import matplotlib.pyplot as plt
 import numpy as np
@@ -38,6 +38,7 @@
 if TYPE_CHECKING:
     from collections.abc import Collection, Iterator, Sequence
     from io import StringIO
+    from typing import Any, Literal
 
     from numpy.typing import ArrayLike
     from typing_extensions import Self

pymatgen/analysis/pourbaix_diagram.py

@@ -12,7 +12,7 @@
 from copy import deepcopy
 from functools import cmp_to_key, partial
 from multiprocessing import Pool
-from typing import TYPE_CHECKING, Any, no_type_check
+from typing import TYPE_CHECKING, no_type_check
 
 import numpy as np
 from monty.json import MontyDecoder, MSONable
@@ -31,6 +31,8 @@
 from pymatgen.util.string import Stringify
 
 if TYPE_CHECKING:
+    from typing import Any
+
     import matplotlib.pyplot as plt
     from typing_extensions import Self
 

pymatgen/analysis/structure_matcher.py

@@ -4,7 +4,7 @@
 
 import abc
 import itertools
-from typing import TYPE_CHECKING, Literal
+from typing import TYPE_CHECKING
 
 import numpy as np
 from monty.json import MSONable
@@ -16,6 +16,7 @@
 
 if TYPE_CHECKING:
     from collections.abc import Mapping, Sequence
+    from typing import Literal
 
     from typing_extensions import Self
 

pymatgen/cli/pmg_config.py

@@ -8,7 +8,7 @@
 import shutil
 import subprocess
 from glob import glob
-from typing import TYPE_CHECKING, Literal
+from typing import TYPE_CHECKING
 from urllib.request import urlretrieve
 
 from monty.json import jsanitize
@@ -21,6 +21,7 @@
 
 if TYPE_CHECKING:
     from argparse import Namespace
+    from typing import Literal
 
 
 def setup_cp2k_data(cp2k_data_dirs: list[str]) -> None:

pymatgen/command_line/bader_caller.py

@@ -22,7 +22,7 @@
 from pathlib import Path
 from shutil import which
 from tempfile import TemporaryDirectory
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING
 
 import numpy as np
 from monty.dev import deprecated
@@ -34,6 +34,8 @@
 from pymatgen.io.vasp.outputs import Chgcar
 
 if TYPE_CHECKING:
+    from typing import Any
+
     from typing_extensions import Self
 
     from pymatgen.core import Structure

pymatgen/core/interface.py

@@ -8,7 +8,7 @@
 from functools import reduce
 from itertools import chain, combinations, product
 from math import cos, floor, gcd
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING
 
 import numpy as np
 from monty.fractions import lcm
@@ -25,6 +25,7 @@
 
 if TYPE_CHECKING:
     from collections.abc import Sequence
+    from typing import Any
 
     from numpy.typing import ArrayLike
     from typing_extensions import Self

pymatgen/core/periodic_table.py

@@ -11,7 +11,7 @@
 from enum import Enum, unique
 from itertools import combinations, product
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Callable, Literal
+from typing import TYPE_CHECKING
 
 import numpy as np
 from monty.json import MSONable
@@ -20,6 +20,8 @@
 from pymatgen.util.string import Stringify, formula_double_format
 
 if TYPE_CHECKING:
+    from typing import Any, Callable, Literal
+
     from typing_extensions import Self
 
     from pymatgen.util.typing import SpeciesLike