ENH: Add an `ndarray` typing protocol

[anuppari]

I'd like to define a protocol that extends the numpy N-D Array. I can't create an abstract base class for unrelated reasons. I'd like to do something like the following:

from typing import Protocol
import numpy.typing import NDArrayProtocol  # Currently doesn't exist

class NdArrayWithMethod(NDArrayProtocol, Protocol):
    def method(self, input) -> None: pass

class ImpementsProtocol(np.ndarray):
    def method(self, input):
        ...

class AlsoImpementsProtocol(np.ndarray):
    def method(self, input):
        ...

def operates on protocol(in: NdArrayWithMethod):
   ...

The protocol cannot subclass ndarray directly, so I need to use something like NDArrayProtocol. Is there an ndarray protocol exposed somewhere?

[BvB93]

Is there an ndarray protocol exposed somewhere?

I'm afraid there isn't, and considering the sheer size of such a protocol I'm a rather reluctant to add one.

Interestingly, based on your examples it does sound like an intersection-type (xref python/typing#213) would resolve your issue, as this would allow you to describe an arbitrary np.ndarray subclass that implements a number of extra methods.

class MethodProtocol(Protocol):
    def method(self, input): ...

# expects an `ndarray` subclass that implements `method`
def operates_on_protocol(in: np.ndarray & MethodProtocol): ...

[vnmabus]

If you are interested in the future array API standard maybe you can also take a look at this:
data-apis/array-api#229

[anuppari]

Interestingly, based on your examples it does sound like an intersection-type (xref python/typing#213) would resolve your issue, as this would allow you to describe an arbitrary np.ndarray subclass that implements a number of extra methods.

Yeah, but until that is implemented, this was another approach I was considering.

If you are interested in the future array API standard maybe you can also take a look at this:
data-apis/array-api#229

Unfortunately, that doesn't nearly cover all the functions and capabilities of an np.ndarray.

@BvB93 Can the type annotations in the pyi files just be provided as a runtime_checkable Protocol rather than a regular class? That way there's not a redundant copy of the same interface and the current capabilities are still maintained.

[BvB93]

@BvB93 Can the type annotations in the pyi files just be provided as a runtime_checkable Protocol rather than a regular class? That way there's not a redundant copy of the same interface and the current capabilities are still maintained.

No, even if we ignore the fact that ndarray is not actually a protocol (and that's a big "if"), making it one would mean that you couldn't directly initialize it (just like with abstract baseclasses).

[anuppari]

What I mean to suggest is to only change the annotation to a protocol, i.e., something like

- class ndarray(_ArrayOrScalarCommon, Generic[_ShapeType, _DType_co]):
+ class ndarrayProtocol(_ArrayOrScalarCommon, Generic[_ShapeType, _DType_co], Protocol):

and then have the actual ndarray implementation explicitly subclass ndarrayProtocol but not Protocol. This way ndarray is still concrete, can be initialized, and type checking still works, but we also have a protocol available to extend as needed.

Might also need to add:

class ndarray(ndarrayProtocol[_ShapeType, _DType_co]):  # <-- not a Protocol
    pass

or something to the same __init__.pyi file

[BvB93]

Honestly, I see little value in adding, exposing and maintaining such a massive protocol.

The only use case were such protocol would be appropriate for the mixed nominal/structural subtyping you're trying to express here. And even this is more of band-aid solution necessitated by the lack of Intersection type.

I can't create an abstract base class for unrelated reasons.

Depending on the exact nature of these constraints, you could potentially define a type-check-only ABC and use that as baseclass.

import abc
import numpy as np
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    class FakeABC(np.ndarray, metaclass=abc.ABCMeta):
        @abc.abstractmethod
        def method(self, input): ...
else:
    FakeABC = np.ndarray
     

class ImpementsMethod(FakeABC):
    def method(self, input): ,..


class AlsoImpementsMethod(FakeABC):
    def method(self, input): ...


def operates_on_fake_abc(in: FakeABC): ...