diff --git a/marimo/_plugins/stateless/mpl/_mpl.py b/marimo/_plugins/stateless/mpl/_mpl.py index 3772c2bc8b2..5f7e9424808 100644 --- a/marimo/_plugins/stateless/mpl/_mpl.py +++ b/marimo/_plugins/stateless/mpl/_mpl.py @@ -52,9 +52,11 @@ def add(self, manager: FigureManagerWebAgg) -> None: self.figure_managers[str(manager.num)] = manager def get(self, figure_id: str) -> FigureManagerWebAgg: - if figure_id not in self.figure_managers: + try: + # Avoid double lookup by using try/except for performance + return self.figure_managers[figure_id] + except KeyError: raise RuntimeError(f"Figure {figure_id} not found.") # noqa: E501 - return self.figure_managers[str(figure_id)] def remove(self, manager: FigureManagerWebAgg) -> None: try: