Is your feature request related to a problem?

Now that python 3.11 has been released, we should start to support it officially.

Describe the solution you'd like

I guess the first step would be to replace python 3.10 as a maximum version in the tests and see what crashes (and get lucky).

Describe alternatives you've considered

No response

Additional context

No response

Is your feature request related to a problem?

Currently, part of the static type of a DataArray or Dataset is a Mapping[Hashable, DataArray]. I'm quite sure that 99% of the users will actually use str key values (aka. variable names), while some exotic people (me included) want to use e.g. Enums for their keys. Currently, we allow to use anything as keys as long as it is hashable, but once the DataArray/set is created, the type information of the keys is lost.

Consider e.g. ```python

for name, da in Dataset({"a": ("t", np.arange(5))}).items(): reveal_type(name) # hashable reveal_type(da.dims) # tuple[hashable, ...] `` Woudn't that be nice if this would actually returnstr`, so you don't have to cast it or assert it everytime?

This could be solved by making these classes generic.

Another related issue is the underlying data. This could be introduced as a Generic type as well. Probably, this should reach some common ground on all wrapping array libs that are out there. Every one should use a Generic Array class that keeps track of the type of the wrapped array, e.g. dask.array.core.Array[np.ndarray]. In return, we could do DataArray[np.ndarray] or then DataArray[dask.array.core.Array[nd.ndarray]].

Describe the solution you'd like

The implementation would be something along the lines of:

```python KeyT = TypeVar("KeyT", bound=Hashable) DataT = TypeVar("DataT", bound=<some protocol?>)

class DataArray(Generic[KeyT, DataT]):

_coords: dict[KeyT, Variable[DataT]]
_indexes: dict[KeyT, Index[DataT]]
_name: KeyT | None
_variable: Variable[DataT]

def __init__(
    self,
    data: DataT = dtypes.NA,
    coords: Sequence[Sequence[DataT] | pd.Index | DataArray[KeyT]]
    | Mapping[KeyT, DataT]
    | None = None,
    dims: str | Sequence[KeyT] | None = None,
    name: KeyT | None = None,
    attrs: Mapping[KeyT, Any] | None = None,
    # internal parameters
    indexes: Mapping[KeyT, Index] | None = None,
    fastpath: bool = False,
) -> None:
...

```

Now you could create a "classical" DataArray: ```python da = DataArray(np.arange(10), {"t": np.arange(10)}, dims=["t"])

will be of type

DataArray[str, np.ndarray]

while you could also create something more fancypython da2 = DataArray(dask.array.array([1, 2, 3]), {}, dims=[("tup1", "tup2),])

will be of type

DataArray[tuple[str, str], dask.array.core.Array]

``` Any whenever you access the dimensions / coord names / underlying data you will get the correct type.

For now I only see three mayor problems: 1) non-array types (like lists or anything iterable) will get cast to a np.ndarray and I have no idea how to tell the type checker that DataArray([1, 2, 3], {}, "a") should be DataArray[str, np.ndarray] and not DataArray[str, list[int]]. Depending on the Protocol in the bound TypeVar this might even fail static type analysis or require tons of special casing and overloads. 2) How does the type checker extract the dimension type for Datasets? This is quite convoluted and I am not sure this can be typed correctly... 3) The parallel compute workflows are quite dynamic and I am not sure if static type checking can keep track of the underlying datatype... What does DataArray([1, 2, 3], dims="a").chunk({"a": 2}) return? Is it DataArray[str, dask.array.core.Array]? But what about other chunking frameworks?

Describe alternatives you've considered

One could even extend this and add more Generic types.

Different types for dimensions and variable names would be a first (and probably quite a nice) feature addition.

One could even go so far and type the keys and values of variables and coords (for Datasets) differently. This came up e.g. in https://github.com/pydata/xarray/issues/3967 However, this would create a ridiculous amount of Generic types and is probably more confusing than helpful.

Additional context

Probably this feature should be done in consecutive PRs that each implement one Generic each, otherwise this will be a giant task!

• [x] Closes #7340
• [x] Tests added
• [x] User visible changes (including notable bug fixes) are documented in whats-new.rst
• [x] ~New functions/methods are listed in api.rst~

• [x] Related to #7270

• [x] Closes #7275
• [x] Tests added
• [x] User visible changes (including notable bug fixes) are documented in whats-new.rst
• [x] ~New functions/methods are listed in api.rst~

• [x] Closes #7111
• [x] Tests added
• [x] User visible changes (including notable bug fixes) are documented in whats-new.rst

xarray.testing.assert_identical and probably more do not work yet.

