Mypy errors when using ArrayLike

I don't understand how I should be using ArrayLike in my code. If check mypy, I keep getting errors when I try to use the variables for anything without calling cast. I am trying to define function signatures that work with ndarray as well as regular lists.

For example, the code below

import numpy.typing as npt
import numpy as np

from typing import Any

def f(a: npt.ArrayLike) -> int:
    return len(a)

def g(a: npt.ArrayLike) -> Any:
    return a[0]

print(f(np.array([0, 1])), g(np.array([0, 1])))
print(f([0, 1]), g([0, 1]))

give me theses errors for f() and g():

Argument 1 to "len" has incompatible type "Union[_SupportsArray[dtype[Any]], _NestedSequence[_SupportsArray[dtype[Any]]], bool, int, float, complex, str, bytes, _NestedSequence[Union[bool, int, float, complex, str, bytes]]]"; expected "Sized"  [arg-type]

Value of type "Union[_SupportsArray[dtype[Any]], _NestedSequence[_SupportsArray[dtype[Any]]], bool, int, float, complex, str, bytes, _NestedSequence[Union[bool, int, float, complex, str, bytes]]]" is not indexable  [index]

The purpose of numpy.typing.ArrayLike is to be able to annotate

> objects that can be coerced into an ndarray.

With that purpose in mind, they defined the type to be the following union:

Union[
    _SupportsArray[dtype[Any]],
    _NestedSequence[_SupportsArray[dtype[Any]]],
    bool,
    int,
    float,
    complex,
    str,
    bytes,
    _NestedSequence[Union[bool, int, float, complex, str, bytes]]
]

_SupportsArray is simply a protocol with an __array__ method. It neither requires __len__ (for use with the len function) nor __getitem__ (for indexing) to be implemented.

_NestedSequence is a more restrictive protocol that does actually require __len__ and __getitem__.

But the problem with this code is that the parameter annotation is that union:

import numpy.typing as npt

...

def f(a: npt.ArrayLike) -> int:
    return len(a)

So a could be a sequence-like object that supports __len__, but it could also be just an object that supports __array__ and nothing else. It could even be just an int for example (see the union again). Therefore the call len(a) is unsafe.

Similarly, here the item access is not type safe because a might not implement __getitem__:

...

def g(a: npt.ArrayLike) -> Any:
    return a[0]

So the reason it is not working for you is that it is not meant to be used as an annotation for things that are numpy arrays or other sequences; it is meant to be used for things that can be turned into numpy arrays.

If you want to annotate your functions f and g to take both lists and numpy arrays, you could just use a union of list and NDArray like list[Any] | npt.NDArray[Any].

If you want to have a wider annotation to accommodate any type that has __len__ and __getitem__, you need to define your own protocol:

from typing import Any, Protocol, TypeVar

import numpy as np

T = TypeVar("T", covariant=True)


class SequenceLike(Protocol[T]):
    def __len__(self) -> int: ...
    def __getitem__(self, item: int) -> T: ...


def f(a: SequenceLike[Any]) -> int:
    return len(a)


def g(a: SequenceLike[T]) -> T:
    return a[0]


print(f(np.array([0, 1])), g(np.array([0, 1])))
print(f([0, 1]), g([0, 1]))

To be more precise __getitem__ should probably also take slice objects, but the overloads may be overkill for you.

Try to comply with the numpy.typing followed by suggested settings under mypy config

[mypy]
plugins = numpy.typing.mypy_plugin