Implementing weakref

doc/source/changelog.rst

@@ -1,6 +1,13 @@
 MontePy Changelog
 =================
 
+#Next Version#
+--------------
+
+**Performance Improvement**
+
+* Fixed cyclic memory reference that lead to memory leak in ``copy.deepcopy`` (:issue:`514`).
+
 0.4.1
 --------------
 

htmlcov

@@ -0,0 +1 @@
+/home/mgale/winHome/Downloads/htmlcov/

montepy/cells.py

@@ -89,6 +89,19 @@ def allow_mcnp_volume_calc(self, value):
             raise TypeError("allow_mcnp_volume_calc must be set to a bool")
         self._volume.is_mcnp_calculated = value
 
+    def link_to_problem(self, problem):
+        """Links the input to the parent problem for this input.
+
+        This is done so that inputs can find links to other objects.
+
+        :param problem: The problem to link this input to.
+        :type problem: MCNP_Problem
+        """
+        super().link_to_problem(problem)
+        inputs_to_property = montepy.Cell._INPUTS_TO_PROPERTY
+        for attr, _ in inputs_to_property.values():
+            getattr(self, attr).link_to_problem(problem)
+
     def update_pointers(
         self, cells, materials, surfaces, data_inputs, problem, check_input=False
     ):

montepy/mcnp_object.py

@@ -18,6 +18,7 @@
 import numpy as np
 import textwrap
 import warnings
+import weakref
 
 
 class MCNP_Object(ABC):
@@ -35,7 +36,7 @@ class MCNP_Object(ABC):
     """
 
     def __init__(self, input, parser):
-        self._problem = None
+        self._problem_ref = None
         self._parameters = ParametersNode()
         self._input = None
         if input:
@@ -245,9 +246,25 @@ def link_to_problem(self, problem):
         :param problem: The problem to link this input to.
         :type problem: MCNP_Problem
         """
-        if not isinstance(problem, montepy.mcnp_problem.MCNP_Problem):
+        if not isinstance(problem, (montepy.mcnp_problem.MCNP_Problem, type(None))):
             raise TypeError("problem must be an MCNP_Problem")
-        self._problem = problem
+        if problem is None:
+            self._problem_ref = None
+        else:
+            self._problem_ref = weakref.ref(problem)
+
+    @property
+    def _problem(self):
+        if self._problem_ref is not None:
+            return self._problem_ref()
+        return None
+
+    @_problem.setter
+    def _problem(self, problem):
+        if problem is None:
+            self._problem_ref = None
+            return
+        self.link_to_problem(problem)
 
     @property
     def trailing_comment(self):
@@ -429,3 +446,14 @@ def allowed_keywords(self):  # pragma: no cover
             stacklevel=2,
         )
         return set()
+
+    def __getstate__(self):
+        state = self.__dict__.copy()
+        weakref_key = "_problem_ref"
+        if weakref_key in state:
+            del state[weakref_key]
+        return state
+
+    def __setstate__(self, crunchy_data):
+        crunchy_data["_problem_ref"] = None
+        self.__dict__.update(crunchy_data)

montepy/mcnp_problem.py

@@ -1,7 +1,8 @@
 # Copyright 2024, Battelle Energy Alliance, LLC All Rights Reserved.
-import os
+import copy
 from enum import Enum
 import itertools
+
 from montepy.data_inputs import mode, transform
 from montepy._cell_data_control import CellDataPrintController
 from montepy.cell import Cell
@@ -16,6 +17,9 @@
 from montepy.input_parser.input_file import MCNP_InputFile
 from montepy.universes import Universes
 from montepy.transforms import Transforms
+import montepy
+
+import os
 import warnings
 
 
@@ -37,6 +41,7 @@ def __init__(self, destination):
             self._input_file = MCNP_InputFile(destination)
         self._title = None
         self._message = None
+        self.__unpickled = False
         self._print_in_data_block = CellDataPrintController()
         self._original_inputs = []
         self._cells = Cells(problem=self)
@@ -48,6 +53,10 @@ def __init__(self, destination):
         self._mcnp_version = DEFAULT_VERSION
         self._mode = mode.Mode()
 
+    def __setstate__(self, nom_nom):
+        self.__dict__.update(nom_nom)
+        self.__unpickled = True
+
     @property
     def original_inputs(self):
         """
@@ -64,6 +73,42 @@ def original_inputs(self):
         """
         return self._original_inputs
 
+    def __relink_objs(self):
+        if self.__unpickled:
+            self._cells.link_to_problem(self)
+            for collection in {"_surfaces", "_data_inputs"}:
+                collection = getattr(self, collection)
+                if isinstance(
+                    collection,
+                    montepy.numbered_object_collection.NumberedObjectCollection,
+                ):
+                    collection.link_to_problem(self)
+                else:
+                    for obj in collection:
+                        obj.link_to_problem(self)
+            self.__unpickled = False
+
+    def __unlink_objs(self):
+        for collection in {"_surfaces", "_data_inputs"}:
+            collection = getattr(self, collection)
+            if isinstance(
+                collection,
+                montepy.numbered_object_collection.NumberedObjectCollection,
+            ):
+                collection.link_to_problem(None)
+            else:
+                for obj in collection:
+                    obj.link_to_problem(None)
+
+    def __deepcopy__(self, memo):
+        cls = type(self)
+        result = cls.__new__(cls)
+        memo[id(self)] = result
+        for k, v in self.__dict__.items():
+            setattr(result, k, copy.deepcopy(v, memo))
+        result.__unlink_objs()
+        return result
+
     @property
     def cells(self):
         """
@@ -72,6 +117,7 @@ def cells(self):
         :return: a collection of the Cell objects, ordered by the order they were in the input file.
         :rtype: Cells
         """
