Array API type hinting package

[honno]

Whilst I've been working on Array API support for Hypothesis and the test suite, I've been interested in type hinting some functions that take/return Array API objects (i.e. data type objects and arrays). I think it would be cool if there was a package to do this "properly".

How I picture it right now:

• Data types:
  • Specific e.g. Bool. UInt64 Int64, Float64 (these are type hinting the data type objects e.g. xp.bool, xp.int64 and xp.float64, but I'm not sure I like using something like BoolType as it's a bit tiresome IMO)
    • or something like Int[64] and Float[64]...?
  • Categorical i.e. Integer, SignedInteger, UnsignedInteger, Floating
  • DataType - any data type objects
• array object Array, a generic which can take data type hints like Array[Int64] or Array[Floating]
  • could also isinstance() by way of checking __array_namespace__, which notably we suggest in the spec for general conformance checking
• ArrayModule for the xp namespace
• Other return types like for iinfo and finfo
• Helper types:
  • Scalar as an alias of Union[bool, int, float] (in the future also complex)
  • RealScalar as an alias of Union[int, float]
  • Pre-emptively ComplexScalar as an alias of Union[int, float, complex]

Obviously the utility of type hints for actual type checking are pretty limited (save checking a generic Array and ArrayModule), unless I'm missing some easy way to grab the array module xp so that something like Bool could use obj == xp.bool. ^{Maybe there could be a make_types_namespace-ish to bound xp so these type hints can actually help mypy lol}

That said I like using them just for a scrappy documentation myself, and I wonder if that could be the case for folks developing array-consuming methods for the Array API. I think a generic Array type could be interesting for writing spec, as we could annotate input and output arrays with their categorial type hints

def sort(x: Array, /, *, axis:int = -1, descending:bool = False, stable: bool = True) -> Array:
    ...

def equal(x1: Array, x2: Array, /) -> Array[Bool]:
    ...

Also a niche use case—such type hints could maybe be used in Hypothesis for inferring strategies via type hints.

Anywho, interested if folks have ideas for what a good API for such a package would look like. Thought to myself I'd like to work on this as a fun side project in the not-so-near future (I have 0 idea how to make stuff like generic type hints so thought it'd be a good excuse to learn haha).

[leofang]

cc: @BvB93

[BvB93]

Data types:

I agree that it might be nice to have something like this at some point.

Interstingly, if the dtype objects are proper classes (or at least treated as such when type checking) then there is no need to distinguish between the object and the annotation, i.e. this would allow the likes of xp.bool and xp.int64 to be directly used as annotation.

from typing import TYPE_CHECKING
import numpy as np

# For example; the `TYPE_CHECKING` may or may not be necessary depending on the particular implementation
if TYPE_CHECKING: 
    class float64(np.dtype[np.float64]): ...
else:
    float64 = np.dtype(np.float64)

foo_dtype: float64

array object Array, a generic which can take data type hints like Array[Int64] or Array[Floating]

I would love to see Array being made generic at some point, but there are definetly a number of challenges here: how to automatically infere the dtype? How the encode casting rules? How many generic parameters? etc.

ArrayModule for the xp namespace

This will likelly be quite challenging, as there is no convenient way of typing namespaces (e.g. via Literal modules; xref data-apis/array-api#267). That is, unless we're willing to considering something like a mypy plugin (which could be rather helpfull here, I suspect).

Helper types:

The aliases (sort of) already exist with the numeric tower from PEP 484, namely: complex is supposed to be treated as a supertype of float, which is in turn treated as a supertype of int (which is a proper superclass of bool). In other words, you can always simplify the proposed Scalar and RealScalar types to float.

[rgommers]

So if you want to have an object that can be used as an annotation, and it should cover any object that has a compatible name (float64 etc.) and does the right thing within some library, then does that require a Protocol with no methods that each library can inherit from?

[BvB93]

Add a protocol here would be quite tricky. The only requirement of data type objects here is that they must be comparible, thus implying that we'd need an __eq__-based protocol here. The problem is that due to the broad definition of object.__eq__ the vast majority of python objects will satisfy this hypothetical protocol, effectivelly making it useless.

No, instead I'd to expose a more concrete type here, something more akin to output of type(float64).

[rgommers]

No, instead I'd to expose a more concrete type here, something more akin to output of type(float64).

I'm not sure how to do that in a way that makes sense. There will not be a type relationship between dtype objects from different libraries, and the type hinting package here is not something that every array library will want as a runtime dependency.

[BvB93]

There will not be a type relationship between dtype objects from different libraries, and the type hinting package here is not something that every array library will want as a runtime dependency.

True, I don't think distributing such an object from a central place will be feasible;
protocols are unviable here due to __eq__ and a common abstract base class even less so.

This is why I was more leaning towards exposing a <DType> type in the API proper. The relevant type should exist somewhere in a private namespace (as it is used in the construction of float64 and the likes), all that would be required is to expose it.

[rgommers]

The relevant type should exist somewhere in a private namespace (as it is used in the construction of float64 and the likes), all that would be required is to expose it.

I don't think that is the case - it assumes some class hierarchy between dtypes that is not necessary and may not be present in all libraries. I don't think requiring such a hierarchy just to making type checkers happy is a good idea, even aside from exposing an object that has no other use.

[patrick-kidger]

Just came across this older thread. I don't know what kind of interest this topic has at the moment but I thought I'd point folks at the TorchTyping package and the newly-released jaxtyping package for interest.