add types

src/pymatgen/command_line/enumlib_caller.py

@@ -8,19 +8,19 @@
 
 If you use this module, please cite:
 
-Gus L. W. Hart and Rodney W. Forcade, "Algorithm for generating derivative
-structures," Phys. Rev. B 77 224115 (26 June 2008)
+    - Gus L. W. Hart and Rodney W. Forcade, "Algorithm for generating derivative
+    structures," Phys. Rev. B 77 224115 (26 June 2008)
 
-Gus L. W. Hart and Rodney W. Forcade, "Generating derivative structures from
-multilattices: Application to hcp alloys," Phys. Rev. B 80 014120 (July 2009)
+    - Gus L. W. Hart and Rodney W. Forcade, "Generating derivative structures from
+    multilattices: Application to hcp alloys," Phys. Rev. B 80 014120 (July 2009)
 
-Gus L. W. Hart, Lance J. Nelson, and Rodney W. Forcade, "Generating
-derivative structures at a fixed concentration," Comp. Mat. Sci. 59
-101-107 (March 2012)
+    - Gus L. W. Hart, Lance J. Nelson, and Rodney W. Forcade, "Generating
+    derivative structures at a fixed concentration," Comp. Mat. Sci. 59
+    101-107 (March 2012)
 
-Wiley S. Morgan, Gus L. W. Hart, Rodney W. Forcade, "Generating derivative
-superstructures for systems with high configurational freedom," Comp. Mat.
-Sci. 136 144-149 (May 2017)
+    - Wiley S. Morgan, Gus L. W. Hart, Rodney W. Forcade, "Generating derivative
+    superstructures for systems with high configurational freedom," Comp. Mat.
+    Sci. 136 144-149 (May 2017)
 """
 
 from __future__ import annotations
@@ -34,6 +34,7 @@
 from glob import glob
 from shutil import which
 from threading import Timer
+from typing import TYPE_CHECKING
 
 import numpy as np
 from monty.dev import requires
@@ -44,17 +45,19 @@
 from pymatgen.io.vasp.inputs import Poscar
 from pymatgen.symmetry.analyzer import SpacegroupAnalyzer
 
+if TYPE_CHECKING:
+    from typing import ClassVar
+
 logger = logging.getLogger(__name__)
 
-# Favor the use of the newer "enum.x" by Gus Hart instead of the older
-# "multienum.x"
-enum_cmd = which("enum.x") or which("multienum.x")
-# prefer makestr.x at present
-makestr_cmd = which("makestr.x") or which("makeStr.x") or which("makeStr.py")
+# Favor the use of the newer "enum.x" by Gus Hart over "multienum.x"
+ENUM_CMD = which("enum.x") or which("multienum.x")
+# Prefer makestr.x at present
+MAKESTR_CMD = which("makestr.x") or which("makeStr.x") or which("makeStr.py")
 
 
 @requires(
-    enum_cmd and makestr_cmd,
+    ENUM_CMD and MAKESTR_CMD,
     "EnumlibAdaptor requires the executables 'enum.x' or 'multienum.x' "
     "and 'makestr.x' or 'makeStr.py' to be in the path. Please download the "
     "library at https://github.com/msg-byu/enumlib and follow the instructions "
@@ -64,26 +67,26 @@ class EnumlibAdaptor:
     """An adaptor for enumlib.
 
     Attributes:
-        structures (list): all enumerated structures.
+        structures (list[Structure]): all enumerated structures.
     """
 
-    amount_tol = 1e-5
+    amount_tol: ClassVar[float] = 1e-5
 
     def __init__(
         self,
-        structure,
-        min_cell_size=1,
-        max_cell_size=1,
-        symm_prec=0.1,
-        enum_precision_parameter=0.001,
-        refine_structure=False,
-        check_ordered_symmetry=True,
-        timeout=None,
-    ):
+        structure: Structure,
+        min_cell_size: int = 1,
+        max_cell_size: int = 1,
+        symm_prec: float = 0.1,
+        enum_precision_parameter: float = 0.001,
+        refine_structure: bool = False,
+        check_ordered_symmetry: bool = True,
+        timeout: float | None = None,
+    ) -> None:
         """Initialize the adapter with a structure and some parameters.
 
         Args:
-            structure: An input structure.
+            structure (Structure): An input structure.
             min_cell_size (int): The minimum cell size wanted. Defaults to 1.
             max_cell_size (int): The maximum cell size wanted. Defaults to 1.
             symm_prec (float): Symmetry precision. Defaults to 0.1.
