How to type DTypeLike |Mapping[Any, DTypeLike]

[headtr1ck]

I am trying to improve typing in xarray and stumbled upon the following problem:

numpys DTypeLike also accepts dict-like types for their field support (https://numpy.org/doc/stable/user/basics.rec.html).

Some parts of xarrays code require a Mapping from variables to dtypes and will create this mapping if a single dtype is provided (i.e. full_like). I.e.:

if not isinstance(dtype, Mapping):
    dtype_ = {k: dtype for k in other.data_vars.keys()}
else:
    dtype_ = dtype

If I now type dtype: DTypeLike | Mapping[Hashable, DTypeLike] mypy will complain since DtypeLike itself can also be a Mapping.

What is the best way to solve this issue? Can one somehow detect if the Mapping is Mapping[Hashable, DTypeLike] or numpys _DTypeDict?

Simple MVCE where mypy raises:

from __future__ import annotations
from numpy.typing import DTypeLike
from typing import Mapping


x: DTypeLike | Mapping[int, DTypeLike]
y: Mapping[int, DTypeLike]

if isinstance(x, Mapping):
    y = x
else:
    y = {1: x}

[max-sixty]

I think the issue there is that y is getting assigned the DTypeLike | Mapping[int, DTypeLike] despite only being Mapping[int, DTypeLike].

Does this help?

from __future__ import annotations

from typing import Any, Mapping, Union

from numpy.typing import DTypeLike

DTypeLikeMap = Union[Mapping[Any, DTypeLike], DTypeLike]

x: DTypeLikeMap
y: DTypeLikeMap

if isinstance(x, Mapping):
    y = x
else:
    # y = {1: x}
    y = {"a": x}

Weirdly it was getting confused by the 1, I think because of some issue with TypedDict, which is part of DTypeLike. I'm not sure whether that's going to cause more problems...

[headtr1ck]

I think it is getting confused since this
np.dtype({'names': ['col1', 'col2'], 'formats': ['i4', 'f4']}) is a valid dtype but also a Mapping.

So we might need to define our own DTypeLike that disallows field dtype definitions and only allow them in form of actual np.dtyes.

[max-sixty]

Good point. Is {'names': ['col1', 'col2'], 'formats': ['i4', 'f4']} (i.e. without the np.dtype) also a valid dtype? It's a shame that mypy can't unify those into the union of the types.

If we can make our own, that sounds like the best path, even though it'll be some work to recreate it. And then hopefully mypy would gain that ability to union the Mappings at some point and we could defer to numpy's.

[headtr1ck]

Yes, this dict is a valid DTypeLike.
That's why mypy complains in the assignment under the isinstance(x, Mapping)

[headtr1ck]

Actually i just realized that you want y to be of type Mapping[Any, DtypeLike] since later on it is used as a Mapping from variables in a Dataset to their dtypes.

[max-sixty]

Yes good point!

[headtr1ck]

I guess I will go with a custom DTypeLike definition that excludes Mappings.

This uses numpys internal definitions (many of them are still Any, but will change in future releases):

    from numpy.typing._dtype_like import _SupportsDType, _DTypeLikeNested, _ShapeLike, DType

    DTypeLikeSave = Union[
        DType[Any],
        # default data type (float64)
        None,
        # array-scalar types and generic types
        Type[Any],
        # character codes, type strings or comma-separated fields, e.g., 'float64'
        str,
        # (flexible_dtype, itemsize)
        Tuple[_DTypeLikeNested, int],
        # (fixed_dtype, shape)
        Tuple[_DTypeLikeNested, _ShapeLike],
        # (base_dtype, new_dtype)
        Tuple[_DTypeLikeNested, _DTypeLikeNested],
        # because numpy does the same?
        list[Any],
        # anything with a dtype attribute
        _SupportsDType[DType[Any]]
    ]

[max-sixty]

That would be great!