fastplotlib · kushalkolar · Apr 13, 2023 · Apr 3, 2023 · Apr 3, 2023 · Apr 3, 2023
diff --git a/examples/garbage_collection.ipynb b/examples/garbage_collection.ipynb
diff --git a/fastplotlib/graphics/_base.py b/fastplotlib/graphics/_base.py
@@ -1,4 +1,5 @@
 from typing import *
+import weakref
 from warnings import warn
 
 import numpy as np
@@ -14,6 +15,11 @@
 from dataclasses import dataclass
 
 
+# dict that holds all world objects for a given python kernel/session
+# Graphic objects only use proxies to WorldObjects
+WORLD_OBJECTS: Dict[str, WorldObject] = dict()  #: {hex id str: WorldObject}
+
+
 PYGFX_EVENTS = [
     "key_down",
     "key_up",
@@ -44,7 +50,8 @@ def __init_subclass__(cls, **kwargs):
 
 class Graphic(BaseGraphic):
     def __init__(
-            self, name: str = None):
+            self, name: str = None
+    ):
         """
 
         Parameters
@@ -58,10 +65,16 @@ def __init__(
         self.registered_callbacks = dict()
         self.present = PresentFeature(parent=self)
 
+        # store hex id str of Graphic instance mem location
+        self.loc: str = hex(id(self))
+
     @property
     def world_object(self) -> WorldObject:
-        """Associated pygfx WorldObject."""
-        return self._world_object
+        """Associated pygfx WorldObject. Always returns a proxy, real object cannot be accessed directly."""
+        return weakref.proxy(WORLD_OBJECTS[hex(id(self))])
+
+    def _set_world_object(self, wo: WorldObject):
+        WORLD_OBJECTS[hex(id(self))] = wo
 
     @property
     def position(self) -> Vector3:
@@ -75,7 +88,7 @@ def visible(self) -> bool:
         return self.world_object.visible
 
     @visible.setter
-    def visible(self, v) -> bool:
+    def visible(self, v: bool):
         """Access or change the visibility."""
         self.world_object.visible = v
 
@@ -100,6 +113,9 @@ def __repr__(self):
         else:
             return rval
 
+    def __del__(self):
+        del WORLD_OBJECTS[self.loc]
+
 
 class Interaction(ABC):
     """Mixin class that makes graphics interactive"""
@@ -216,8 +232,13 @@ def _event_handler(self, event):
 
                     # for now we only have line collections so this works
                     else:
-                        for i, item in enumerate(self._graphics):
-                            if item.world_object is event.pick_info["world_object"]:
+                        # get index of world object that made this event
+                        for i, item in enumerate(self.graphics):
+                            wo = WORLD_OBJECTS[item.loc]
+                            # we only store hex id of worldobject, but worldobject `pick_info` is always the real object
+                            # so if pygfx worldobject triggers an event by itself, such as `click`, etc., this will be
+                            # the real world object in the pick_info and not the proxy
+                            if wo is event.pick_info["world_object"]:
                                 indices = i
                     target_info.target._set_feature(feature=target_info.feature, new_data=target_info.new_data, indices=indices)
                 else:
@@ -264,22 +285,21 @@ class PreviouslyModifiedData:
     indices: Any
 
 
+COLLECTION_GRAPHICS: dict[str, Graphic] = dict()
+
+
 class GraphicCollection(Graphic):
     """Graphic Collection base class"""
 
     def __init__(self, name: str = None):
         super(GraphicCollection, self).__init__(name)
-        self._graphics: List[Graphic] = list()
-
-    @property
-    def world_object(self) -> Group:
-        """Returns the underling pygfx WorldObject."""
-        return self._world_object
+        self._graphics: List[str] = list()
 
     @property
     def graphics(self) -> Tuple[Graphic]:
-        """returns the Graphics within this collection"""
-        return tuple(self._graphics)
+        """The Graphics within this collection. Always returns a proxy to the Graphics."""
+        proxies = [weakref.proxy(COLLECTION_GRAPHICS[loc]) for loc in self._graphics]
+        return tuple(proxies)
 
     def add_graphic(self, graphic: Graphic, reset_index: True):
         """Add a graphic to the collection"""
@@ -289,17 +309,31 @@ def add_graphic(self, graphic: Graphic, reset_index: True):
                 f"You can only add {self.child_type} to a {self.__class__.__name__}, "
                 f"you are trying to add a {graphic.__class__.__name__}."
             )
-        self._graphics.append(graphic)
+
+        loc = hex(id(graphic))
+        COLLECTION_GRAPHICS[loc] = graphic
+
+        self._graphics.append(loc)
         if reset_index:
             self._reset_index()
         self.world_object.add(graphic.world_object)
 
     def remove_graphic(self, graphic: Graphic, reset_index: True):
         """Remove a graphic from the collection"""
-        self._graphics.remove(graphic)
+        self._graphics.remove(graphic.loc)
+
         if reset_index:
             self._reset_index()
-        self.world_object.remove(graphic)
+
+        self.world_object.remove(graphic.world_object)
+
+    def __del__(self):
+        self.world_object.clear()
+
+        for loc in self._graphics:
+            del COLLECTION_GRAPHICS[loc]
+
+        super().__del__()
 
     def _reset_index(self):
         for new_index, graphic in enumerate(self._graphics):
@@ -312,7 +346,7 @@ def __getitem__(self, key):
         if isinstance(key, slice):
             key = cleanup_slice(key, upper_bound=len(self))
             selection_indices = range(key.start, key.stop, key.step)
-            selection = self._graphics[key]
+            selection = self.graphics[key]
 
         # fancy-ish indexing
         elif isinstance(key, (tuple, list, np.ndarray)):
@@ -324,7 +358,7 @@ def __getitem__(self, key):
             selection = list()
 
             for ix in key:
-                selection.append(self._graphics[ix])
+                selection.append(self.graphics[ix])
 
             selection_indices = key
         else:
@@ -365,7 +399,7 @@ def __init__(
         selection_indices: Union[list, range]
             the corresponding indices from the parent GraphicCollection that were selected
         """
