Merge pull request #488 from idaholab/mat_append

Made method to auto insert new materials to data_inputs.

doc/source/changelog.rst

@@ -6,6 +6,7 @@ MontePy Changelog
 **Features Added**
 
 * Write problems to either file paths or streams (file handles) with MCNP_Problem.write_problem() (:issue:`492`).
+* When adding a material to problem.materials it will also be added to problem.data_inputs, ensuring it is printed to the file (:pull:`488`).
 
 **Bug Fixes**
 

doc/source/developing.rst

@@ -397,6 +397,12 @@ For example the init function for ``Cells``
         def __init__(self, cells=None):
             super().__init__(montepy.Cell, cells)
 
+Collection: :class:`~montepy.numbered_object_collection.NumberedDataObjectCollection`
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+This is a subclass of :class:`~montepy.numbered_object_collection.NumberedObjectCollection`,
+which is designed for :class:`~montepy.data_inputs.data_input.DataInputAbstract` instances.
+It is a wrapper that will ensure that all of its items are also in :func:`~montepy.mcnp_problem.MCNP_Problem.data_inputs`.
+
 
 Numbered Object :class:`~montepy.numbered_mcnp_object.Numbered_MCNP_Object`
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

montepy/materials.py

@@ -1,14 +1,18 @@
 # Copyright 2024, Battelle Energy Alliance, LLC All Rights Reserved.
 import montepy
-from montepy.numbered_object_collection import NumberedObjectCollection
+from montepy.numbered_object_collection import NumberedDataObjectCollection
 
 Material = montepy.data_inputs.material.Material
 
 
-class Materials(NumberedObjectCollection):
+class Materials(NumberedDataObjectCollection):
     """
     A container of multiple :class:`~montepy.data_inputs.material.Material` instances.
 
+    .. note::
+        When items are added to this (and this object is linked to a problem),
+        they will also be added to :func:`montepy.mcnp_problem.MCNP_Problem.data_inputs`.
+
     :param objects: the list of materials to start with if needed
     :type objects: list
     """

montepy/mcnp_problem.py

@@ -286,7 +286,7 @@ def parse_input(self, check_input=False, replace=True):
                             else:
                                 raise e
                         if isinstance(obj, Material):
-                            self._materials.append(obj)
+                            self._materials.append(obj, False)
                         if isinstance(obj, transform.Transform):
                             self._transforms.append(obj)
                     if trailing_comment is not None and last_obj is not None:

montepy/numbered_object_collection.py

@@ -331,8 +331,7 @@ def __delitem__(self, idx):
             raise TypeError("index must be an int")
         obj = self[idx]
         self.__num_cache.pop(obj.number, None)
-        idx = self._objects.index(obj)
-        del self._objects[idx]
+        self._objects.remove(obj)
 
     def __setitem__(self, key, newvalue):
         if not isinstance(key, int):
@@ -364,10 +363,8 @@ def __iadd__(self, other):
                 )
             else:
                 self.__num_cache[obj.number] = obj
-        self._objects += other_list
-        if self._problem:
-            for obj in other_list:
-                obj.link_to_problem(self._problem)
+        for obj in other_list:
+            self.append(obj)
         return self
 
     def __contains__(self, other):
