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.

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
@@ -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
 

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)

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
 
 # ======================================================================== #
@@ -1232,6 +1235,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):