Merge pull request #218 from nlesc-nano/periodic

ENH: Automatically set the `MultiMolecule.lattice` property in ARMC

FOX/armc/package_manager.py

@@ -63,6 +63,7 @@ class PkgDict(TypedDict):
 
 
 T = TypeVar('T')
+RT = TypeVar('RT', bound=Result)
 
 MolLike = Iterable[Tuple[float, float, float]]
 Value = Tuple[PkgDict, ...]
@@ -322,7 +323,7 @@ def __call__(
         """
         # Construct the logger
         if logger is None:
-            logger = DummyLogger()
+            logger = cast(Logger, DummyLogger())
 
         # Check if a hook has been specified
         if self.hook is not None:
@@ -343,8 +344,11 @@ def __call__(
 
     @staticmethod
     @schedule
-    def assemble_job(job: PkgDict, old_results: Optional[Result] = None,
-                     name: Optional[str] = None) -> PromisedObject:
+    def assemble_job(
+        job: PkgDict,
+        old_results: Optional[Result] = None,
+        name: Optional[str] = None,
+    ) -> PromisedObject:
         """Create a :class:`PromisedObject` from a qmflow :class:`Package` instance."""
         job_name = name if name is not None else ''
         obj_type = job['type']
@@ -357,7 +361,7 @@ def assemble_job(job: PkgDict, old_results: Optional[Result] = None,
 
         if isinstance(obj_type, CP2K) and isinstance(old_results, CP2K_Result):
             try:
-                lattice = old_results.lattice
+                lattice: np.ndarray[Any, np.dtype[np.float64]] = old_results.lattice
                 assert lattice is not None
             except (AssertionError, FileNotFoundError):
                 pass
@@ -409,10 +413,19 @@ def update_settings(self, dct: Sequence[Tuple[str, Mapping]], new_keys: bool = T
                     del df['guess']
                 df.update(df_update)
 
+    @overload
+    @staticmethod
+    def _extract_mol(results: None, logger: Logger) -> Tuple[None, None]: ...
+    @overload  # noqa: E301
     @staticmethod
     def _extract_mol(
-        results: Optional[Iterable[Result]], logger: Logger
-    ) -> Union[Tuple[None, None], Tuple[List[MultiMolecule], List[Any]]]:
+        results: Iterable[RT], logger: Logger
+    ) -> Tuple[None, None] | Tuple[List[MultiMolecule], List[RT]]: ...
+    @staticmethod  # noqa: E301
+    def _extract_mol(
+        results: None | Iterable[RT],
+        logger: Logger,
+    ) -> Tuple[None, None] | Tuple[List[MultiMolecule], List[RT]]:
         """Create a list of MultiMolecule from the passed **results**."""
         # `noodles.run_parallel()` can return `None` under certain circumstances
         if results is None:
@@ -421,16 +434,21 @@ def _extract_mol(
         mol_list = []
         results_list = list(results)
         for result in results_list:
+            try:
+                lattice: None | np.ndarray[Any, np.dtype[np.float64]] = result.lattice
+                assert lattice is not None
+            except (AssertionError, FileNotFoundError):
+                lattice = None
+
             try:  # Construct and return a MultiMolecule object
                 path: str = get_xyz_path(result.archive['work_dir'])  # type: ignore
                 mol = MultiMolecule.from_xyz(path)
+                mol.lattice = lattice
                 mol.round(3, inplace=True)
                 mol_list.append(mol)
-
             except XYZError:  # The .xyz file is unreadable for some reason
                 logger.warning(f"Failed to parse {path!r}")
                 return None, None
-
             except Exception as ex:
                 logger.warning(ex)
                 return None, None