From d8b76448e3d20556ed0107dab8f702cd7c9d70f6 Mon Sep 17 00:00:00 2001
From: Illviljan <14371165+Illviljan@users.noreply.github.com>
Date: Sat, 13 Jul 2024 22:35:29 +0200
Subject: [PATCH] Fix typing for test_plot.py (#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
---
 xarray/tests/test_plot.py | 37 +++++++++++++++++++++----------------
 1 file changed, 21 insertions(+), 16 deletions(-)

diff --git a/xarray/tests/test_plot.py b/xarray/tests/test_plot.py
index fa08e9975ab..578e6bcc18e 100644
--- a/xarray/tests/test_plot.py
+++ b/xarray/tests/test_plot.py
@@ -158,9 +158,10 @@ def setup(self) -> Generator:
         plt.close("all")
 
     def pass_in_axis(self, plotmethod, subplot_kw=None) -> None:
-        fig, axs = plt.subplots(ncols=2, subplot_kw=subplot_kw)
-        plotmethod(ax=axs[0])
-        assert axs[0].has_data()
+        fig, axs = plt.subplots(ncols=2, subplot_kw=subplot_kw, squeeze=False)
+        ax = axs[0, 0]
+        plotmethod(ax=ax)
+        assert ax.has_data()
 
     @pytest.mark.slow
     def imshow_called(self, plotmethod) -> bool:
@@ -240,9 +241,9 @@ def test_1d_x_y_kw(self) -> None:
 
         xy: list[list[None | str]] = [[None, None], [None, "z"], ["z", None]]
 
-        f, ax = plt.subplots(3, 1)
+        f, axs = plt.subplots(3, 1, squeeze=False)
         for aa, (x, y) in enumerate(xy):
-            da.plot(x=x, y=y, ax=ax.flat[aa])
+            da.plot(x=x, y=y, ax=axs.flat[aa])
 
         with pytest.raises(ValueError, match=r"Cannot specify both"):
             da.plot(x="z", y="z")
@@ -1566,7 +1567,9 @@ def test_colorbar_kwargs(self) -> None:
         assert "MyLabel" in alltxt
         assert "testvar" not in alltxt
         # change cbar ax
-        fig, (ax, cax) = plt.subplots(1, 2)
+        fig, axs = plt.subplots(1, 2, squeeze=False)
+        ax = axs[0, 0]
+        cax = axs[0, 1]
         self.plotmethod(
             ax=ax, cbar_ax=cax, add_colorbar=True, cbar_kwargs={"label": "MyBar"}
         )
@@ -1576,7 +1579,9 @@ def test_colorbar_kwargs(self) -> None:
         assert "MyBar" in alltxt
         assert "testvar" not in alltxt
         # note that there are two ways to achieve this
-        fig, (ax, cax) = plt.subplots(1, 2)
+        fig, axs = plt.subplots(1, 2, squeeze=False)
+        ax = axs[0, 0]
+        cax = axs[0, 1]
         self.plotmethod(
             ax=ax, add_colorbar=True, cbar_kwargs={"label": "MyBar", "cax": cax}
         )
@@ -3371,16 +3376,16 @@ def test_plot1d_default_rcparams() -> None:
         # see overlapping markers:
         fig, ax = plt.subplots(1, 1)
         ds.plot.scatter(x="A", y="B", marker="o", ax=ax)
-        np.testing.assert_allclose(
-            ax.collections[0].get_edgecolor(), mpl.colors.to_rgba_array("w")
-        )
+        actual: np.ndarray = mpl.colors.to_rgba_array("w")
+        expected: np.ndarray = ax.collections[0].get_edgecolor()  # type: ignore[assignment] # mpl error?
+        np.testing.assert_allclose(actual, expected)
 
         # Facetgrids should have the default value as well:
         fg = ds.plot.scatter(x="A", y="B", col="x", marker="o")
         ax = fg.axs.ravel()[0]
-        np.testing.assert_allclose(
-            ax.collections[0].get_edgecolor(), mpl.colors.to_rgba_array("w")
-        )
+        actual = mpl.colors.to_rgba_array("w")
+        expected = ax.collections[0].get_edgecolor()  # type: ignore[assignment] # mpl error?
+        np.testing.assert_allclose(actual, expected)
 
         # scatter should not emit any warnings when using unfilled markers:
         with assert_no_warnings():
@@ -3390,9 +3395,9 @@ def test_plot1d_default_rcparams() -> None:
         # Prioritize edgecolor argument over default plot1d values:
         fig, ax = plt.subplots(1, 1)
         ds.plot.scatter(x="A", y="B", marker="o", ax=ax, edgecolor="k")
-        np.testing.assert_allclose(
-            ax.collections[0].get_edgecolor(), mpl.colors.to_rgba_array("k")
-        )
+        actual = mpl.colors.to_rgba_array("k")
+        expected = ax.collections[0].get_edgecolor()  # type: ignore[assignment] # mpl error?
+        np.testing.assert_allclose(actual, expected)
 
 
 @requires_matplotlib