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