Python - typing - union of subscriptable type

I'd like to create an Array type which should be subscriptable and be a Union of typing.List and numpy.ndarray types. I know numpy doesn't come with stubs, but those numpy stubs (by Machinalis) should work fine as they are subsriptable.

This is the expected behavior:

def foo(bar: Array[int])->None:
    pass

foo([1,2,3])          # No typing error
foo(numpy.arange(4))  # No typing error
foo((1,2,3))          # Error: Expected Array[int], got Tuple[int]
foo([1.,2.,3.])       # Error: Expected Array[int], got Array[float]

I've tried a few things but none of them work as expected.

How would you do that in Python 3.7?

I'll also accept some kind of duck-typing solution, even if it doesn't satisfy the Tuple error. The main point of focus is to create subscriptable union of subscriptable types.

Thanks.

My best attempt : (mypy errors in comments)

class _meta_getitem(type):
    def __getitem__(cls, x):
        return cls.__getitem__(cls, x)

class Array(metaclass=_meta_getitem):

    def __getitem__(self, element_type: type) -> type:
        array_type = typing.Union[List[element_type],  # error: Invalid type "element_type"
                                  numpy.ndarray[element_type]]
        return typing.NewType("Array[{}]".format(element_type.__name__), 
                              array_type)  # The type alias to Union is invalid in runtime context

if __name__ == "__name__":
    x: Array[int] = numpy.arange(4) # "Array" expects no type arguments, but 1 given

Solution

Creating a type alias of Union[List[T], Array[T]] ought to work:

from typing import TypeVar, Union, List

T = TypeVar('T')
Array = Union[List[T], numpy.ndarray[T]]

def foo(bar: Array[int]) -> None: pass

See the mypy docs on generic type aliases for more info about this technique.

This code may potentially fail at runtime since numpy.ndarray isn't actually subscriptable at runtime, only in the type-hinting world. You can work around this by hiding your custom type hint behind a typing.TYPE_CHECKING guard, which is always false at runtime and true at type-check time.

You can do this relatively cleanly in Python 3.7+:

from __future__ import annotations
from typing import TypeVar, Union, List, TYPE_CHECKING

if TYPE_CHECKING:
    T = TypeVar('T')
    Array = Union[List[T], numpy.ndarray[T]]

def foo(bar: Array[int]) -> None: pass

You have to wrap your Array[int] inside of a string for older versions of Python 3, however:

from typing import TypeVar, Union, List, TYPE_CHECKING

if TYPE_CHECKING:
    T = TypeVar('T')
    Array = Union[List[T], numpy.ndarray[T]]

def foo(bar: "Array[int]") -> None: pass

Note that attempting to construct your own Array type hint by composing together several other type hints at runtime is unlikely to work: static analysis tools like mypy work by actually analyzing your code without running it: it's not actually going to attempt to evaluate anything inside of your custom Array class.

A little more generally speaking, attempting to "use" type hints at runtime tends to be fraught with peril: they're really only meant to be used as type hints.

Finally, you appear to be misunderstanding what NewType is for. I recommend reading the relevant docs for more info.