Fix typing for test_plot.py

[Illviljan]

Fixes mypy errors related to plotting seen in CI.

xref: #9231

[Illviljan]

These looks like an error in matplotlib:

xarray/tests/test_plot.py: note: In function "test_plot1d_default_rcparams":
xarray/tests/test_plot.py:3380: error: Incompatible types in assignment (expression has type "tuple[float, float, float] | tuple[float, float, float, float] | tuple[tuple[float, float, float] | str, float] | tuple[tuple[float, float, float, float], float] | Sequence[tuple[float, float, float] | str | str | tuple[float, float, float, float] | tuple[tuple[float, float, float] | str, float] | tuple[tuple[float, float, float, float], float]]", variable has type "ndarray[Any, Any]")  [assignment]
xarray/tests/test_plot.py:3387: error: Incompatible types in assignment (expression has type "tuple[float, float, float] | tuple[float, float, float, float] | tuple[tuple[float, float, float] | str, float] | tuple[tuple[float, float, float, float], float] | Sequence[tuple[float, float, float] | str | str | tuple[float, float, float, float] | tuple[tuple[float, float, float] | str, float] | tuple[tuple[float, float, float, float], float]]", variable has type "ndarray[Any, Any]")  [assignment]
xarray/tests/test_plot.py:3399: error: Incompatible types in assignment (expression has type "tuple[float, float, float] | tuple[float, float, float, float] | tuple[tuple[float, float, float] | str, float] | tuple[tuple[float, float, float, float], float] | Sequence[tuple[float, float, float] | str | str | tuple[float, float, float, float] | tuple[tuple[float, float, float] | str, float] | tuple[tuple[float, float, float, float], float]]", variable has type "ndarray[Any, Any]")  [assignment]

The test boils down to:

import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt
import xarray as xr

ds = xr.tutorial.scatter_example_dataset(seed=42)

fig, ax = plt.subplots(1, 1)
ds.plot.scatter(x="A", y="B", marker="o", ax=ax)
actual: np.ndarray = mpl.colors.to_rgba_array("w")
assert isinstance(actual, np.ndarray)
expected: np.ndarray = ax.collections[0].get_edgecolor() # mypy complains here.
assert isinstance(expected, np.ndarray)

np.testing.assert_allclose(actual, expected)

get_edgecolor returns an np.ndarray and the asserts proves that but mypy disagrees, ignoring these.

A test with matplotlib only:

def test_debug() -> None:
    import numpy as np
    import matplotlib as mpl
    import matplotlib.pyplot as plt


    fig, ax = plt.subplots(1, 1)
    ax.scatter(x=np.array([1, 2, 3]), y=np.array([6, 7, 8]), color="k")
    actual: np.ndarray = mpl.colors.to_rgba_array("k")
    assert isinstance(actual, np.ndarray)
    expected: np.ndarray = ax.collections[0].get_edgecolor() # error: Incompatible types in assignment (expression has type "tuple[float, float, float] | tuple[float, float, float, float] | tuple[tuple[float, float, float] | str, float] | tuple[tuple[float, float, float, float], float] | Sequence[tuple[float, float, float] | str | str | tuple[float, float, float, float] | tuple[tuple[float, float, float] | str, float] | tuple[tuple[float, float, float, float], float]]", variable has type "ndarray[Any, Any]")  [assignment]
    assert isinstance(expected, np.ndarray)
    np.testing.assert_allclose(actual, expected)

Which also fails

xarray/tests/test_plot.py: note: In function "test_debug":
xarray/tests/test_plot.py:3377: error: Incompatible types in assignment (expression has type "tuple[float, float, float] | tuple[float, float, float, float] | tuple[tuple[float, float, float] | str, float] | tuple[tuple[float, float, float, float], float] | Sequence[tuple[float, float, float] | str | str | tuple[float, float, float, float] | tuple[tuple[float, float, float] | str, float] | tuple[tuple[float, float, float, float], float]]", variable has type "ndarray[Any, Any]")  [assignment]

Cc @ksunden if this still interests you.

CI has installed this:

INSTALLED VERSIONS
------------------
commit: ac3f09fdcd9d406784303240b534e31fbfd411d7
python: 3.11.9 | packaged by conda-forge | (main, Apr 19 2024, 18:36:13) [GCC 12.3.0]
python-bits: 64
OS: Linux
OS-release: 6.5.0-1023-azure
machine: x86_64
processor: x86_64
byteorder: little
LC_ALL: None
LANG: C.UTF-8
LOCALE: ('en_US', 'UTF-8')
libhdf5: 1.14.3
libnetcdf: 4.9.2

xarray: 2024.6.1.dev51+gac3f09fd
pandas: 2.2.2
numpy: 2.0.0
scipy: 1.14.0
netCDF4: 1.7.1
pydap: None
h5netcdf: 1.3.0
h5py: 3.11.0
zarr: 2.18.2
cftime: 1.6.4
nc_time_axis: 1.4.1
iris: 3.9.0
bottleneck: 1.4.0
dask: 2024.7.0
distributed: 2024.7.0
matplotlib: 3.9.1
cartopy: 0.23.0
seaborn: 0.13.2
numbagg: 0.8.1
fsspec: 2024.6.1
cupy: None
pint: None
sparse: 0.15.4
flox: 0.9.8
numpy_groupies: 0.11.1
setuptools: 70.2.0
pip: 24.0
conda: 24.5.0
pytest: 8.2.2
mypy: None
IPython: None
sphinx: None

[headtr1ck]

Are we trying to support numpy 1 and 2 in typing? Or do we simply switch to 2 entirely?

[max-sixty]

Are we trying to support numpy 1 and 2 in typing? Or do we simply switch to 2 entirely?

Possibly others have decided already — I defer to that if so.

But if not — I would advocate for minimizing developer burden in typing, and so only supporting a single version of dependencies. To the extent we're moving to numpy=2, that seems preferable...

[headtr1ck]

But if not — I would advocate for minimizing developer burden in typing, and so only supporting a single version of dependencies. To the extent we're moving to numpy=2, that seems preferable...

I totally agree. For internal typing checks we can restrict to numpy2. But anything exposed to user side should be compatible, I guess.

But actually I am not sure how to achieve that because mypy does not understand package version checks, only python version checks.

[ksunden]

This is another case of "np.ndarray doesn't count as a "Sequence", but is used as one regularly"

[max-sixty]

But anything exposed to user side should be compatible, I guess.

Though do we expose much of the numpy type system ourselves?

To the extent we do — I think it might be really difficult to maintain compat — I might give this up quite easily...