Add an Array Protocol & improve static typing support

[nstarman]

Continuing from #584, I've done a quick refactor of the array (and dtype) objects to make them typing Protocols.
@rgommers, I agree that making Array generic wrt the dtype is good for a followup.

It's not perfect, but mypy doesn't complain a whole lot, so I think this is a good start. I'm happy to make refinements.

[honno]

Thanks for your work on this :)

The Array protocl stuff looks grand. And personally I like having a DType class just for consistency, even if practically type hinting an object has __eq__ isn't useful.

On namespace type hints (e.g. _CreatoinFuncs, ArrayAPINamespace): IMO I'm not sure on this yet, as they make contributing to the spec less ergonomic given we repeat the signatures, although they are pretty nifty to have. ^{Is there a world where these could be dynamically construct them from the stubs instead?} We might want to put these in a follow-up PR if just to ship the concretely non-controversial stuff first—interested in what others think.

[honno]

So I see methods/properties which don't return arrays don't get any type hint for their self arg—could these be hinted with Array?

[nstarman]

Static type checkers don't like this because the TypeVar is not properly bound — see https://peps.python.org/pep-0484/#scoping-rules-for-type-variables ("Unbound type variables should not appear in the bodies of generic functions, or in the class bodies apart from method definitions:").
In general, there's no need to type hint the self arg unless the return type depends on the type of self — e.g. see https://peps.python.org/pep-0484/#user-defined-generic-types.

[nstarman]

A nice thing about using Protocol is that self is understood by static type checkers to be the Protocol.

[asmeurer]

Why is this only done for certain methods?

[nstarman]

I did it for any methods I looked at over the course of the PR. This is by no means comprehensive and should be done: here or elsewhere.

[nstarman]

We might want to put these in a follow-up PR if just to ship the concretely non-controversial stuff first—interested in what others think.

I was thinking the same thing last evening about moving this to a followup PR.

On namespace type hints (e.g. _CreatoinFuncs, ArrayAPINamespace): IMO I'm not sure on this yet, as they make contributing to the spec less ergonomic given we repeat the signatures, although they are pretty nifty to have. Is there a world where these could be dynamically construct them from the stubs instead?

Maybe? I haven't experimented yet but perhaps mypy will be satisfied with:

class ArrayAPINamespace(Protocol):

    zeros_like = staticmethod(creation_funcs.zeros_like)

Also for a future PR: zeros_like could be promoted to a Protocol,

class Array(Protocol):
    def __array_namespace__(self, ...) -> ArrayAPINamespace[Self]: ...


class ArrayAPINamespace(Protocol[Array]):

    zeros_like = staticmethod(zeros_like)


class zeros_like(Protocol[Array]):
    def __call__(self, x: Array, /, *, dtype: Optional[dtype] = None, device: Optional[device] = None) -> Array: ...

The advantage of this is that functions can be checked to see if their signature is compatible with zeros_like.

[rgommers]

Thanks @nstarman, and apologies for the delay in reviewing (your PR came in right at the start of my 2 week holiday). This is looking quite good. I pushed a couple of small commits to resolve some Mypy and Sphinx complaints.

This is close to ready I'd say. One thing that isn't nice is that Self now shows up instead of array in the signatures of the html docs. I tried to fiddle with autodoc_type_aliases in src/_array_api_conf.py, but was not successful in mapping that back to array. Using "self" will be pretty confusing in docs for regular methods that return arrays. I think we'll have to figure that one out.

The DType protocol is fine with me - it won't add much to type-checking, but it should be fine to have it here for consistency.

The Array protocol looks good.

[rgommers]

@BvB93 if you have any feedback on this PR, that would be great to see.

[rgommers]

@nstarman there are 3 errors left when running mypy array_api_stubs/_draft/array_object.py (from lots more on main). There are more errors in other files - but if we are going to maintain static typing support in this repo, adding a CI job running Mypy seems feasible. We could start with this file, and expand it to all files later. WDYT?

[nstarman]

One thing that isn't nice is that Self now shows up instead of array in the signatures of the html docs.

If the array type is just text in the html docs and not 'clickable', then we can rename Self to array.

adding a CI job running Mypy seems feasible. We could start with this file, and expand it to all files later. WDYT?

That sounds like a good idea. I can try adding to the CI and create an ignore file that can later be whittled down as the package is made more type friendly.

[rgommers]

If the array type is just text in the html docs and not 'clickable', then we can rename Self to array.

That works. It's not clickable - there is no array object in the standard on purpose, because in the end we don't care whether it's named ndarray, Tensor, or something else.

[nstarman]

@rgommers Sorry for the long interlude. I can't seem to resolve 3 linking errors in the docs https://app.circleci.com/pipelines/github/data-apis/array-api/1088/workflows/ed2dfa75-c3c7-4845-8f28-e01fe35d0c8d/jobs/996?invite=true#step-103-1646

[asmeurer]

Sphinx is trying to create links for the type hints. I think to fix that you need to update the list(s) here https://github.com/data-apis/array-api/blob/main/src/_array_api_conf.py#L52-L77

[nstarman]

I can squash commits if you don't use Squash & Merge (assuming this is approved, of course :) ).
Thanks @asmeurer for the pointer.

[honno]

Nice work!

[honno]

Just to say more for myself than anyone else—I initially wasn't too sure if we wanted to use the py3.9+ type hint features, but:

• Forgot that from future import __annotations__ makes this work for py3.8
• py3.8 doesn't get updates end of next year
• The scientific ecosystem already is adopting py3.9+

[honno]

Looks good on my end!

I think pre-commit is a good addition (would love to add hooks for other spec things myself), but probably not too useful until we get a workflow that runs it. That said, I'm not too sure if we'd want to do that just yet 🤔

[nstarman]

I think pre-commit is a good addition (would love to add hooks for other spec things myself), but probably not too useful until we get a workflow that runs it. That said, I'm not too sure if we'd want to do that just yet 🤔

pre-commit.ci can be turned on, sans additional GH workflows.
black, and ruff would be wonderful as well!

[nstarman]

Looks good on my end!

Thanks @honno for the reviews and discussion.

[asmeurer]

Any reason to change array to Array? I think we might prefer to keep it as array like it is now.

[TomNicholas]

It would be great to see this released. It looks very developed to me, and would be useful for us within xarray.

[jorenham]

Perhaps this is out of scope, but this Literal is redundant (and as this change implies, also incorrect), since inf and -inf are both assignable to float.

[nstarman]

@rgommers. I'd be happy to finish this PR. I don't recall why this stalled out.

[nstarman]

Rebased and applied fixes for things that had changed.