Change DataLibrary to behave like list and add remove_by_material method

openmc/data/library.py

@@ -5,24 +5,29 @@ import pathlib
 import h5py
 
 import openmc
-from openmc.mixin import EqualityMixin
 from openmc._xml import clean_indentation, reorder_attributes
 
 
-class DataLibrary(EqualityMixin):
+class DataLibrary(list):
     """Collection of cross section data libraries.
 
-    Attributes
-    ----------
-    libraries : list of dict
-        List in which each item is a dictionary summarizing cross section data
-        from a single file. The dictionary has keys 'path', 'type', and
-        'materials'.
+    This clas behaves like a list where each item is a dictionary summarizing
+    cross section data from a single file. The dictionary has keys 'path',
+    'type', and 'materials'.
+
+    .. versionchanged:: 0.13.4
+        This class now behaves like a list rather than requiring you to access
+        the list of libraries through a special attribute.
 
     """
 
     def __init__(self):
-        self.libraries = []
+        super().__init__()
+
+    @property
+    def libraries(self):
+        # For backwards compatibility
+        return self
 
     def get_by_material(self, name, data_type='neutron'):
         """Return the library dictionary containing a given material.
@@ -43,11 +48,26 @@ class DataLibrary(EqualityMixin):
             the dictionary has keys 'path', 'type', and 'materials'.
 
         """
-        for library in self.libraries:
+        for library in self:
             if name in library['materials'] and data_type in library['type']:
                 return library
         return None
 
+    def remove_by_material(self, name: str, data_type='neutron'):
+        """Remove the library dictionary containing a specific material
+
+        Parameters
+        ----------
+        name : str
+            Name of material, e.g. 'Am241'
+        data_type : str
+            Name of data type, e.g. 'neutron', 'photon', 'wmp', or 'thermal'
+
+        """
+        library = self.get_by_material(name, data_type)
+        if library is not None:
+            self.remove(library)
+
     def register_file(self, filename):
         """Register a file with the data library.
 
@@ -77,7 +97,7 @@ class DataLibrary(EqualityMixin):
                 .format(path.name, self.__class__.__name__))
 
         library = {'path': str(path), 'type': filetype, 'materials': materials}
-        self.libraries.append(library)
+        self.append(library)
 
     def export_to_xml(self, path='cross_sections.xml'):
         """Export cross section data library to an XML file.
@@ -92,7 +112,7 @@ class DataLibrary(EqualityMixin):
 
         # Determine common directory for library paths
         common_dir = os.path.dirname(os.path.commonprefix(
-            [lib['path'] for lib in self.libraries]))
+            [lib['path'] for lib in self]))
         if common_dir == '':
             common_dir = '.'
 
@@ -100,7 +120,7 @@ class DataLibrary(EqualityMixin):
             dir_element = ET.SubElement(root, "directory")
             dir_element.text = os.path.realpath(common_dir)
 
-        for library in self.libraries:
+        for library in self:
             if library['type'] == "depletion_chain":
                 lib_element = ET.SubElement(root, "depletion_chain")
             else:

tests/unit_tests/test_data_misc.py

@@ -11,7 +11,7 @@ import openmc.data
 
 def test_data_library(tmpdir):
     lib = openmc.data.DataLibrary.from_xml()
-    for f in lib.libraries:
+    for f in lib:
         assert sorted(f.keys()) == ['materials', 'path', 'type']
 
     f = lib.get_by_material('U235')
@@ -22,6 +22,9 @@ def test_data_library(tmpdir):
     assert f['type'] == 'thermal'
     assert 'c_H_in_H2O' in f['materials']
 
+    lib.remove_by_material('Pu239')
+    assert lib.get_by_material('Pu239') is None
+
     filename = str(tmpdir.join('test.xml'))
     lib.export_to_xml(filename)
     assert os.path.exists(filename)
@@ -29,9 +32,9 @@ def test_data_library(tmpdir):
     new_lib = openmc.data.DataLibrary()
     directory = os.path.dirname(os.environ['OPENMC_CROSS_SECTIONS'])
     new_lib.register_file(os.path.join(directory, 'H1.h5'))
-    assert new_lib.libraries[-1]['type'] == 'neutron'
+    assert new_lib[-1]['type'] == 'neutron'
     new_lib.register_file(os.path.join(directory, 'c_Zr_in_ZrH.h5'))
-    assert new_lib.libraries[-1]['type'] == 'thermal'
+    assert new_lib[-1]['type'] == 'thermal'
 
 
 def test_depletion_chain_data_library(run_in_tmpdir):