+        self.__relink_objs()
         return self._cells
 
     @cells.setter
@@ -141,6 +187,7 @@ def surfaces(self):
         :return: a collection of the Surface objects, ordered by the order they were in the input file.
         :rtype: Surfaces
         """
+        self.__relink_objs()
         return self._surfaces
 
     @property
@@ -151,6 +198,7 @@ def materials(self):
         :return: a colection of the Material objects, ordered by the order they were in the input file.
         :rtype: Materials
         """
+        self.__relink_objs()
         return self._materials
 
     @materials.setter
@@ -159,6 +207,7 @@ def materials(self, mats):
             raise TypeError("materials must be of type list and Materials")
         if isinstance(mats, list):
             mats = Materials(mats)
+        mats.link_to_problem(self)
         self._materials = mats
 
     @property
@@ -183,6 +232,7 @@ def data_inputs(self):
         :return: a list of the :class:`~montepy.data_cards.data_card.DataCardAbstract` objects, ordered by the order they were in the input file.
         :rtype: list
         """
+        self.__relink_objs()
         return self._data_inputs
 
     @property
@@ -298,7 +348,7 @@ def parse_input(self, check_input=False, replace=True):
                         if isinstance(obj, Material):
                             self._materials.append(obj, False)
                         if isinstance(obj, transform.Transform):
-                            self._transforms.append(obj)
+                            self._transforms.append(obj, False)
                     if trailing_comment is not None and last_obj is not None:
                         obj._grab_beginning_comment(trailing_comment)
                         last_obj._delete_trailing_comment()
@@ -400,9 +450,9 @@ def add_cell_children_to_problem(self):
         surfaces = sorted(surfaces)
         materials = sorted(materials)
         transforms = sorted(transforms)
-        self._surfaces = Surfaces(surfaces)
-        self._materials = Materials(materials)
-        self._transforms = Transforms(transforms)
+        self._surfaces = Surfaces(surfaces, problem=self)
+        self._materials = Materials(materials, problem=self)
+        self._transforms = Transforms(transforms, problem=self)
         self._data_inputs = sorted(set(self._data_inputs + materials + transforms))
 
     def write_problem(self, destination, overwrite=False):

montepy/numbered_object_collection.py

@@ -1,6 +1,8 @@
 # Copyright 2024, Battelle Energy Alliance, LLC All Rights Reserved.
 from abc import ABC, abstractmethod
 import typing
+import weakref
+
 import montepy
 from montepy.numbered_mcnp_object import Numbered_MCNP_Object
 from montepy.errors import *
@@ -41,7 +43,9 @@ def __init__(self, obj_class, objects=None, problem=None):
         assert issubclass(obj_class, Numbered_MCNP_Object)
         self._obj_class = obj_class
         self._objects = []
-        self._problem = problem
+        self._problem_ref = None
+        if problem is not None:
+            self._problem_ref = weakref.ref(problem)
         if objects:
             if not isinstance(objects, list):
                 raise TypeError("NumberedObjectCollection must be built from a list")
@@ -68,9 +72,31 @@ def link_to_problem(self, problem):
         :param problem: The problem to link this card to.
         :type problem: MCNP_Problem
         """
-        if not isinstance(problem, montepy.mcnp_problem.MCNP_Problem):
+        if not isinstance(problem, (montepy.mcnp_problem.MCNP_Problem, type(None))):
             raise TypeError("problem must be an MCNP_Problem")
-        self._problem = problem
+        if problem is None:
+            self._problem_ref = None
+        else:
+            self._problem_ref = weakref.ref(problem)
+        for obj in self:
+            obj.link_to_problem(problem)
+
+    @property
+    def _problem(self):
+        if self._problem_ref is not None:
+            return self._problem_ref()
+        return None
+
+    def __getstate__(self):
+        state = self.__dict__.copy()
+        weakref_key = "_NumberedObjectCollection__problem"
+        if weakref_key in state:
+            del state[weakref_key]
+        return state
+
+    def __setstate__(self, crunchy_data):
+        crunchy_data["_NumberedObjectCollection__problem"] = None
+        self.__dict__.update(crunchy_data)
 
     @property
     def numbers(self):

montepy/transforms.py

@@ -1,9 +1,9 @@
 # Copyright 2024, Battelle Energy Alliance, LLC All Rights Reserved.
-from montepy.numbered_object_collection import NumberedObjectCollection
+from montepy.numbered_object_collection import NumberedDataObjectCollection
 from montepy.data_inputs.transform import Transform
 
 
-class Transforms(NumberedObjectCollection):
+class Transforms(NumberedDataObjectCollection):
     """
     A container of multiple :class:`~montepy.data_inputs.transform.Transform` instances.
     """

tests/test_interface.py

@@ -32,3 +32,10 @@ def test_pickling_all_models(self):
                     new_problem = pickle.load(fh)
                 data = pickle.dumps(problem)
                 new_problem = pickle.loads(data)
+                for attr in {"cells", "surfaces", "materials", "data_inputs"}:
+                    assert len(getattr(problem, attr)) == len(
+                        getattr(new_problem, attr)
+                    )
+                    for obj in getattr(new_problem, attr):
+                        assert obj._problem is not None
+                        assert obj._problem is new_problem

tests/test_numbered_collection.py

@@ -118,8 +118,10 @@ def test_append_renumber(self):
         self.assertEqual(len(cells), size + 2)
 
     def test_append_renumber_problems(self):
+        print(hex(id(self.simple_problem.materials._problem)))
         prob1 = copy.deepcopy(self.simple_problem)
         prob2 = copy.deepcopy(self.simple_problem)
+        print(hex(id(self.simple_problem.materials._problem)))
         # Delete Material 2, making its number available.
         prob2.materials.remove(prob2.materials[2])
         len_mats = len(prob2.materials)