@@ -124,9 +127,9 @@ def __init__(
         self.check_ordered_symmetry = check_ordered_symmetry
         self.timeout = timeout
 
-    def run(self):
+    def run(self) -> None:
         """Run the enumeration."""
-        # Create a temporary directory for working.
+        # Work in a temporary directory
         with ScratchDir(".") as tmp_dir:
             logger.debug(f"Temp dir : {tmp_dir}")
             # Generate input files
@@ -139,31 +142,32 @@ def run(self):
             else:
                 raise EnumError("Unable to enumerate structure.")
 
-    def _gen_input_file(self):
+    def _gen_input_file(self) -> None:
         """Generate the necessary struct_enum.in file for enumlib. See enumlib
         documentation for details.
         """
         coord_format = "{:.6f} {:.6f} {:.6f}"
-        # Using symmetry finder, get the symmetrically distinct sites.
+        # Use symmetry finder to get the symmetrically distinct sites
         fitter = SpacegroupAnalyzer(self.structure, self.symm_prec)
         symmetrized_structure = fitter.get_symmetrized_structure()
         logger.debug(
             f"Spacegroup {fitter.get_space_group_symbol()} ({fitter.get_space_group_number()}) "
             f"with {len(symmetrized_structure.equivalent_sites)} distinct sites"
         )
-        """Enumlib doesn"t work when the number of species get too large. To
-        simplify matters, we generate the input file only with disordered sites
-        and exclude the ordered sites from the enumeration. The fact that
-        different disordered sites with the exact same species may belong to
-        different equivalent sites is dealt with by having determined the
-        spacegroup earlier and labelling the species differently.
-        """
-        # index_species and index_amounts store mappings between the indices
+
+        # Enumlib doesn"t work when the number of species get too large. To
+        # simplify matters, we generate the input file only with disordered sites
+        # and exclude the ordered sites from the enumeration. The fact that
+        # different disordered sites with the exact same species may belong to
+        # different equivalent sites is dealt with by having determined the
+        # spacegroup earlier and labelling the species differently.
+
+        # `index_species` and `index_amounts` store mappings between the indices
         # used in the enum input file, and the actual species and amounts.
         index_species = []
         index_amounts = []
 
-        # Stores the ordered sites, which are not enumerated.
+        # Store the ordered sites, which are not enumerated.
         ordered_sites = []
         disordered_sites = []
         coord_str = []
@@ -191,24 +195,24 @@ def _gen_input_file(self):
                     coord_str.append(f"{coord_format.format(*site.coords)} {sp_label}")
                 disordered_sites.append(sites)
 
-        def get_sg_info(ss):
+        def get_sg_number(ss) -> int:
             finder = SpacegroupAnalyzer(Structure.from_sites(ss), self.symm_prec)
             return finder.get_space_group_number()
 
-        target_sg_num = get_sg_info(list(symmetrized_structure))
+        target_sg_num = get_sg_number(list(symmetrized_structure))
         curr_sites = list(itertools.chain.from_iterable(disordered_sites))
-        sg_num = get_sg_info(curr_sites)
+        sg_num = get_sg_number(curr_sites)
         ordered_sites = sorted(ordered_sites, key=len)
         logger.debug(f"Disordered sites has sg # {sg_num}")
         self.ordered_sites = []
 
-        # progressively add ordered sites to our disordered sites
+        # Progressively add ordered sites to our disordered sites
         # until we match the symmetry of our input structure
         if self.check_ordered_symmetry:
             while sg_num != target_sg_num and len(ordered_sites) > 0:
                 sites = ordered_sites.pop(0)
                 temp_sites = list(curr_sites) + sites
-                new_sg_num = get_sg_info(temp_sites)
+                new_sg_num = get_sg_number(temp_sites)
                 if sg_num != new_sg_num:
                     logger.debug(f"Adding {sites[0].specie} in enum. New sg # {new_sg_num}")
                     index_species.append(sites[0].specie)
@@ -274,12 +278,18 @@ def get_sg_info(ss):
                 min_conc = math.floor(conc * base)
                 output.append(f"{min_conc - 1} {min_conc + 1} {base}")
         output.append("")
+
         logger.debug("Generated input file:\n" + "\n".join(output))
         with open("struct_enum.in", mode="w", encoding="utf-8") as file:
             file.write("\n".join(output))
 
-    def _run_multienum(self):
-        with subprocess.Popen([enum_cmd], stdout=subprocess.PIPE, stdin=subprocess.PIPE, close_fds=True) as process:
+    def _run_multienum(self) -> int:
+        """Run enumlib to get multiple structure.
+
+        Returns:
+            int: number of structures.
+        """
+        with subprocess.Popen([ENUM_CMD], stdout=subprocess.PIPE, stdin=subprocess.PIPE, close_fds=True) as process:
             if self.timeout:
                 timed_out = False
                 timer = Timer(self.timeout * 60, lambda p: p.kill(), [process])
@@ -308,25 +318,26 @@ def _run_multienum(self):
         logger.debug(f"Enumeration resulted in {count} structures")
         return count
 
-    def _get_structures(self, num_structs):
-        structs = []
+    def _get_structures(self, num_structs: int) -> list[Structure]:
+        structs: list[Structure] = []
 
-        if ".py" in makestr_cmd:
-            options = ["-input", "struct_enum.out", str(1), str(num_structs)]
+        if ".py" in MAKESTR_CMD:
+            options = ("-input", "struct_enum.out", str(1), str(num_structs))
         else:
-            options = ["struct_enum.out", str(0), str(num_structs - 1)]
+            options = ("struct_enum.out", str(0), str(num_structs - 1))
 
         with subprocess.Popen(
-            [makestr_cmd, *options],
+            [MAKESTR_CMD, *options],
             stdout=subprocess.PIPE,
             stdin=subprocess.PIPE,
             close_fds=True,
         ) as rs:
             _stdout, stderr = rs.communicate()
+
         if stderr:
             logger.warning(stderr.decode())
 
-        # sites retrieved from enumlib will lack site properties
+        # Sites retrieved from enumlib will lack site properties
         # to ensure consistency, we keep track of what site properties
         # are missing and set them to None
         # TODO: improve this by mapping ordered structure to original