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()