Allow StepResult.get_material to accept integer material ID (#3872)

Co-authored-by: Paul Romano <paul.k.romano@gmail.com>

Author: AlvaroCubi

openmc/deplete/stepresult.py

@@ -12,8 +12,9 @@
 import numpy as np
 
 import openmc
-from openmc.mpi import comm, MPI
 from openmc.checkvalue import PathLike
+from openmc.mpi import MPI, comm
+
 from .reaction_rates import ReactionRates
 
 VERSION_RESULTS = (1, 2)
@@ -196,15 +197,15 @@ def distribute(self, local_materials, ranges):
         new.rates = self.rates[ranges]
         return new
 
-    def get_material(self, mat_id):
+    def get_material(self, mat_id: str | int) -> openmc.Material:
         """Return material object for given depleted composition
 
         .. versionadded:: 0.13.2
 
         Parameters
         ----------
-        mat_id : str
-            Material ID as a string
+        mat_id : str or int
+            Material ID as a string or integer
 
         Returns
         -------
@@ -217,6 +218,9 @@ def get_material(self, mat_id):
             If specified material ID is not found in the StepResult
 
         """
+        # Coerce to str since internal dictionaries use str keys
+        mat_id = str(mat_id)
+
         with warnings.catch_warnings():
             warnings.simplefilter('ignore', openmc.IDWarning)
             material = openmc.Material(material_id=int(mat_id))

tests/unit_tests/test_deplete_resultslist.py

@@ -1,10 +1,11 @@
 """Tests the Results class"""
 
-from pathlib import Path
 from math import inf
+from pathlib import Path
 
 import numpy as np
 import pytest
+
 import openmc.deplete
 
 
@@ -221,3 +222,11 @@ def test_stepresult_get_material(res):
     densities = mat1.get_nuclide_atom_densities()
     assert densities['Xe135'] == pytest.approx(1e-14)
     assert densities['U234'] == pytest.approx(1.00506e-05)
+
+
+def test_stepresult_get_material_mat_id_as_int(res):
+    # Get material at first timestep using int mat_id
+    step_result = res[0]
+    mat1 = step_result.get_material(1)
+    assert mat1.id == 1
+    assert mat1.volume == step_result.volume["1"]