From 97ef592294ed5c9606b077700eb8ea9b52c29f27 Mon Sep 17 00:00:00 2001 From: Caitlin Lewis Date: Tue, 13 Dec 2022 12:09:16 -0500 Subject: [PATCH] storing graphics in dict --- fastplotlib/layouts/_base.py | 32 +++++++++++++++++++++----------- fastplotlib/layouts/_subplot.py | 5 +++++ 2 files changed, 26 insertions(+), 11 deletions(-) diff --git a/fastplotlib/layouts/_base.py b/fastplotlib/layouts/_base.py index 68221ad11..ee0601a70 100644 --- 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 index bb847c4ec..add79b76d 100644 --- 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)