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)