@@ -424,3 +421,82 @@ def items(
         """
         for o in self._objects:
             yield o.number, o
+
+
+class NumberedDataObjectCollection(NumberedObjectCollection):
+    def __init__(self, obj_class, objects=None, problem=None):
+        self._last_index = None
+        if problem and objects:
+            self._last_index = problem.data_inputs.index(objects[-1])
+        super().__init__(obj_class, objects, problem)
+
+    def append(self, obj, insert_in_data=True):
+        """Appends the given object to the end of this collection.
+
+        :param obj: the object to add.
+        :type obj: Numbered_MCNP_Object
+        :param insert_in_data: Whether to add the object to the linked problem's data_inputs.
+        :type insert_in_data: bool
+        :raises NumberConflictError: if this object has a number that is already in use.
+        """
+        if self._problem:
+            print(self._last_index, len(self))
+            if self._last_index:
+                index = self._last_index
+            elif len(self) > 0:
+                try:
+                    index = self._problem.data_inputs.index(self._objects[-1])
+                except ValueError:
+                    index = len(self._problem.data_inputs)
+            else:
+                index = len(self._problem.data_inputs)
+            if insert_in_data:
+                self._problem.data_inputs.insert(index + 1, obj)
+            self._last_index = index + 1
+        super().append(obj)
+
+    def __delitem__(self, idx):
+        if not isinstance(idx, int):
+            raise TypeError("index must be an int")
+        obj = self[idx]
+        super().__delitem__(idx)
+        if self._problem:
+            self._problem.data_inputs.remove(obj)
+
+    def remove(self, delete):
+        """
+        Removes the given object from the collection.
+
+        :param delete: the object to delete
+        :type delete: Numbered_MCNP_Object
+        """
+        super().remove(delete)
+        if self._problem:
+            self._problem.data_inputs.remove(delete)
+
+    def pop(self, pos=-1):
+        """
+        Pop the final items off of the collection
+
+        :param pos: The index of the element to pop from the internal list.
+        :type pos: int
+        :return: the final elements
+        :rtype: Numbered_MCNP_Object
+        """
+        if not isinstance(pos, int):
+            raise TypeError("The index for popping must be an int")
+        obj = self._objects.pop(pos)
+        super().pop(pos)
+        if self._problem:
+            self._problem.data_inputs.remove(obj)
+        return obj
+
+    def clear(self):
+        """
+        Removes all objects from this collection.
+        """
+        if self._problem:
+            for obj in self._objects:
+                self._problem.data_inputs.remove(obj)
+        self._last_index = None
+        super().clear()

tests/test_numbered_collection.py

@@ -4,6 +4,8 @@
 import montepy.cells
 from montepy.errors import NumberConflictError
 import unittest
+import pytest
+import os
 
 
 class TestNumberedObjectCollection(unittest.TestCase):
@@ -288,3 +290,85 @@ def test_str(self):
         ]
         for phrase in key_phrases:
             self.assertIn(phrase, repr(cells))
+
+
+# test data numbered object
+@pytest.fixture(scope="module")
+def read_simple_problem():
+    return montepy.read_input(os.path.join("tests", "inputs", "test.imcnp"))
+
+
+@pytest.fixture
+def cp_simple_problem(read_simple_problem):
+    return copy.deepcopy(read_simple_problem)
+
+
+def test_data_init(cp_simple_problem):
+    new_mats = montepy.materials.Materials(
+        list(cp_simple_problem.materials), problem=cp_simple_problem
+    )
+    assert list(new_mats) == list(cp_simple_problem.materials)
+
+
+def test_data_append(cp_simple_problem):
+    prob = cp_simple_problem
+    new_mat = copy.deepcopy(next(iter(prob.materials)))
+    new_mat.number = prob.materials.request_number()
+    prob.materials.append(new_mat)
+    assert new_mat in prob.materials
+    assert new_mat in prob.data_inputs
+    # trigger getting data_inputs end
+    prob.materials.clear()
+    prob.materials.append(new_mat)
+    assert new_mat in prob.materials
+    assert new_mat in prob.data_inputs
+    prob.data_inputs.clear()
+    prob.materials._last_index = None
+    new_mat = copy.deepcopy(next(iter(prob.materials)))
+    new_mat.number = prob.materials.request_number()
+    prob.materials.append(new_mat)
+    assert new_mat in prob.materials
+    assert new_mat in prob.data_inputs
+    # trigger getting index of last material
+    prob.materials._last_index = None
+    new_mat = copy.deepcopy(next(iter(prob.materials)))
+    new_mat.number = prob.materials.request_number()
+    prob.materials.append(new_mat)
+    assert new_mat in prob.materials
+    assert new_mat in prob.data_inputs
+
+
+def test_data_remove(cp_simple_problem):
+    prob = cp_simple_problem
+    old_mat = next(iter(prob.materials))
+    prob.materials.remove(old_mat)
+    assert old_mat not in prob.materials
+    assert old_mat not in prob.data_inputs
+
+
+def test_data_delete(cp_simple_problem):
+    prob = cp_simple_problem
+    old_mat = next(iter(prob.materials))
+    del prob.materials[old_mat.number]
+    assert old_mat not in prob.materials
+    assert old_mat not in prob.data_inputs
+    with pytest.raises(TypeError):
+        del prob.materials["foo"]
+
+
+def test_data_clear(cp_simple_problem):
+    data_len = len(cp_simple_problem.data_inputs)
+    mat_len = len(cp_simple_problem.materials)
+    cp_simple_problem.materials.clear()
+    assert len(cp_simple_problem.materials) == 0
+    assert len(cp_simple_problem.data_inputs) == data_len - mat_len
+
+
+def test_data_pop(cp_simple_problem):
+    old_mat = next(reversed(list(cp_simple_problem.materials)))
+    popper = cp_simple_problem.materials.pop()
+    assert popper is old_mat
+    assert old_mat not in cp_simple_problem.materials
+    assert old_mat not in cp_simple_problem.data_inputs
+    with pytest.raises(TypeError):
+        cp_simple_problem.materials.pop("foo")