Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
26 changes: 18 additions & 8 deletions spatialmath/base/graphics.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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]
Expand All @@ -1760,6 +1767,11 @@ def expand_dims(dim: ArrayLike = None, nd: int = 2) -> NDArray:
* [A,B] -> [A, B, A, B, A, B]
* [A,B,C,D,E,F] -> [A, B, C, D, E, F]
"""
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:
Expand All @@ -1780,8 +1792,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 isnotebook() -> bool:
"""
Expand Down
26 changes: 26 additions & 0 deletions tests/base/test_graphics.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,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)

@pytest.mark.skipif(
sys.platform.startswith("darwin") and sys.version_info < (3, 11),
reason="tkinter bug with mac",
Expand Down
Loading