diff --git a/spatialmath/base/graphics.py b/spatialmath/base/graphics.py index 2baf94b7..97368801 100644 --- a/spatialmath/base/graphics.py +++ b/spatialmath/base/graphics.py @@ -1513,7 +1513,10 @@ def _axes_dimensions(ax: plt.Axes) -> int: # handle the case of Animate objects pretending to be Axes classname = ax.__class__.__name__ base_classes = ax.__class__.__bases__ - if classname in ("Axes3DSubplot", "Animate") or any(base_class.__name__ in ("Axes3DSubplot", "Animate") for base_class in base_classes): + if classname in ("Axes3DSubplot", "Animate") or any( + base_class.__name__ in ("Axes3DSubplot", "Animate") + for base_class in base_classes + ): return 3 elif classname in ("AxesSubplot", "Animate2"): return 2 @@ -1736,20 +1739,24 @@ def plotvol3( ax._plotvol = True return ax - def expand_dims(dim: ArrayLike = None, nd: int = 2) -> NDArray: + def expand_dims(dim: ArrayLike | None = None, nd: int = 2) -> NDArray | None: """ Expand compact axis dimensions - :param dim: dimensions, defaults to None + :param dim: dimensions, or None if unspecified, defaults to None :type dim: scalar, array_like(2), array_like(4), array_like(6), optional - :param nd: number of axes dimensions, defaults to 2 + :param nd: number of axes dimensions (2 or 3), defaults to 2 :type nd: int, optional :raises ValueError: bad arguments - :return: 2d or 3d dimensions vector - :rtype: ndarray(4) or ndarray(6) + :return: 2d or 3d dimensions vector, or None if unspecified + :rtype: ndarray(4), ndarray(6), or None Compute bounding dimensions for plots from shorthand notation. + ``None`` is returned unchanged so the caller can choose whether to + autoscale or obtain dimensions from another source. No default extent + is inferred. + If ``nd==2``, [xmin, xmax, ymin, ymax]: * A -> [-A, A, -A, A] * [A,B] -> [A, B, A, B] @@ -1766,6 +1773,11 @@ def expand_dims(dim: ArrayLike = None, nd: int = 2) -> NDArray: corresponding maximum, matplotlib will invert that axis when the limits are applied, which is a convenient way to flip an axis direction. """ + if nd not in (2, 3): + raise ValueError("nd is 2 or 3") + if dim is None: + return None + dim = smb.getvector(dim) if nd == 2: @@ -1788,8 +1800,6 @@ def expand_dims(dim: ArrayLike = None, nd: int = 2) -> NDArray: return dim else: raise ValueError("bad dimension specified") - else: - raise ValueError("nd is 2 or 3") def set_window_title(fig: plt.Figure, title: str) -> bool: """ diff --git a/tests/base/test_graphics.py b/tests/base/test_graphics.py index 17b88735..2d5223c4 100644 --- a/tests/base/test_graphics.py +++ b/tests/base/test_graphics.py @@ -46,6 +46,32 @@ class TestGraphics(unittest.TestCase): def teardown_method(self, method): plt.close("all") + def test_expand_dims_unspecified(self): + self.assertIsNone(expand_dims()) + for nd in (2, 3): + with self.subTest(nd=nd): + self.assertIsNone(expand_dims(nd=nd)) + self.assertIsNone(expand_dims(None, nd=nd)) + + def test_expand_dims_specified(self): + for nd in (2, 3): + with self.subTest(nd=nd): + # Zero is an explicit extent, not an unspecified one. + np.testing.assert_array_equal(expand_dims(0, nd), [0, 0] * nd) + np.testing.assert_array_equal(expand_dims(3, nd), [-3, 3] * nd) + np.testing.assert_array_equal(expand_dims([-2, 5], nd), [-2, 5] * nd) + limits = np.arange(2 * nd) + np.testing.assert_array_equal(expand_dims(limits, nd), limits) + with self.assertRaises(ValueError): + expand_dims([1, 2, 3], nd) + + def test_expand_dims_invalid_nd(self): + for nd in (1, 4): + for dim in (None, 1): + with self.subTest(nd=nd, dim=dim): + with self.assertRaisesRegex(ValueError, "nd is 2 or 3"): + expand_dims(dim, nd) + def test_expand_dims(self): np.testing.assert_array_equal(expand_dims(3), [-3, 3, -3, 3]) np.testing.assert_array_equal(expand_dims(-3), [-3, 3, -3, 3]) @@ -56,7 +82,9 @@ def test_expand_dims(self): np.testing.assert_array_equal(expand_dims(-3, nd=3), [-3, 3, -3, 3, -3, 3]) np.testing.assert_array_equal(expand_dims([1, 2], nd=3), [1, 2] * 3) - np.testing.assert_array_equal(expand_dims([1, 2, 3, 4, 5, 6], nd=3), [1, 2, 3, 4, 5, 6]) + np.testing.assert_array_equal( + expand_dims([1, 2, 3, 4, 5, 6], nd=3), [1, 2, 3, 4, 5, 6] + ) with self.assertRaises(ValueError): expand_dims([1, 2, 3])