While adding type hints to test_backends I noticed that the open_dataset method and the abstract BackendEntrypoint were missing stream types as inputs.

I was playing around a bit with a mypy plugin and this was the best I could come up with. Unfortunately the mypy docu about the plugins is not very detailed...

This plugin makes mypy recognize the user defined accessors.

There is a quite severe bug in there (due to my lack of understanding of mypy internals probably) which makes it work only on the first run but when you change a line in your code and run mypy again it will crash... (you can delete the cache to make it work one more time again :)

Any chance that a mypy expert can figure this out? haha

This PR removes several warnings from the tests and improves their typing on the way.

This PR adds typing support for groupby, coarsen, rolling, weighted and resample.

There are several open points:

1. Coarsen is missing type annotations for reductions like max, they get added dynamically.
2. The Groupby group-key type is quite wide. Does anyone have any idea on how to type it correctly? For now it is still Any.
3. Several function signatures were inconsistent between the DataArray and Dataset versions (looking at you: map). I took the freedom to align them (required for mypy), hopefully this does not break too much.
4. I was moving the generation functions from DataWithCoords to DataArray and Dataset, which adds some copy-paste of code (I tried to keep it minimal) but was unavoidable for typing support. (Adds the bonus that the corresponding modules are now only imported when required).

• [x] Closes #6707
• [x] ~~Tests added~~
• [x] User visible changes (including notable bug fixes) are documented in whats-new.rst

I don't really know how to test this... Is it ok to leave it untested?

(partial) typing for functions full_like, zeros_like, ones_like.

I could not figure out how to properly use TypeVars so many things are "hardcoded" with overloads.

I have added a DTypeLikeSave to npcompat, not sure that this file is supposed to be edited.

Problem1: TypeVar["T", Dataset, DataArray, Variable] can only be one of these three, but never with Union[Dataset, DataArray] which is used in several other places in xarray. Problem2: The official mypy support says: use TypeVar["T", bound=Union[Dataset, DataArray, Variable] but the the isinstance(obj, Dataset) could not be correctly resolved (is that a mypy issue?).

So if anyone can get it to work with TypeVars, feel free to change it. :)

Attempt to fix the typing issues in xr.polyval.

Some problems are still occuring and require a type: ignore. They seem more like mypy issues to me.

What happened?

When using DataArray/Dataset.isel with drop=True with a scalar DataArray as indexer (see example) resulting scalar coordinates do not get dropped. When using an integer the behavior is as expected.

What did you expect to happen?

I expect that using a scalar DataArray behaves the same as an integer.

Minimal Complete Verifiable Example

```Python import xarray as xr

da = xr.DataArray([1, 2, 3], dims="x", coord={"k": ("x", [0, 1, 2])})

<xarray.DataArray (x: 3)>

array([1, 2, 3])

Coordinates:

k (x) int32 0 1 2

da.isel({"x": 1}, drop=True)

works

<xarray.DataArray ()>

array(2)

da.isel({"x": xr.DataArray(1)}, drop=True)

does not drop "k" coordinate

<xarray.DataArray ()>

array(2)

Coordinates:

k int32 1

```

Relevant log output

No response

Anything else we need to know?

No response

Environment

What happened?

I have some polynomial coefficients and want to evaluate them at some values using xr.polyval.

As described in the docstring/docu I created a 1D coordinate DataArray and pass it to xr.polyval but it raises a KeyError (see example).

What did you expect to happen?

I expected that the polynomial would be evaluated at the given points.

Minimal Complete Verifiable Example

```Python import xarray as xr

coeffs = xr.DataArray([1, 2, 3], dims="degree")

With a "handmade" coordinate it fails:

coord = xr.DataArray([0, 1, 2], dims="x")

xr.polyval(coord, coeffs)

raises:

Traceback (most recent call last):

File "<stdin>", line 1, in <module>

File "xarray/core/computation.py", line 1847, in polyval

x = get_clean_interp_index(coord, coord.name, strict=False)

File "xarray/core/missing.py", line 252, in get_clean_interp_index

index = arr.get_index(dim)

File "xarray/core/common.py", line 404, in get_index

raise KeyError(key)

KeyError: None

If one adds a name to the coord that is called like the dimension:

coord2 = xr.DataArray([0, 1, 2], dims="x", name="x")

xr.polyval(coord2, coeffs)

works

```

Relevant log output

No response

Anything else we need to know?

I assume that the "standard" workflow is to obtain the coord argument from an existing DataArrays coordinate, where the name would be correctly set already. However, that is not clear from the description, and also prevents my "manual" workflow.

It could be that the problem will be solved by replacing the coord DataArray argument by an explicit Index in the future.

Environment