Skip to content

Commit

Permalink
Test type annotations for Variable.get_axis_num (pydata#9832)
Browse files Browse the repository at this point in the history
This is a regression test for GH9822.
  • Loading branch information
bmerry authored Nov 27, 2024
1 parent f3a65d5 commit 0b97969
Showing 1 changed file with 6 additions and 1 deletion.
7 changes: 6 additions & 1 deletion xarray/tests/test_variable.py
Original file line number Diff line number Diff line change
Expand Up @@ -1604,14 +1604,19 @@ def test_squeeze(self):
with pytest.raises(ValueError, match=r"cannot select a dimension"):
v.squeeze("y")

def test_get_axis_num(self):
def test_get_axis_num(self) -> None:
v = Variable(["x", "y", "z"], np.random.randn(2, 3, 4))
assert v.get_axis_num("x") == 0
assert v.get_axis_num(["x"]) == (0,)
assert v.get_axis_num(["x", "y"]) == (0, 1)
assert v.get_axis_num(["z", "y", "x"]) == (2, 1, 0)
with pytest.raises(ValueError, match=r"not found in array dim"):
v.get_axis_num("foobar")
# Test the type annotations: mypy will complain if the inferred
# type is wrong
v.get_axis_num("x") + 0
v.get_axis_num(["x"]) + ()
v.get_axis_num(("x", "y")) + ()

def test_set_dims(self):
v = Variable(["x"], [0, 1])
Expand Down

0 comments on commit 0b97969

Please sign in to comment.