Skip to content

perf: defer Matplotlib import until something actually plots#198

Open
petercorke wants to merge 1 commit into
rai-opensource:masterfrom
petercorke:perf/lazy-matplotlib-import
Open

perf: defer Matplotlib import until something actually plots#198
petercorke wants to merge 1 commit into
rai-opensource:masterfrom
petercorke:perf/lazy-matplotlib-import

Conversation

@petercorke

Copy link
Copy Markdown
Collaborator

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.

`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.
@codecov-commenter

Copy link
Copy Markdown

⚠️ Please install the 'codecov app svg image' to ensure uploads and comments are reliably processed by Codecov.

Codecov Report

❌ Patch coverage is 87.50000% with 4 lines in your changes missing coverage. Please review.

Files with missing lines Patch % Lines
spatialmath/base/__init__.py 92.30% 1 Missing ⚠️
spatialmath/base/transforms3d.py 83.33% 1 Missing ⚠️
spatialmath/geom2d.py 66.66% 1 Missing ⚠️
spatialmath/spline.py 66.66% 1 Missing ⚠️

📢 Thoughts on this report? Let us know!

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants