Update signature for _arrayfunction.__array__

[Illviljan]

I believe numpy updated the signature of __array__ in numpy2.
It currently looks like this now: https://github.com/numpy/numpy/blob/de1ca2676e0f1db5a1b95e1d05e86ca87321cac2/numpy/__init__.pyi#L1471-L1478

xref #9231, #9252

[Illviljan]

Attempting to fix these:

xarray/tests/test_namedarray.py: note: In function "check_duck_array_typevar":
xarray/tests/test_namedarray.py:82: note: Revealed type is "Union[xarray.namedarray._typing._arrayfunction[Any, _DType`-1], xarray.namedarray._typing._arrayapi[Any, _DType`-1]]"
xarray/tests/test_namedarray.py: note: In class "TestNamedArray":
xarray/tests/test_namedarray.py:185: error: Argument 2 to "NamedArray" has incompatible type "ndarray[Any, dtype[Any]]"; expected "_arrayfunction[Any, Never] | _arrayapi[Any, Never]"  [arg-type]
xarray/tests/test_namedarray.py: note: In member "test_duck_array_class" of class "TestNamedArray":
xarray/tests/test_namedarray.py:346: error: Argument 1 to "check_duck_array_typevar" has incompatible type "ndarray[Any, dtype[signedinteger[_64Bit]]]"; expected "_arrayfunction[Any, Never] | _arrayapi[Any, Never]"  [arg-type]
xarray/tests/test_namedarray.py: note: In member "test_warn_on_repeated_dimension_names" of class "TestNamedArray":
xarray/tests/test_namedarray.py:565: error: Argument 2 to "NamedArray" has incompatible type "ndarray[Any, dtype[signedinteger[Any]]]"; expected "_arrayfunction[Any, Never] | _arrayapi[Any, Never]"  [arg-type]

[Illviljan]

I'm struggling to understand what mypy complains about here.

• mypy seems unable to narrow to the correct dtype for some reason. What are typical reasons to get Never ?
• mypy passes if the typing is explictly defined, fails if it's only implicit. What are typical reasons to require annotating variables ?
• pyright seems to handle this example fine though.

Here's a simplified example:

def test_mypy_typing_fail() -> None:
    from typing import Generic, TypeVar
    from xarray.namedarray._typing import _Dims

    _ST = TypeVar("_ST", bound=Any)
    _DT = TypeVar("_DT", bound=Any)

    class TestArray(Generic[_ST, _DT]):
        __slots__ = ("_data", "_dims", "_attrs")

        _data: duckarray[_ST, _DT]
        _dims: _Dims
        _attrs: dict[Any, Any] | None

        def __init__(
            self,
            dims: _Dims,
            data: duckarray[_ST, _DT],
            attrs: dict[Any, Any] = {},
        ):
            self._data = data
            self._dims = dims
            self._attrs = attrs

    data: np.ndarray[Any, np.dtype[np.int64]] = np.array([1, 2, 3], dtype=np.int64)
    # ta: TestArray[Any, np.dtype[np.int64]] # mypy fails if commented out, why?
    ta = TestArray(("time",), data)

    # mypy:
    # error: Need type annotation for "ta"  [var-annotated]
    # error: Argument 2 to "TestArray" has incompatible type "ndarray[Any, dtype[signedinteger[_64Bit]]]"; expected "_arrayfunction[Never, Never] | _arrayapi[Never, Never]"  [arg-type]

    # pyright - pass:
    # reveal_type(ta)  # information: Type of "ta" is "TestArray[_ShapeType_co@_arrayapi, dtype[signedinteger[_64Bit]]]"

[headtr1ck]

Thats weird... I cannot reproduce a minimal example that does not use duckarray...

Also reveal_type(duckarray)

Revealed type is "typing._SpecialForm"

is not very helpful...

[Illviljan]

Also reveal_type(duckarray)

Revealed type is "typing._SpecialForm"

is not very helpful...

Interesting, maybe due to the Union? This simpler example does the same at least:

int_and_float = Union[int, float]
reveal_type(int_and_float)  #  note: Revealed type is "typing._SpecialForm"

Compare pyright:

int_and_float = Union[int, float]
reveal_type(int_and_float)  #  information: Type of "int_and_float" is "type[int] | type[float]"

[Illviljan]

An interesting note. The duckarray protocol has to use all the typevars to be able to properly inherit values without explicit definition. Since numpy doesn't do that yet with _ShapeType I haven't pushed it either.

This example works but commenting out shape from the Protocol will start forcing explicit definitions:

from typing import Any, Generic, Protocol, TypeVar


_ST = TypeVar("_ST", bound=Any)
_ST_co = TypeVar("_ST_co", bound=Any, covariant=True)

_DT = TypeVar("_DT", bound=Any)
_DT_co = TypeVar("_DT_co", bound=Any, covariant=True)


class ArrayProtocol(Protocol[_ST_co, _DT_co]):
    @property
    def dtype(self) -> _DT_co: ...

    @property
    def shape(self) -> _ST_co: ...  # Commenting out shape will start to require explicit definition


class Array(Generic[_ST, _DT]):
    def __init__(self, shape: _ST, dtype: _DT) -> None:
        self._shape = shape
        self._dtype = dtype

    @property
    def shape(self) -> _ST:
        return self._shape

    @property
    def dtype(self) -> _DT:
        return self._dtype

    def count(self, value: Any, /) -> int:
        return 2


def test_pass_typevars(a: ArrayProtocol[_ST, _DT]) -> ArrayProtocol[_ST, _DT]:
    return a


arr = Array(shape=(3,), dtype=int)
reveal_type(arr)  # information: Type of "arr" is "Array[tuple[Literal[3]], type[int]]"

array = test_pass_typevars(a=arr)
reveal_type(array)  #  information: Type of "array" is "ArrayProtocol[tuple[Literal[3]], type[int]]"

[headtr1ck]

This example works but commenting out shape from the Protocol will start forcing explicit definitions:

That is expected, because the shape TypeVar is unbound then.

Maybe you could add a definition/overload of __new__ that always returns it with shape=Any?
Not sure how numpy handles this problem currently...

[Illviljan]

Wow, what a journey. Turns out it was _arrayfunction.__array__ that didn't have the correct signature, that caused the typevars to somehow stay unbounded (Never/Any).

Debug method was to turn off methods in _arrayfunction until the dtype typevars worked again. The clue was I managed to get it to work as intended by replacing with very basic Protocols.

[headtr1ck]

Thanks for figuring this out!

Some verbosity of mypy on why it decides stuff would be useful ...