From 5b6c6667cbc082d438ea38708fd4da34b89b1a18 Mon Sep 17 00:00:00 2001 From: lewisgross1296 Date: Tue, 16 Aug 2022 13:43:15 -0500 Subject: [PATCH] put _get_tally_index(self,key) back and forgot return in __getitem__ --- openmc/lib/tally.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/openmc/lib/tally.py b/openmc/lib/tally.py index dc782539bc..5ddee90e49 100644 --- a/openmc/lib/tally.py +++ b/openmc/lib/tally.py @@ -396,6 +396,7 @@ class _TallyMapping(Mapping): except (AllocationError, InvalidIDError) as e: # __contains__ expects a KeyError to work correctly raise KeyError(str(e)) + return Tally(index=index.value) def __iter__(self): for i in range(len(self)): @@ -411,4 +412,14 @@ class _TallyMapping(Mapping): """Delete a tally from tally vector and remove the ID,index pair from tally""" _dll.openmc_remove_tally(self[key]._index) + def _get_tally_index(self,key): + """Given the ID, return the index""" + index = c_int32() + try: + _dll.openmc_get_tally_index(key, index) + except (AllocationError, InvalidIDError) as e: + # __contains__ expects a KeyError to work correctly + raise KeyError(str(e)) + return index.value + tallies = _TallyMapping()