From 72d6fd2d891be399ffa1bb42b9b53599edada75b Mon Sep 17 00:00:00 2001
From: Janosh Riebesell <janosh.riebesell@gmail.com>
Date: Mon, 11 Mar 2024 08:05:26 +0100
Subject: [PATCH] add type hints

---
 pymatgen/apps/borg/queen.py  | 22 +++++++++++++---------
 tests/apps/borg/test_hive.py | 12 ++++++------
 2 files changed, 19 insertions(+), 15 deletions(-)

diff --git a/pymatgen/apps/borg/queen.py b/pymatgen/apps/borg/queen.py
index 4e8058df246..d5a257d5007 100644
--- a/pymatgen/apps/borg/queen.py
+++ b/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,25 +64,25 @@ 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))
 
@@ -86,7 +90,7 @@ 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:
diff --git a/tests/apps/borg/test_hive.py b/tests/apps/borg/test_hive.py
index 817e0913481..a174803bf00 100644
--- a/tests/apps/borg/test_hive.py
+++ b/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()