perf: defer Matplotlib import until something actually plots (#198)

## Problem

`import spatialmath` unconditionally pulls in all of Matplotlib, even
for
users who never plot anything (RAI wishlist item: "Don't depend on
matplotlib unless you're actually plotting").
`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 matplotlib turned out to be missing.

Measured on my 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's
the
cost this PR removes from `import spatialmath`. (Caveat: on this machine
matplotlib isn't the single biggest contributor to overall import time —
`scipy.interpolate` and `sympy`, when installed, are comparable — so
this
alone may not fully explain reports of ~9s import time on other
machines;
worth keeping an eye on.)

## Approach

Deliberately **not** touching matplotlib usage inside `graphics.py`/
`animate.py` themselves — dozens of `plt.Axes`-style annotations live in
there, and rewriting those would be a much bigger, riskier change than
this
needs to be. Instead, defer the **module import** of
`graphics`/`animate`
as a whole, plus a few other files with the same eager-import pattern:

- **`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.), not 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.
- **`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 — `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
(module-import time, before the deferred import ever runs) — annotations
become lazy strings instead, which type checkers still read fine.
Matches
the existing convention already used in 8 other files in this codebase.

One real (minor) bug this surfaced: `tests/test_geom3d.py` used
`plt.figure()` without importing `matplotlib.pyplot` itself — it only
worked because `from spatialmath.geom3d import *` used to leak `plt` in
as
a wildcard-imported name. Fixed by importing it directly.

## Testing

- `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 whose local-run behaviour
depends
  on a separate, already-open PR, unrelated to this change).
- Manual smoke test with a forced Agg backend: `trplot`, `trplot2`,
  `Polygon2.plot`/`.animate`/`.contains`, `Ellipse.plot`, `Line3.plot`,
  `Plane3.plot`, `BSplineSE3.visualize` all still produce real output.
- `black --check` clean at the pinned 23.10.0.

Author: petercorke

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,

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
@@ -17,21 +19,24 @@
 import sys
 import math
 import warnings
+import importlib.util
+from typing import TYPE_CHECKING
 import numpy as np
 import warnings
 
-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
@@ -1311,10 +1316,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],
@@ -1569,6 +1570,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
 

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
 import warnings
 
@@ -46,12 +50,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
 
 # ---------------------------------------------------------------------------------------#
@@ -2994,13 +3001,10 @@ def trprint(
     return s
 
 
-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:
 
@@ -3191,6 +3195,8 @@ def trplot(
         # animation
         # anaglyph
 
+        from spatialmath.base.graphics import plotvol3, axes_logic
+
         if dims is None:
             ax = axes_logic(ax, 3, projection)
         else:
@@ -3504,6 +3510,8 @@ def tranimate(T: Union[SO3Array, SE3Array], **kwargs) -> str:
 
         :seealso: `trplot`, `plotvol3`
         """
+        from spatialmath.base.animate import Animate
+
         # accept dim (matches Animate/plotvol2/plotvol3), keep dims as an
         # alias for backward compatibility with the previous docstring/API
         dim = kwargs.pop("dim", kwargs.pop("dims", None))

spatialmath/geom2d.py

@@ -9,15 +9,14 @@
 
 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
 import numpy as np
 
 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):

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
 
 # ======================================================================== #
@@ -1233,6 +1236,8 @@ def plot(
 
         :seealso: :meth:`intersect_volume`
         """
+        import matplotlib.pyplot as plt
+
         if ax is None:
             ax = plt.gca()
 

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,16 +8,18 @@
 """
 
 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
 
 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")

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):