Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Use array_api compliant dtype #8933

Draft
wants to merge 52 commits into
base: main
Choose a base branch
from

Conversation

Illviljan
Copy link
Contributor

@Illviljan Illviljan commented Apr 12, 2024

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

Notes:

  • For duckarrays, use _dtype[_generic]
  • For actual np.ndarrays, use np.dtype[np.generic]
  • np.dtype is too specific in general and it's probably not needed in most of their array_api functions.

@TomNicholas TomNicholas added the array API standard Support for the Python array API standard label Apr 12, 2024
@Illviljan
Copy link
Contributor Author

I don't understand why mypy doesn't narrow down the data typing here, I'm expecting the ExplicitlyIndexed to be gone downstream the isinstance-check:

@overload
def to_duck_array(
    data: ExplicitlyIndexed, **kwargs: dict[str, Any]
) -> duckarray[Any, _dtype[_generic]]: ...
@overload
def to_duck_array(
    data: chunkedduckarray[_ShapeType, _DType], **kwargs: dict[str, Any]
) -> chunkedduckarray[_ShapeType, _DType]: ...
@overload
def to_duck_array(
    data: duckarray[_ShapeType, _DType], **kwargs: dict[str, Any]
) -> duckarray[_ShapeType, _DType]: ...
@overload
def to_duck_array(
    data: np.typing.ArrayLike, **kwargs: dict[str, Any]
) -> duckarray[Any, _dtype[_generic]]: ...
def to_duck_array(
    data: (
        ExplicitlyIndexed
        | chunkedduckarray[_ShapeType, _DType]
        | duckarray[_ShapeType, _DType]
        | np.typing.ArrayLike
    ),
    **kwargs: dict[str, Any],
) -> (
    duckarray[_ShapeType, _DType]
    | duckarray[Any, _dtype[_generic]]
    | np.ndarray[Any, np.dtype[np.generic]]
):
    from xarray.core.indexing import ExplicitlyIndexed
    from xarray.namedarray.parallelcompat import get_chunked_array_type

    if isinstance(data, ExplicitlyIndexed):
        reveal_type(data)
        return data.get_duck_array()  # type: ignore[no-untyped-call, no-any-return]

    reveal_type(data)  # ExplicitlyIndexed should be gone?
    if is_chunked_array(data):
        chunkmanager = get_chunked_array_type(data)
        loaded_data, *_ = chunkmanager.compute(data, **kwargs)  # type: ignore[var-annotated]
        return loaded_data

    if is_duck_array(data):
        return data
    else:
        return np.asarray(data)

But mypy doesn't understand that for some reason:

xarray\namedarray\pycompat.py: note: In function "to_duck_array":
xarray\namedarray\pycompat.py:171: note: Revealed type is "xarray.core.indexing.ExplicitlyIndexed"
xarray\namedarray\pycompat.py:174: note: Revealed type is "Union[xarray.core.indexing.ExplicitlyIndexed, xarray.namedarray._typing._arrayfunction[_ShapeType`-1, _DType`-2], xarray.namedarray._typing._arrayapi[_ShapeType`-1, _DType`-2], _collections_abc.Buffer, numpy._typing._array_like._SupportsArray[numpy.dtype[Any]], numpy._typing._nested_sequence._NestedSequence[numpy._typing._array_like._SupportsArray[numpy.dtype[Any]]], builtins.int, builtins.float, builtins.complex, builtins.str, numpy._typing._nested_sequence._NestedSequence[Union[builtins.bool, builtins.int, builtins.float, builtins.complex, builtins.str, builtins.bytes]]]"
xarray\namedarray\pycompat.py:175: error: Argument 1 to "is_chunked_array" has incompatible type "ExplicitlyIndexed | _arrayfunction[_ShapeType, _DType] | _arrayapi[_ShapeType, _DType] | Buffer | _SupportsArray[dtype[Any]] | _NestedSequence[_SupportsArray[dtype[Any]]] | int | float | complex | str | _NestedSequence[bool | int | float | complex | str | bytes]"; expected "_arrayfunction[Any, Any] | _arrayapi[Any, Any]"  [arg-type]
Found 1 error in 1 file (checked 1 source file)

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
array API standard Support for the Python array API standard
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants