Merge pull request #2207 from shimwell/type_hinting_for_results

type hints added to openmc.deplete.Results

openmc/checkvalue.py

@@ -1,8 +1,13 @@
 import copy
+import os
+from typing import Union
 from collections.abc import Iterable
 
 import numpy as np
 
+# Type for arguments that accept file paths
+PathLike = Union[str, os.PathLike]
+
 
 def check_type(name, value, expected_type, expected_iter_type=None, *, none_ok=False):
     """Ensure that an object is of an expected type. Optionally, if the object is

openmc/deplete/results.py

@@ -1,6 +1,7 @@
 import numbers
 import bisect
 import math
+from typing import Iterable, Optional, Tuple, Union
 from warnings import warn
 
 import h5py
@@ -11,11 +12,12 @@ import openmc.checkvalue as cv
 from openmc.data.library import DataLibrary
 from openmc.material import Material, Materials
 from openmc.exceptions import DataError
+from openmc.checkvalue import PathLike
 
 __all__ = ["Results", "ResultsList"]
 
 
-def _get_time_as(seconds, units):
+def _get_time_as(seconds: float, units: str) -> float:
     """Converts the time in seconds to time in different units
 
     Parameters
@@ -70,7 +72,7 @@ class Results(list):
 
 
     @classmethod
-    def from_hdf5(cls, filename):
+    def from_hdf5(cls, filename: PathLike):
         """Load in depletion results from a previous file
 
         Parameters
@@ -91,7 +93,13 @@ class Results(list):
         )
         return cls(filename)
 
-    def get_atoms(self, mat, nuc, nuc_units="atoms", time_units="s"):
+    def get_atoms(
+        self,
+        mat: Union[Material, str],
+        nuc: str,
+        nuc_units: str = "atoms",
+        time_units: str = "s"
+    ) -> Tuple[np.ndarray, np.ndarray]:
         """Get number of nuclides over time from a single material
 
         .. note::
@@ -155,7 +163,12 @@ class Results(list):
 
         return times, concentrations
 
-    def get_reaction_rate(self, mat, nuc, rx):
+    def get_reaction_rate(
+        self,
+        mat: Union[Material, str],
+        nuc: str,
+        rx: str
+    ) -> Tuple[np.ndarray, np.ndarray]:
         """Get reaction rate in a single material/nuclide over time
 
         .. note::
@@ -200,7 +213,7 @@ class Results(list):
 
         return times, rates
 
-    def get_keff(self, time_units='s'):
+    def get_keff(self, time_units: str = 's') -> Tuple[np.ndarray, np.ndarray]:
         """Evaluates the eigenvalue from a results list.
 
         .. versionadded:: 0.13.1
@@ -236,12 +249,12 @@ class Results(list):
         times = _get_time_as(times, time_units)
         return times, eigenvalues
 
-    def get_eigenvalue(self, time_units='s'):
+    def get_eigenvalue(self, time_units: str = 's') -> Tuple[np.ndarray, np.ndarray]:
         warn("The get_eigenvalue(...) function has been renamed get_keff and "
              "will be removed in a future version of OpenMC.", FutureWarning)
         return self.get_keff(time_units)
 
-    def get_depletion_time(self):
+    def get_depletion_time(self) -> np.ndarray:
         """Return an array of the average time to deplete a material
 
         .. note::
@@ -269,7 +282,7 @@ class Results(list):
             times[ix] = res.proc_time
         return times
 
-    def get_times(self, time_units="d") -> np.ndarray:
+    def get_times(self, time_units: str = "d") -> np.ndarray:
         """Return the points in time that define the depletion schedule
 
         .. versionadded:: 0.12.1
@@ -298,7 +311,7 @@ class Results(list):
         return _get_time_as(times, time_units)
 
     def get_step_where(
-        self, time, time_units="d", atol=1e-6, rtol=1e-3
+        self, time, time_units: str = "d", atol: float = 1e-6, rtol: float = 1e-3
     ) -> int:
         """Return the index closest to a given point in time
 
@@ -358,7 +371,11 @@ class Results(list):
                 time, time_units, atol, rtol)
         )
 
-    def export_to_materials(self, burnup_index, nuc_with_data=None) -> Materials:
+    def export_to_materials(
+        self,
+        burnup_index: int,
+        nuc_with_data: Optional[Iterable[str]] = None
+    ) -> Materials:
         """Return openmc.Materials object based on results at a given step
 
         .. versionadded:: 0.12.1

openmc/material.py

@@ -3,7 +3,6 @@ from collections.abc import Iterable
 from copy import deepcopy
 from numbers import Real
 from pathlib import Path
-import os
 import re
 import typing  # imported separately as py3.8 requires typing.Iterable
 import warnings
@@ -18,6 +17,7 @@ import openmc.data
 import openmc.checkvalue as cv
 from ._xml import clean_indentation, reorder_attributes
 from .mixin import IDManagerMixin
+from openmc.checkvalue import PathLike
 
 
 # Units for density supported by OpenMC
@@ -1382,7 +1382,7 @@ class Materials(cv.CheckedList):
         for material in self:
             material.make_isotropic_in_lab()
 
-    def export_to_xml(self, path: Union[str, os.PathLike] = 'materials.xml'):
+    def export_to_xml(self, path: PathLike = 'materials.xml'):
         """Export material collection to an XML file.
 
         Parameters
@@ -1429,7 +1429,7 @@ class Materials(cv.CheckedList):
             fh.write('</materials>\n')
 
     @classmethod
-    def from_xml(cls, path: Union[str, os.PathLike] = 'materials.xml'):
+    def from_xml(cls, path: PathLike = 'materials.xml'):
         """Generate materials collection from XML file
 
         Parameters

openmc/settings.py

@@ -13,6 +13,7 @@ import openmc.checkvalue as cv
 
 from . import RegularMesh, Source, VolumeCalculation, WeightWindows
 from ._xml import clean_indentation, get_text, reorder_attributes
+from openmc.checkvalue import PathLike
 
 
 class RunMode(Enum):
@@ -1535,7 +1536,7 @@ class Settings:
         if text is not None:
             self.max_tracks = int(text)
 
-    def export_to_xml(self, path: Union[str, os.PathLike] = 'settings.xml'):
+    def export_to_xml(self, path: PathLike = 'settings.xml'):
         """Export simulation settings to an XML file.
 
         Parameters
@@ -1607,7 +1608,7 @@ class Settings:
         tree.write(str(p), xml_declaration=True, encoding='utf-8')
 
     @classmethod
-    def from_xml(cls, path: Union[str, os.PathLike] = 'settings.xml'):
+    def from_xml(cls, path: PathLike = 'settings.xml'):
         """Generate settings from XML file
 
         .. versionadded:: 0.13.0