Skip to content

Commit d8b7644

Browse files
authored
Fix typing for test_plot.py (pydata#9234)
* Fix typing for test_plot.py * Update test_plot.py * make sure we actually get ndarrays here, I get it locally at least * Add a minimal test and ignore in real test * Update test_plot.py * Update test_plot.py
1 parent 0eac740 commit d8b7644

File tree

1 file changed

+21
-16
lines changed

1 file changed

+21
-16
lines changed

xarray/tests/test_plot.py

Lines changed: 21 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -158,9 +158,10 @@ def setup(self) -> Generator:
158158
plt.close("all")
159159

160160
def pass_in_axis(self, plotmethod, subplot_kw=None) -> None:
161-
fig, axs = plt.subplots(ncols=2, subplot_kw=subplot_kw)
162-
plotmethod(ax=axs[0])
163-
assert axs[0].has_data()
161+
fig, axs = plt.subplots(ncols=2, subplot_kw=subplot_kw, squeeze=False)
162+
ax = axs[0, 0]
163+
plotmethod(ax=ax)
164+
assert ax.has_data()
164165

165166
@pytest.mark.slow
166167
def imshow_called(self, plotmethod) -> bool:
@@ -240,9 +241,9 @@ def test_1d_x_y_kw(self) -> None:
240241

241242
xy: list[list[None | str]] = [[None, None], [None, "z"], ["z", None]]
242243

243-
f, ax = plt.subplots(3, 1)
244+
f, axs = plt.subplots(3, 1, squeeze=False)
244245
for aa, (x, y) in enumerate(xy):
245-
da.plot(x=x, y=y, ax=ax.flat[aa])
246+
da.plot(x=x, y=y, ax=axs.flat[aa])
246247

247248
with pytest.raises(ValueError, match=r"Cannot specify both"):
248249
da.plot(x="z", y="z")
@@ -1566,7 +1567,9 @@ def test_colorbar_kwargs(self) -> None:
15661567
assert "MyLabel" in alltxt
15671568
assert "testvar" not in alltxt
15681569
# change cbar ax
1569-
fig, (ax, cax) = plt.subplots(1, 2)
1570+
fig, axs = plt.subplots(1, 2, squeeze=False)
1571+
ax = axs[0, 0]
1572+
cax = axs[0, 1]
15701573
self.plotmethod(
15711574
ax=ax, cbar_ax=cax, add_colorbar=True, cbar_kwargs={"label": "MyBar"}
15721575
)
@@ -1576,7 +1579,9 @@ def test_colorbar_kwargs(self) -> None:
15761579
assert "MyBar" in alltxt
15771580
assert "testvar" not in alltxt
15781581
# note that there are two ways to achieve this
1579-
fig, (ax, cax) = plt.subplots(1, 2)
1582+
fig, axs = plt.subplots(1, 2, squeeze=False)
1583+
ax = axs[0, 0]
1584+
cax = axs[0, 1]
15801585
self.plotmethod(
15811586
ax=ax, add_colorbar=True, cbar_kwargs={"label": "MyBar", "cax": cax}
15821587
)
@@ -3371,16 +3376,16 @@ def test_plot1d_default_rcparams() -> None:
33713376
# see overlapping markers:
33723377
fig, ax = plt.subplots(1, 1)
33733378
ds.plot.scatter(x="A", y="B", marker="o", ax=ax)
3374-
np.testing.assert_allclose(
3375-
ax.collections[0].get_edgecolor(), mpl.colors.to_rgba_array("w")
3376-
)
3379+
actual: np.ndarray = mpl.colors.to_rgba_array("w")
3380+
expected: np.ndarray = ax.collections[0].get_edgecolor() # type: ignore[assignment] # mpl error?
3381+
np.testing.assert_allclose(actual, expected)
33773382

33783383
# Facetgrids should have the default value as well:
33793384
fg = ds.plot.scatter(x="A", y="B", col="x", marker="o")
33803385
ax = fg.axs.ravel()[0]
3381-
np.testing.assert_allclose(
3382-
ax.collections[0].get_edgecolor(), mpl.colors.to_rgba_array("w")
3383-
)
3386+
actual = mpl.colors.to_rgba_array("w")
3387+
expected = ax.collections[0].get_edgecolor() # type: ignore[assignment] # mpl error?
3388+
np.testing.assert_allclose(actual, expected)
33843389

33853390
# scatter should not emit any warnings when using unfilled markers:
33863391
with assert_no_warnings():
@@ -3390,9 +3395,9 @@ def test_plot1d_default_rcparams() -> None:
33903395
# Prioritize edgecolor argument over default plot1d values:
33913396
fig, ax = plt.subplots(1, 1)
33923397
ds.plot.scatter(x="A", y="B", marker="o", ax=ax, edgecolor="k")
3393-
np.testing.assert_allclose(
3394-
ax.collections[0].get_edgecolor(), mpl.colors.to_rgba_array("k")
3395-
)
3398+
actual = mpl.colors.to_rgba_array("k")
3399+
expected = ax.collections[0].get_edgecolor() # type: ignore[assignment] # mpl error?
3400+
np.testing.assert_allclose(actual, expected)
33963401

33973402

33983403
@requires_matplotlib

0 commit comments

Comments
 (0)