Tidy annotations

src/ifermi/analysis.py

@@ -1,7 +1,8 @@
 """Isosurface and isoline analysis functions."""
 
+from __future__ import annotations
+
 import warnings
-from typing import Union
 
 import numpy as np
 
@@ -41,7 +42,7 @@ def isosurface_area(vertices: np.ndarray, faces: np.ndarray) -> float:
 
 def average_properties(
     vertices: np.ndarray, faces: np.ndarray, properties: np.ndarray, norm: bool = False
-) -> Union[float, np.ndarray]:
+) -> float | np.ndarray:
     """Average property across an isosurface.
 
     Args:

src/ifermi/brillouin_zone.py

@@ -1,12 +1,16 @@
 """Brillouin zone and slice geometries."""
 
+from __future__ import annotations
+
 import itertools
 from dataclasses import dataclass, field
-from typing import Optional
+from typing import TYPE_CHECKING
 
 import numpy as np
 from monty.json import MSONable
-from pymatgen.core.structure import Structure
+
+if TYPE_CHECKING:
+    from pymatgen.core.structure import Structure
 
 __all__ = ["ReciprocalSlice", "ReciprocalCell", "WignerSeitzCell"]
 
@@ -23,10 +27,10 @@ class ReciprocalSlice(MSONable):
             the 3D Brillouin zone to points on the 2D slice.
     """
 
-    reciprocal_space: "ReciprocalCell"
+    reciprocal_space: ReciprocalCell
     vertices: np.ndarray
     transformation: np.ndarray
-    _edges: Optional[np.ndarray] = field(default=None, init=False)
+    _edges: np.ndarray | None = field(default=None, init=False)
 
     def __post_init__(self):
         """Ensure all inputs are numpy arrays."""
@@ -66,7 +70,7 @@ class ReciprocalCell(MSONable):
     faces: list[list[int]]
     centers: np.ndarray
     normals: np.ndarray
-    _edges: Optional[np.ndarray] = field(default=None, init=False)
+    _edges: np.ndarray | None = field(default=None, init=False)
 
     def __post_init__(self):
         """Ensure all inputs are numpy arrays."""
@@ -76,7 +80,7 @@ def __post_init__(self):
         self.normals = np.array(self.normals)
 
     @classmethod
-    def from_structure(cls, structure: Structure) -> "ReciprocalCell":
+    def from_structure(cls, structure: Structure) -> ReciprocalCell:
         """Initialise the reciprocal cell from a structure.
 
         Args:
@@ -195,7 +199,7 @@ class WignerSeitzCell(ReciprocalCell):
     """
 
     @classmethod
-    def from_structure(cls, structure: Structure) -> "WignerSeitzCell":
+    def from_structure(cls, structure: Structure) -> WignerSeitzCell:
         """Initialise the Wigner-Seitz cell from a structure.
 
         Args:

src/ifermi/interpolate.py

@@ -1,10 +1,14 @@
 """Tools for Fourier and linear interpolation."""
 
-from typing import Optional
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
 
 import numpy as np
 from pymatgen.electronic_structure.bandstructure import BandStructure
-from pymatgen.electronic_structure.core import Spin
+
+if TYPE_CHECKING:
+    from pymatgen.electronic_structure.core import Spin
 
 __all__ = ["FourierInterpolator", "LinearInterpolator", "trim_bandstructure"]
 
@@ -23,8 +27,8 @@ class FourierInterpolator:
     def __init__(
         self,
         band_structure: BandStructure,
-        magmom: Optional[np.ndarray] = None,
-        mommat: Optional[np.ndarray] = None,
+        magmom: np.ndarray | None = None,
+        mommat: np.ndarray | None = None,
     ):
         from BoltzTraP2.units import Angstrom
         from pymatgen.io.ase import AseAtomsAdaptor
@@ -160,7 +164,7 @@ def __init__(
         kpoints: np.ndarray,
         energies: np.ndarray,
         lattice_matrix: np.ndarray,
-        mommat: Optional[np.ndarray] = None,
+        mommat: np.ndarray | None = None,
     ):
         self.kpoints = kpoints
         self.ebands = energies

src/ifermi/kpoints.py

@@ -1,12 +1,17 @@
 """k-point manipulation functions."""
 
+from __future__ import annotations
+
 import warnings
+from typing import TYPE_CHECKING
 
 import numpy as np
-from pymatgen.electronic_structure.bandstructure import BandStructure
 
 from ifermi.defaults import KTOL
 
+if TYPE_CHECKING:
+    from pymatgen.electronic_structure.bandstructure import BandStructure
+
 __all__ = [
     "kpoints_to_first_bz",
     "kpoints_from_bandstructure",