Search code examples
pythonnumpytyping

vector-like input - numpy - typing


I have a problem with python typing of the function with vector-like (1-D) input, which will be looped. The input should be one of list, tuple, numpy.ndarray or similar. What is important I will be checking the length for the input. Thus the typing.Sequence should be perfectly suited. But is not as np.array is not a Sequence, it is e.g. Iterable. The Iterable looks to not be the right choice as I need to use len().

Thank you for recommendations.

# arg1 and arg2 should be one of `list`, `tuple`, `numpy.ndarray` or similar
def fun(arg1: ?, arg2: ?) -> None:
    # assert isinstance(arg1, Sequence-Like), ""
    # assert isinstance(arg12, Sequence-Like), ""
    assert len(arg1) > 0, "length of the arg1 has to be at least 1"
    assert len(arg1) == len(arg2), "..."


Solution

  • A straightforward (perhaps hacky) solution is to create a Union specifying the types your function supports:

    def fun(arg1: Union[Sequence, numpy.ndarray, pandas.Series], arg2: ...
    

    You can also pre-define the type so you don't take up so much room in the function signature:

    SequenceLike = Union[Sequence, numpy.ndarray, pandas.Series]
    
    def fun(arg1: SequenceLike, arg2: SequenceLike) -> None:
        ...