perf: defer Matplotlib import until something actually plots

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