From 8ee82e5d39efec77339ea72db57b143f24e713a6 Mon Sep 17 00:00:00 2001 From: Peter Corke Date: Mon, 27 Jul 2026 15:13:34 +1000 Subject: [PATCH] perf: defer Matplotlib import until something actually plots `import spatialmath` unconditionally pulled in all of Matplotlib, even for users who never plot anything: `spatialmath/base/__init__.py` did `from spatialmath.base.animate import *` / `from spatialmath.base.graphics import *` at package-import time, and those two modules do `import matplotlib.pyplot`. `geom2d.py`, `geom3d.py`, `spline.py`, and both `transforms2d.py`/`transforms3d.py` had their own copies of the same pattern, some via a `try: import matplotlib.pyplot ... except ImportError:` optional-dependency check that still imported eagerly, just without crashing if it failed. Measured on this machine: bare `import matplotlib` costs ~96ms, `matplotlib.path` (needed structurally by Polygon2's geometry, not just plotting) adds ~0ms on top of that, but `matplotlib.pyplot` specifically (backend resolution + figure/state setup) adds another ~160ms. That pyplot cost is what this change removes from `import spatialmath`. Approach: don't touch matplotlib usage inside graphics.py/animate.py at all (dozens of `plt.Axes`-style annotations in there - rewriting those would be the disproportionate version of this fix). Instead, defer the module import itself: - spatialmath/base/__init__.py: replace the two blanket `import *` lines with a PEP 562 module `__getattr__` that imports animate.py/graphics.py lazily on first access of one of their names (Animate, plot_box, tranimate, etc.) rather than unconditionally at package import time. - spatialmath/base/transforms2d.py, transforms3d.py: both had a module-level `try: import matplotlib.pyplot ... except ImportError: _matplotlib_exists = False` used purely to decide whether to define trplot/trplot2/tranimate/tranimate2 at all. Replaced with a cheap `importlib.util.find_spec("matplotlib")` check (no real import), and moved the small number of actual `plotvol2/3`, `axes_logic`, `Animate`/`Animate2` and `plt.show()` call sites into the specific functions that use them - they were already the only things those two files needed matplotlib for. - spatialmath/spline.py, geom3d.py: same pattern, moved `import matplotlib.pyplot as plt` from module top into the one or two methods that actually plot. - spatialmath/geom2d.py: nuance here - Polygon2.__init__ uses `matplotlib.path.Path` for real geometry (point containment etc.), and Polygon2.transformed() uses `matplotlib.transforms.Affine2D` for the same reason, so both stay at module level (cheap anyway, per the measurement above). Only `matplotlib.pyplot` itself, and the `plot_ellipse` import (only used by Ellipse.plot), moved. All of the above needed `from __future__ import annotations` added where not already present, so a signature like `ax: Optional[plt.Axes] = None` doesn't force `plt` to be a real bound name at function *definition* time (which happens at module-import time, before the deferred import ever runs) - annotations become lazy strings instead, which type checkers still read fine. Matches the existing convention in 8 other files in this codebase. One real (if minor) bug this surfaced: tests/test_geom3d.py used `plt.figure()` without importing matplotlib.pyplot itself - it was only working because `from spatialmath.geom3d import *` used to leak `plt` in as a wildcard-imported name. Fixed by importing it directly, which is what the test should have done regardless of this change. Verified: - `import spatialmath` no longer puts `matplotlib.pyplot` in sys.modules; `matplotlib` (the ~96ms base package, structurally needed by Polygon2's geometry) still does. - Wall-clock `import spatialmath`, same machine, same warm caches: ~670ms average on 3 runs before this change (upstream/master, via a throwaway worktree), ~447ms average after. - Full test suite green in both CI-like mode (`CI=true MPLBACKEND=Agg`: 338 passed, 4 skipped) and a local-run simulation (`MPLBACKEND=Agg`, `CI` unset): 332 passed (excluding two files that depend on a separate, already-open PR for their own local-run fixes, unrelated to this change). - Manual smoke test: trplot, trplot2, Polygon2.plot/.animate/.contains, Ellipse.plot, Line3.plot, Plane3.plot, BSplineSE3.visualize all still produce real output under a forced Agg backend. - `black --check` clean at the pinned 23.10.0. --- spatialmath/base/__init__.py | 52 ++++++++++++++++++++++++++++++-- spatialmath/base/transforms2d.py | 23 ++++++++------ spatialmath/base/transforms3d.py | 26 ++++++++++------ spatialmath/geom2d.py | 10 ++++-- spatialmath/geom3d.py | 7 ++++- spatialmath/spline.py | 10 ++++-- tests/test_geom3d.py | 1 + 7 files changed, 103 insertions(+), 26 deletions(-) diff --git a/spatialmath/base/__init__.py b/spatialmath/base/__init__.py index 9e9fbcbe..9c1799cc 100644 --- a/spatialmath/base/__init__.py +++ b/spatialmath/base/__init__.py @@ -9,10 +9,58 @@ from spatialmath.base.transformsNd import * # lgtm [py/polluting-import] from spatialmath.base.vectors import * # lgtm [py/polluting-import] from spatialmath.base.symbolic import * # lgtm [py/polluting-import] -from spatialmath.base.animate import * # lgtm [py/polluting-import] -from spatialmath.base.graphics import * # lgtm [py/polluting-import] from spatialmath.base.numeric import * # lgtm [py/polluting-import] +import importlib + +# `animate` and `graphics` both import Matplotlib, which is slow to import +# (mostly backend resolution in matplotlib.pyplot) and is only actually +# needed once something tries to plot. Load them lazily, on first access +# of one of their names below, instead of unconditionally at package +# import time. +_LAZY_SUBMODULES = ("animate", "graphics") + +_LAZY_ATTRS = { + "Animate": "animate", + "Animate2": "animate", + "plot_text": "graphics", + "plot_point": "graphics", + "plot_homline": "graphics", + "plot_box": "graphics", + "plot_arrow": "graphics", + "plot_polygon": "graphics", + "circle": "graphics", + "plot_circle": "graphics", + "ellipse": "graphics", + "plot_ellipse": "graphics", + "sphere": "graphics", + "plot_sphere": "graphics", + "ellipsoid": "graphics", + "plot_ellipsoid": "graphics", + "cylinder": "graphics", + "plot_cylinder": "graphics", + "plot_cone": "graphics", + "plot_cuboid": "graphics", + "axes_logic": "graphics", + "plotvol2": "graphics", + "plotvol3": "graphics", + "expand_dims": "graphics", + "isnotebook": "graphics", +} + + +def __getattr__(name): + if name in _LAZY_SUBMODULES: + return importlib.import_module(f"spatialmath.base.{name}") + modname = _LAZY_ATTRS.get(name) + if modname is None: + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + module = importlib.import_module(f"spatialmath.base.{modname}") + value = getattr(module, name) + globals()[name] = value # cache so future lookups skip __getattr__ + return value + + from spatialmath.base.argcheck import ( assertmatrix, ismatrix, diff --git a/spatialmath/base/transforms2d.py b/spatialmath/base/transforms2d.py index ac0696cd..a0655eef 100644 --- a/spatialmath/base/transforms2d.py +++ b/spatialmath/base/transforms2d.py @@ -1,3 +1,5 @@ +from __future__ import annotations + # Part of Spatial Math Toolbox for Python # Copyright (c) 2000 Peter Corke # MIT Licence, see details in top-level file: LICENCE @@ -16,20 +18,23 @@ import sys import math +import importlib.util +from typing import TYPE_CHECKING import numpy as np -try: - import matplotlib.pyplot as plt - - _matplotlib_exists = True -except ImportError: - _matplotlib_exists = False +# cheap existence check, doesn't actually import matplotlib: the real +# import happens lazily inside trplot2()/tranimate2() when a plot is +# actually made +_matplotlib_exists = importlib.util.find_spec("matplotlib") is not None import spatialmath.base as smb from spatialmath.base.types import * from spatialmath.base.transformsNd import rt2tr from spatialmath.base.vectors import unitvec +if TYPE_CHECKING: + from matplotlib.axes import Axes + _eps = np.finfo(np.float64).eps try: # pragma: no cover @@ -1227,10 +1232,6 @@ def _FindCorrespondences( if _matplotlib_exists: - import matplotlib.pyplot as plt - - # from mpl_toolkits.axisartist import Axes - from matplotlib.axes import Axes def trplot2( T: Union[SO2Array, SE2Array], @@ -1485,6 +1486,8 @@ def trplot2( if block is not None: # calling this at all, causes FuncAnimation to fail so when invoked from tranimate2 skip this bit + import matplotlib.pyplot as plt + plt.show(block=block) return ax diff --git a/spatialmath/base/transforms3d.py b/spatialmath/base/transforms3d.py index 04f9e1b8..a8d887d9 100644 --- a/spatialmath/base/transforms3d.py +++ b/spatialmath/base/transforms3d.py @@ -1,3 +1,5 @@ +from __future__ import annotations + # Part of Spatial Math Toolbox for Python # Copyright (c) 2000 Peter Corke # MIT Licence, see details in top-level file: LICENCE @@ -15,8 +17,10 @@ # pylint: disable=invalid-name import sys +import importlib.util from collections.abc import Iterable import math +from typing import TYPE_CHECKING import numpy as np from spatialmath.base.argcheck import getunit, getvector, isvector, isscalar, ismatrix @@ -44,12 +48,15 @@ Ab2M, ) from spatialmath.base.quaternions import r2q, q2r, qeye, qslerp, qunit -from spatialmath.base.graphics import plotvol3, axes_logic -from spatialmath.base.animate import Animate import spatialmath.base.symbolic as sym from spatialmath.base.types import * +if TYPE_CHECKING: + # for static type checkers only, both are only ever really imported + # lazily, inside trplot()/tranimate(), when a plot is actually made + from mpl_toolkits.mplot3d import Axes3D + _eps = np.finfo(np.float64).eps # ---------------------------------------------------------------------------------------# @@ -2910,13 +2917,10 @@ def _vec2s(fmt, v): return ", ".join([fmt.format(x) for x in v]) -try: - import matplotlib.pyplot as plt - from mpl_toolkits.mplot3d import Axes3D - - _matplotlib_exists = True -except ImportError: - _matplotlib_exists = False +# cheap existence check, doesn't actually import matplotlib: the real +# import happens lazily inside trplot()/tranimate() when a plot is +# actually made +_matplotlib_exists = importlib.util.find_spec("matplotlib") is not None if _matplotlib_exists: @@ -3107,6 +3111,8 @@ def trplot( # animation # anaglyph + from spatialmath.base.graphics import plotvol3, axes_logic + if dims is None: ax = axes_logic(ax, 3, projection) else: @@ -3419,6 +3425,8 @@ def tranimate(T: Union[SO3Array, SE3Array], **kwargs) -> str: :seealso: `trplot`, `plotvol3` """ + from spatialmath.base.animate import Animate + dim = kwargs.pop("dims", None) ax = kwargs.pop("ax", None) anim = Animate(dim=dim, ax=ax, **kwargs) diff --git a/spatialmath/geom2d.py b/spatialmath/geom2d.py index 55eccb2a..631ef331 100755 --- a/spatialmath/geom2d.py +++ b/spatialmath/geom2d.py @@ -9,7 +9,7 @@ from functools import reduce import warnings -import matplotlib.pyplot as plt +from typing import TYPE_CHECKING from matplotlib.path import Path from matplotlib.patches import PathPatch from matplotlib.transforms import Affine2D @@ -17,7 +17,6 @@ from spatialmath import SE2 import spatialmath.base as smb -from spatialmath.base import plot_ellipse from spatialmath.base.types import ( Points2, Optional, @@ -37,6 +36,9 @@ cast, ) +if TYPE_CHECKING: + import matplotlib.pyplot as plt + _eps = np.finfo(np.float64).eps @@ -450,6 +452,8 @@ def plot(self, ax: Optional[plt.Axes] = None, **kwargs) -> None: :seealso: :meth:`animate` :func:`matplotlib.PathPatch` """ + import matplotlib.pyplot as plt + self.patch = PathPatch(self.path, **kwargs) ax = smb.axes_logic(ax, 2) ax.add_patch(self.patch) @@ -1063,6 +1067,8 @@ def plot(self, **kwargs) -> None: :seealso: :func:`~spatialmath.base.graphics.plot_ellipse` """ + from spatialmath.base import plot_ellipse + return plot_ellipse(self._E, centre=self._centre, **kwargs) def contains(self, p): diff --git a/spatialmath/geom3d.py b/spatialmath/geom3d.py index 896192dc..beceae5b 100755 --- a/spatialmath/geom3d.py +++ b/spatialmath/geom3d.py @@ -6,12 +6,15 @@ import numpy as np import math from collections import namedtuple -import matplotlib.pyplot as plt +from typing import TYPE_CHECKING import spatialmath.base as base from spatialmath.base.types import * from spatialmath.baseposelist import BasePoseList import warnings +if TYPE_CHECKING: + import matplotlib.pyplot as plt + _eps = np.finfo(np.float64).eps # ======================================================================== # @@ -1232,6 +1235,8 @@ def plot( :seealso: :meth:`intersect_volume` """ + import matplotlib.pyplot as plt + if ax is None: ax = plt.gca() diff --git a/spatialmath/spline.py b/spatialmath/spline.py index c2da982e..5fd66b39 100644 --- a/spatialmath/spline.py +++ b/spatialmath/spline.py @@ -1,3 +1,5 @@ +from __future__ import annotations + # Copyright (c) 2024 Robotics and AI Institute LLC dba RAI Institute. # MIT Licence, see details in top-level file: LICENCE @@ -6,9 +8,8 @@ """ from abc import ABC, abstractmethod -from typing import List, Optional, Tuple +from typing import TYPE_CHECKING, List, Optional, Tuple -import matplotlib.pyplot as plt import numpy as np from scipy.interpolate import BSpline, CubicSpline from scipy.spatial.transform import Rotation, RotationSpline @@ -16,6 +17,9 @@ from spatialmath import SE3, SO3, Twist3 from spatialmath.base.transforms3d import tranimate +if TYPE_CHECKING: + import matplotlib.pyplot as plt + class SplineSE3(ABC): def __init__(self) -> None: @@ -39,6 +43,8 @@ def visualize( Args: sample_times: which times to sample the spline at and plot """ + import matplotlib.pyplot as plt + if ax is None: fig = plt.figure(figsize=(10, 10)) ax = fig.add_subplot(projection="3d") diff --git a/tests/test_geom3d.py b/tests/test_geom3d.py index 7a743dd5..81a50fdb 100755 --- a/tests/test_geom3d.py +++ b/tests/test_geom3d.py @@ -14,6 +14,7 @@ import spatialmath.base as base import pytest import sys +import matplotlib.pyplot as plt class Line3Test(unittest.TestCase):