From d1775181e905054677d904acb9f4a652d0195325 Mon Sep 17 00:00:00 2001
From: lewisgross1296 <ligross@wisc.edu>
Date: Fri, 12 Aug 2022 15:49:08 -0500
Subject: [PATCH] added deletion operator to TallyMapping class. used in tests
 and changed delete invalid id to access invalid id

---
 openmc/lib/tally.py          | 22 +++++++++++++++-------
 tests/unit_tests/test_lib.py | 13 ++++++++-----
 2 files changed, 23 insertions(+), 12 deletions(-)

diff --git a/openmc/lib/tally.py b/openmc/lib/tally.py
index 534e6ebaae..c1c8ab0941 100644
--- a/openmc/lib/tally.py
+++ b/openmc/lib/tally.py
@@ -390,13 +390,7 @@ class Tally(_FortranObjectWithID):
 
 class _TallyMapping(Mapping):
     def __getitem__(self, key):
-        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 Tally(index=index.value)
+        return Tally(index=self._get_tally_index(key))
 
     def __iter__(self):
         for i in range(len(self)):
@@ -408,4 +402,18 @@ class _TallyMapping(Mapping):
     def __repr__(self):
         return repr(dict(self))
 
+    def __delitem__(self,key):
+        """Delete a tally from tally vector and remove the ID,index pair from tally"""
+        _dll.openmc_remove_tally(self._get_tally_index(key))
+
+    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()
diff --git a/tests/unit_tests/test_lib.py b/tests/unit_tests/test_lib.py
index 96d54c416b..5063646f7d 100644
--- a/tests/unit_tests/test_lib.py
+++ b/tests/unit_tests/test_lib.py
@@ -339,16 +339,19 @@ def test_new_tally(lib_init):
     new_tally_with_id.scores = ['flux']
     assert len(openmc.lib.tallies) == 5
 
+
 def test_delete_tally(lib_init):
     # delete tally 10 which was added in the above test
     # check length is one less than before
-    openmc.lib.openmc_remove_tally(openmc.lib.get_tally_index(10))
+    del openmc.lib.tallies[10]
     assert len(openmc.lib.tallies) == 4
 
-def test_delete_invalid_id(lib_init):
-    # attempt to delete a tally that is guaranteed not to have a valid index
-    with pytest.raises(exc.InvalidIDError):
-        openmc.lib.openmc_remove_tally(np.max(id)+1)
+
+def test_invalid_tally_id(lib_init):
+    # attempt to access a tally that is guaranteed not to have a valid index
+    max_id = max(openmc.lib.tallies.keys())
+    with pytest.raises(KeyError):
+        openmc.lib.tallies[max_id+1]
 
 
 def test_tally_activate(lib_simulation_init):