John
John

Reputation: 497

Python type hint for int / array of ints

I have a function of an integer and a numpy array. It accepts a single int and a 1D array, or alternatively it accepts an array of ints and a 2D array which are then broadcast, e.g.

def norm_matrix(l, cov):
    norm_cov = cov / (l*(l+1))
    return norm_cov

My question is how to write the type hints in this case. Does this work?

def norm_matrix(
    l: Union[int, np.ndarray],
    cov: np.ndarray
):

and is there a better way to do it?

Upvotes: 0

Views: 687

Answers (2)

scr
scr

Reputation: 962

Your way is correct. It is the default way.

If you find it too cumbersome you can do use the bitwise operator | instead as mentioned in the Union documentation:

l: int | np.ndarray

But generally in Python we frown upon mixing input types since it complicates things. But like all rules, this one has an exception: multimethods - Here is an article by Guido explaining how they work.

Upvotes: 1

Paul_0
Paul_0

Reputation: 378

Your solution coul be improved by using numpy.typing.NDarray From the typing module of numpy

Upvotes: 1

Related Questions