np.clip typing is not as specific as ideal; existing code doesn't type-check

[bdowning]

The typing definition of np.clip always seems to return a union of number or ndarray for any array-like input. Unfortunately this breaks typechecking in existing code we had when we pulled in numpy 1.20.0.

        x0_cr, x1_cr = np.clip((x0_r, x1_r), 0, width)
        y0_cr, y1_cr = np.clip((y0_r, y1_r), 0, height)

It'd be much nicer if it could know when the type returned would be a scalar or a vector based on the input type, which I think should be achievable with the right overloads.

Reproducing code example:

from typing import TYPE_CHECKING

import numpy as np

result = np.clip((0, 10), 4, 6)
if TYPE_CHECKING:
    reveal_type(result)
a, b = result
print(a, b)

$ poetry run python foo.py
4 6
$ poetry run mypy foo.py  
foo.py:7: note: Revealed type is 'Union[numpy.number[Any], numpy.ndarray]'
foo.py:8: error: 'numpy.number[Any]' object is not iterable
Found 1 error in 1 file (checked 1 source file)

Error message:

See above for reported typechecking error on destructuring.

NumPy/Python version information:

$ poetry run python --version
Python 3.9.1
$ poetry show
mypy              0.800   Optional static typing for Python
mypy-extensions   0.4.3   Experimental type system extensions for programs checked with the mypy typechecker.
numpy             1.20.0  NumPy is the fundamental package for array computing with Python.
typed-ast         1.4.2   a fork of Python 2 and 3 ast modules with type comment support
typing-extensions 3.7.4.3 Backported and Experimental Type Hints for Python 3.5+

[jklaise]

This behaviour is more widespread than that involving reduction functions like np.sum, e.g. the following doesn't type-check:

import numpy as np

def sum_axis(x: np.ndarray) -> np.ndarray:
    return np.sum(x, axis=1)

as the expected type is Union[numpy.number[Any], numpy.ndarray].

I've been trying to find the reasoning behind this, I'm not sure if it is ever possible for a reduction function along an axis (as long as axis is not 0 for 1-dimensional arrays) to produce a number. Maybe it's difficult to decide without performing a runtime check on the shape, but it should be possible to refine the return types in case a non-zero axis argument is passed.

EDIT: to expand a little bit, the following pattern either needs an explicit cast which is rather annoying or to ignore the type:

import numpy as np
from typing import cast


def fun(x: np.ndarray) -> np.ndarray:
    return x


# type annotation needed here as otherwised inferred as `Any`
arr3d = np.random.rand(3, 4, 5)  # type: np.ndarray

# the following is inferred as `Union[number, np.ndarray]`.
arr2d = np.sum(arr3d, axis=2)

# the following will not typecheck:
# error: Argument 1 to "fun" has incompatible type "Union[number[Any], ndarray]"; expected "ndarray"
result = fun(arr2d)

# the following will typecheck
result = fun(arr2d)  # type: ignore

# the following will also typecheck
arr2d = cast(np.ndarray, arr2d)
result = fun(arr2d)

[BvB93]

So a PR addressing the issue is up at #18885, the respective backport following later.

To summarize (the PR contains more details): the likes of np.number[Any] | np.ndarray were used a shortcut to deal with the difficult task of typing arbitrary shaped array-likes whose dtype is not quite clear. Over the course of time it has become more apparent that the use of these impromptu unsafe unions is somewhat problematic; honestly it's the worst of both worlds.

Without the ability the describe 0D arrays (see #16544) it does mean that a sacrifice has to be made, i.e. we cannot describe numpy's 0D-to-scalar casting (or even attempt to) as can't type the array shape/dimensionality in the first place. This is honestly the lesser of two evils, as it will be an inconvenience exclusively only 0D arrays, rather than all arrays of any arbitrary dimensionality.