add type hints

pymatgen/apps/borg/queen.py

@@ -9,10 +9,14 @@
 import logging
 import os
 from multiprocessing import Manager, Pool
+from typing import TYPE_CHECKING
 
 from monty.io import zopen
 from monty.json import MontyDecoder, MontyEncoder
 
+if TYPE_CHECKING:
+    from pathlib import Path
+
 logger = logging.getLogger("BorgQueen")
 
 
@@ -60,33 +64,33 @@ def parallel_assimilate(self, rootpath):
         status["count"] = 0
         status["total"] = len(valid_paths)
         logger.info(f"{len(valid_paths)} valid paths found.")
-        with Pool(self._num_drones) as p:
-            p.map(
+        with Pool(self._num_drones) as pool:
+            pool.map(
                 order_assimilation,
                 ((path, self._drone, data, status) for path in valid_paths),
             )
-            for d in data:
-                self._data.append(json.loads(d, cls=MontyDecoder))
+            for string in data:
+                self._data.append(json.loads(string, cls=MontyDecoder))
 
-    def serial_assimilate(self, rootpath):
+    def serial_assimilate(self, root: str | Path) -> None:
         """Assimilate the entire subdirectory structure in rootpath serially."""
         valid_paths = []
-        for parent, subdirs, files in os.walk(rootpath):
+        for parent, subdirs, files in os.walk(root):
             valid_paths.extend(self._drone.get_valid_paths((parent, subdirs, files)))
-        data = []
+        data: list[str] = []
         total = len(valid_paths)
         for idx, path in enumerate(valid_paths, 1):
             new_data = self._drone.assimilate(path)
             self._data.append(new_data)
-            logger.info(f"{idx}/{total} ({idx / total:.2%}) done")
+            logger.info(f"{idx}/{total} ({idx / total:.1%}) done")
         for json_str in data:
             self._data.append(json.loads(json_str, cls=MontyDecoder))
 
     def get_data(self):
         """Returns an list of assimilated objects."""
         return self._data
 
-    def save_data(self, filename):
+    def save_data(self, filename: str | Path) -> None:
         """Save the assimilated data to a file.
 
         Args:

tests/apps/borg/test_hive.py

@@ -76,16 +76,16 @@ def test_get_valid_paths(self):
     def test_assimilate(self):
         test_file = f"{TEST_FILES_DIR}/molecules/methane.log"
         entry = self.drone.assimilate(test_file)
-        for p in [
+        for param in [
             "functional",
             "basis_set",
             "charge",
             "spin_multiplicity",
             "route_parameters",
         ]:
-            assert p in entry.parameters
-        for p in ["corrections"]:
-            assert p in entry.data
+            assert param in entry.parameters
+        for param in ["corrections"]:
+            assert param in entry.data
 
         assert entry.reduced_formula == "H4C"
         assert entry.energy == approx(-39.9768775602)
@@ -94,8 +94,8 @@ def test_assimilate(self):
         assert entry.energy == approx(-39.9768775602)
         assert isinstance(entry, ComputedStructureEntry)
         assert entry.structure is not None
-        for p in ["properly_terminated", "stationary_type"]:
-            assert p in entry.data
+        for param in ["properly_terminated", "stationary_type"]:
+            assert param in entry.data
 
     def test_as_from_dict(self):
         dct = self.structure_drone.as_dict()