-        self._parent = parent
+        self._parent = weakref.proxy(parent)
         self._selection = selection
         self._selection_indices = selection_indices
 

diff --git a/fastplotlib/graphics/features/_base.py b/fastplotlib/graphics/features/_base.py
@@ -2,6 +2,7 @@
 from inspect import getfullargspec
 from warnings import warn
 from typing import *
+import weakref
 
 import numpy as np
 from pygfx import Buffer, Texture
@@ -71,7 +72,7 @@ def __init__(self, parent, data: Any, collection_index: int = None):
             if part of a collection, index of this graphic within the collection
 
         """
-        self._parent = parent
+        self._parent = weakref.proxy(parent)
 
         self._data = to_gpu_supported_dtype(data)
 

diff --git a/fastplotlib/graphics/image.py b/fastplotlib/graphics/image.py
@@ -100,11 +100,13 @@ def __init__(
             self.cmap = ImageCmapFeature(self, cmap)
             material = pygfx.ImageBasicMaterial(clim=(vmin, vmax), map=self.cmap())
 
-        self._world_object: pygfx.Image = pygfx.Image(
+        world_object = pygfx.Image(
             geometry,
             material
         )
 
+        self._set_world_object(world_object)
+
         self.data = ImageDataFeature(self, data)
         # TODO: we need to organize and do this better
         if isolated_buffer:
@@ -272,7 +274,8 @@ def __init__(
         start_ixs = [list(map(lambda c: c * chunk_size, chunk)) for chunk in chunks]
         stop_ixs = [list(map(lambda c: c + chunk_size, chunk)) for chunk in start_ixs]
 
-        self._world_object = pygfx.Group()
+        world_object = pygfx.Group()
+        self._set_world_object(world_object)
 
         if (vmin is None) or (vmax is None):
             vmin, vmax = quick_min_max(data)

diff --git a/fastplotlib/graphics/line.py b/fastplotlib/graphics/line.py
@@ -85,12 +85,14 @@ def __init__(
 
         self.thickness = ThicknessFeature(self, thickness)
 
-        self._world_object: pygfx.Line = pygfx.Line(
+        world_object: pygfx.Line = pygfx.Line(
             # self.data.feature_data because data is a Buffer
             geometry=pygfx.Geometry(positions=self.data(), colors=self.colors()),
             material=material(thickness=self.thickness(), vertex_colors=True)
         )
 
+        self._set_world_object(world_object)
+
         if z_position is not None:
             self.world_object.position.z = z_position
 

diff --git a/fastplotlib/graphics/line_collection.py b/fastplotlib/graphics/line_collection.py
@@ -157,7 +157,7 @@ def __init__(
                         "or must be a str of tuple/list with the same length as the data"
                     )
 
-        self._world_object = pygfx.Group()
+        self._set_world_object(pygfx.Group())
 
         for i, d in enumerate(data):
             if isinstance(z_position, list):
@@ -343,6 +343,6 @@ def __init__(
         )
 
         axis_zero = 0
-        for i, line in enumerate(self._graphics):
+        for i, line in enumerate(self.graphics):
             getattr(line.position, f"set_{separation_axis}")(axis_zero)
             axis_zero = axis_zero + line.data()[:, axes[separation_axis]].max() + separation
diff --git a/fastplotlib/graphics/line_slider.py b/fastplotlib/graphics/line_slider.py
@@ -74,7 +74,7 @@ def __init__(
         else:
             material = pygfx.LineMaterial
 
-        colors_inner = np.repeat([Color("w")], 2, axis=0).astype(np.float32)
+        colors_inner = np.repeat([Color(color)], 2, axis=0).astype(np.float32)
         colors_outer = np.repeat([Color([1., 1., 1., 0.25])], 2, axis=0).astype(np.float32)
 
         line_inner = pygfx.Line(
@@ -88,17 +88,19 @@ def __init__(
             material=material(thickness=thickness + 4, vertex_colors=True)
         )
 
-        self._world_object = pygfx.Group()
+        world_object = pygfx.Group()
 
-        self._world_object.add(line_outer)
-        self._world_object.add(line_inner)
+        world_object.add(line_outer)
+        world_object.add(line_inner)
+
+        self._set_world_object(world_object)
 
         self.position.x = x_pos
 
         self.slider = slider
         self.slider.observe(self.set_position, "value")
 
-        self.name = name
+        super().__init__(name=name)
 
     def set_position(self, change):
         self.position.x = change["new"]

diff --git a/fastplotlib/graphics/scatter.py b/fastplotlib/graphics/scatter.py
@@ -72,9 +72,11 @@ def __init__(
 
         super(ScatterGraphic, self).__init__(*args, **kwargs)
 
-        self._world_object: pygfx.Points = pygfx.Points(
+        world_object = pygfx.Points(
             pygfx.Geometry(positions=self.data(), sizes=sizes, colors=self.colors()),
             material=pygfx.PointsMaterial(vertex_colors=True, vertex_sizes=True)
         )
 
+        self._set_world_object(world_object)
+
         self.world_object.position.z = z_position
diff --git a/fastplotlib/graphics/text.py b/fastplotlib/graphics/text.py
@@ -38,11 +38,13 @@ def __init__(
         """
         super(TextGraphic, self).__init__(name=name)
 
-        self._world_object = pygfx.Text(
+        world_object = pygfx.Text(
             pygfx.TextGeometry(text=text, font_size=size, screen_space=False),
             pygfx.TextMaterial(color=face_color, outline_color=outline_color, outline_thickness=outline_thickness)
         )
 
+        self._set_world_object(world_object)
+
         self.world_object.position.set(*position)
 
         self.name = None