pydata / pydata/xarray

Use Generic Types instead of Hashable or Any

Open
#8,199 2 comments 6 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

enhancement topic-typing
Dominant language
Python
Stars
4.2k
Forks
1.4k
Avg merge
2d 15h
Merged PRs (30d)
14

Description

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.


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 return str, 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:

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:

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 fancy

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!

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start by reading the existing type annotations for DataArray, Dataset, Variable, and Index, then trace how dimensions, coordinates, and underlying data are exposed. Define a smaller generic milestone first, and verify that its inferred types remain correct for NumPy arrays, Dask arrays, mappings, and chunked DataArray instances.

Written by the indexing model from the issue text.

Assessment

Tech stack
numpy, pandas, python
Domain
data
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Stale
Clarity
Needs clarification
Newbie friendliness
25/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.