How to type annotate mathematical operations that supports built-in numerics, collections and numpy arrays?

Python Help
,
1

In scientific code that deals with mathematical operations, before type annotations, it is common to assume that any mathematical function could take in a numpy array. To be explicit you just add to your docstring that you can take array_like or ndarray (as is conventional in NumFocus project docs).

Today, with type annotations, I am at a lost on how to annotate something as simple as:

import numpy as np

def get_variance_plus_one(sample):  # A stupid function, just for example
    return 1 + np.var(sample)

What is the proper type for sample?

  • Any | Sequence[Any]: Too generic for numerical math. Does not support numpy arrays

  • float | Sequence[float]: Does not support complex number (rare use-case) or numpy arrays.

  • number.Number | Sequence[number.Number]: Does not support numpy arrays.

  • np.typing.ArrayLike: Too generic so much so that __add__ and __mul__ aren’t necessary. Look at this Pylance error:

      Operator "+" not supported for types "ArrayLike" and "ArrayLike"
  Operator "+" not supported for types "_SupportsArray[dtype[Unknown]]" and "_SupportsArray[dtype[Unknown]]"
  Operator "+" not supported for types "_SupportsArray[dtype[Unknown]]" and "_NestedSequence[_SupportsArray[dtype[Unknown]]]"
  Operator "+" not supported for types "_SupportsArray[dtype[Unknown]]" and "bool"
  Operator "+" not supported for types "_SupportsArray[dtype[Unknown]]" and "int"
  Operator "+" not supported for types "_SupportsArray[dtype[Unknown]]" and "float"
  Operator "+" not supported for types "_SupportsArray[dtype[Unknown]]" and "complex"
  Operator "+" not supported for types "_SupportsArray[dtype[Unknown]]" and "str"
  Operator "+" not supported for types "_SupportsArray[dtype[Unknown]]" and "bytes

  • np.typing.NDArray[np.number]: Similar problem to ArrayLike

  • np.typing.NDArray[np.integer | np.floating | np.complexfloating | ...? ] | number.Number
    This is way too verbose and prone to mistakes!

Apologies if this issue is better addressed to NumPy. Feel free to send me a link to the relevant community channel.

1 Like
2

Here’s a sensible solution which provides more type stability with np.asanyarray()

import numpy as np
import numpy.typing as npt

def get_variance_plus_one(sample: npt.ArrayLike) -> np.number:
    sample = np.asanyarray(sample)
    return 1 + np.var(sample)