fastplotlib · clewis7 · Dec 13, 2022
diff --git a/fastplotlib/layouts/_base.py b/fastplotlib/layouts/_base.py
@@ -39,7 +39,9 @@ def __init__(

        self.renderer.add_event_handler(self.set_viewport_rect, "resize")

-        self._graphics: List[Graphic] = list()
+        self._graphics = {}
+        # \
+        #     List[Graphic] = list()

        self.name = name

@@ -101,13 +103,19 @@ def add_graphic(self, graphic, center: bool = True):
        if graphic.name is not None:  # skip for those that have no name
            graphic_names = list()

-            for g in self._graphics:
-                graphic_names.append(g.name)
+            for key in self._graphics.keys():
+                for g in self._graphics[key]:
+                    graphic_names.append(g.name)

            if graphic.name in graphic_names:
                raise ValueError(f"graphics must have unique names, current graphic names are:\n {graphic_names}")

-        self._graphics.append(graphic)
+        if graphic.__class__.__name__ in self._graphics.keys():
+            self._graphics[graphic.__class__.__name__].append(graphic)
+        else:
+            self._graphics[graphic.__class__.__name__] = list()
+            self._graphics[graphic.__class__.__name__].append(graphic)
+
        self.scene.add(graphic.world_object)

        if center:
@@ -155,14 +163,16 @@ def remove_graphic(self, graphic):
        self.scene.remove(graphic.world_object)

    def __getitem__(self, name: str):
-        for graphic in self._graphics:
-            if graphic.name == name:
-                return graphic
+        for key in self._graphics.keys():
+            for graphic in self._graphics[key]:
+                if graphic.name == name:
+                    return graphic

        graphic_names = list()
-        for g in self._graphics:
-            graphic_names.append(g.name)
-        raise IndexError(f"no graphic of given name, the current graphics are:\n {graphic_names}")
+        for key in self._graphics.keys():
+            for g in self._graphics[key]:
+                graphic_names.append(g.name)
+            raise IndexError(f"no graphic of given name, the current graphics are:\n {graphic_names}")

    def __str__(self):
        if self.name is None:
@@ -178,5 +188,5 @@ def __repr__(self):
        return f"{self}\n" \
               f"  parent: {self.parent}\n" \
               f"  Graphics:\n" \
-               f"\t{newline.join(graphic.__repr__() for graphic in self.get_graphics())}" \
+               f"\t{newline.join(graphic.__repr__() for key in self.get_graphics().keys() for graphic in self.get_graphics()[key])}" \
               f"\n"
diff --git a/fastplotlib/layouts/_subplot.py b/fastplotlib/layouts/_subplot.py
@@ -108,6 +108,11 @@ def add_graphic(self, graphic, center: bool = True):
        if isinstance(graphic, HeatmapGraphic):
            self.controller.scale.y = copysign(self.controller.scale.y, -1)

+    def remove_graphic(self, graphic):
+        super(Subplot, self).remove_graphic(graphic)
+
+        self.get_graphics()[graphic.__class__.__name__].remove(graphic)
+
    def set_axes_visibility(self, visible: bool):
        if visible:
            self.scene.add(self._axes)