Should non-generic classes with `__class_getitem__` be subscriptable in a type expression?

Typing
17

Numba uses a pretty neat syntax for declaring N-dimensional arrays of certain dtype when specifying the signature of a @jit function, where you write e.g. float32[:] for a one-dimensional single-precision array, or int8[:, :] for a two-dimensional array of 8-bit integers (docs). I first thought that this might be a good use-case for allowing non-generic __class_getitem__, but I then realized that these numba types like numba.float32 are instances rather than types, so they can just use __getitem__.

However, for a long time I have wondered whether it would be somehow possible to also use a similar terse syntax in in NumPy, i.e. being able to write np.int8[:, :] instead of the painfully verbose np.ndarray[tuple[int, int], np.dtype[np.float32]] [1]. But this always seemed unrealistic to me, and thought that it would require numpy-specific special-casing to be accepted as part of the typing spec.
This thread made me realize that this is actually not as far-fetched as I initially thought. Because by allowing __class_getitem__ to also be used to return TypeForms that are allowed to be used in annotations, then we’d effectively be able to define custom dependent type constructors. For the shaped dtype constructors in NumPy, that could look something like this:

type Dim = slice[None, None, None]
type Float32ND[ShapeT: tuple[int, ...]] = ndarray[ShapeT, dtype[float32]]

class float32:
    # --snip--

    @overload
    def __class_getitem__(cls, key: tuple[()], /) -> TypeForm[Float32ND[tuple[()]]: ...
    @overload
    def __class_getitem__(cls, key: Dim, /) -> TypeForm[Float32ND[tuple[int]]: ...
    @overload
    def __class_getitem__(cls, key: tuple[Dim, Dim], /) -> TypeForm[Float32ND[tuple[int, int]]: ...
    @overload
    def __class_getitem__(cls, key: tuple[Dim, Dim, Dim], /) -> TypeForm[Float32ND[tuple[int, int, int]]: ...

    # --snip--

    @overload
    def __class_getitem__(cls, key: tuple[Dim, ...], /) -> TypeForm[Float32ND[tuple[Any, ...]]: ...

I realize that this wouldn’t be easy to implement in type-checkers, to say the least. But it would make for an immensely powerful addition to Python’s type system, as it will allow us to overload types themselves. It would also be direct solution to What would it take to implement type mapping for generics? without requiring any language changes whatsoever.

Clearly, this __class_getitem__ + TypeForm stuff will require more research and discussion before going forward with it. So for now let’s just consider it a “potential use-case” that allowing non-generic __class_getitem__ definitions will enable. Then after that, we could discuss this TypeForm-driven dependent type constructor stuff separately.

And to be clear: it’s not my intention for this to be a serious proposal, just a potential one. I’m only bringing it up here because it requires us to allow __class_getitem__ in non-generic-type contexts. Then, if we decide on this, I’ll open a thread where we can discuss this in more detail [2].

  1. This could be made less painful using type aliases, but then you’d end up with a huge number of type aliases (one for each combination of dtype and rank of up to 3 or 4 or something), which clearly also isn’t ideal. ↩︎

  2. That is, unless it turns out not to be feasible, after all. ↩︎

3 Likes