Skip to content

TYP: np.array improved shape-typing - #32246

Merged
charris merged 1 commit into
numpy:mainfrom
jorenham:typing/array/improve-shape-typing-some-more
Aug 11, 2026
Merged

TYP: np.array improved shape-typing#32246
charris merged 1 commit into
numpy:mainfrom
jorenham:typing/array/improve-shape-typing-some-more

Conversation

@jorenham

@jorenham jorenham commented Aug 10, 2026 •
edited
Loading

Copy link
Copy Markdown
Member

There were still some shape-typing gaps left for np.array (the most popular numpy function) after #32081:

  • 0d builtin scalars (bool/int/float/complex)
  • 3d nested sequences of builtin scalars (only 1d and 2d sequences were covered)
  • 0d-3d nested sequences with dtype= given

There are some things that still aren't shape-typed, such as ndmin larger than the input, and >=4d nested sequences. I don't think adding overloads for that is worth it though, because code search shows that those patterns aren't used very often

Fully written by me, pre-reviewed by AI.

@jorenham jorenham added this to the 2.6.0 Release milestone Aug 10, 2026
@github-actions

github-actions Bot commented Aug 10, 2026

Copy link
Copy Markdown

Diff from mypy_primer, showing the effect of this PR on type check results on a corpus of open source code:

jax (https://github.com/google/jax)
- jax/_src/internal_test_util/test_harnesses.py:839: error: Incompatible types in assignment (expression has type "ndarray[tuple[Any, ...], dtype[number[_32Bit, float]]]", variable has type "ndarray[tuple[Any, ...], dtype[floating[_32Bit]]]")  [assignment]
+ jax/_src/internal_test_util/test_harnesses.py:839: error: Incompatible types in assignment (expression has type "ndarray[tuple[int], dtype[number[_32Bit, float]]]", variable has type "ndarray[tuple[int, int], dtype[floating[_32Bit]]]")  [assignment]
- jax/_src/internal_test_util/test_harnesses.py:1115: error: Incompatible types in assignment (expression has type "ndarray[tuple[Any, ...], dtype[integer[_32Bit]]]", variable has type "ndarray[tuple[Any, ...], dtype[signedinteger[_32Bit]]]")  [assignment]
+ jax/_src/internal_test_util/test_harnesses.py:1115: error: Incompatible types in assignment (expression has type "ndarray[tuple[int, ...], dtype[integer[_32Bit]]]", variable has type "ndarray[tuple[()], dtype[signedinteger[_32Bit]]]")  [assignment]
- jax/_src/internal_test_util/test_harnesses.py:1657: error: Incompatible types in assignment (expression has type "ndarray[tuple[Any, ...], dtype[number[_32Bit, float]]]", variable has type "ndarray[tuple[Any, ...], dtype[floating[_32Bit]]]")  [assignment]
+ jax/_src/internal_test_util/test_harnesses.py:1657: error: Incompatible types in assignment (expression has type "ndarray[tuple[int], dtype[number[_32Bit, float]]]", variable has type "ndarray[tuple[int, int], dtype[floating[_32Bit]]]")  [assignment]
- jax/_src/internal_test_util/test_harnesses.py:1718: error: Incompatible types in assignment (expression has type "list[object]", variable has type "list[ndarray[tuple[Any, ...], dtype[floating[_32Bit]]]]")  [assignment]
+ jax/_src/internal_test_util/test_harnesses.py:1718: error: Incompatible types in assignment (expression has type "list[object]", variable has type "list[ndarray[tuple[int], dtype[floating[_32Bit]]]]")  [assignment]
- jax/_src/internal_test_util/test_harnesses.py:2231: error: Incompatible types in assignment (expression has type "tuple[ndarray[tuple[Any, ...], dtype[unsignedinteger[_8Bit]]] | ndarray[tuple[Any, ...], dtype[unsignedinteger[_32Bit]]] | int, ...]", variable has type "list[list[int]]")  [assignment]
+ jax/_src/internal_test_util/test_harnesses.py:2231: error: Incompatible types in assignment (expression has type "tuple[ndarray[tuple[()], dtype[unsignedinteger[_8Bit]]] | ndarray[tuple[()], dtype[unsignedinteger[_32Bit]]] | int, ...]", variable has type "list[list[int]]")  [assignment]
- jax/_src/internal_test_util/test_harnesses.py:2231: error: Incompatible types in assignment (expression has type "tuple[ndarray[tuple[Any, ...], dtype[unsignedinteger[_8Bit]]] | ndarray[tuple[Any, ...], dtype[unsignedinteger[_32Bit]]] | int, ...]", variable has type "tuple[int, ...]")  [assignment]
+ jax/_src/internal_test_util/test_harnesses.py:2231: error: Incompatible types in assignment (expression has type "tuple[ndarray[tuple[()], dtype[unsignedinteger[_8Bit]]] | ndarray[tuple[()], dtype[unsignedinteger[_32Bit]]] | int, ...]", variable has type "tuple[int, ...]")  [assignment]
- jax/_src/internal_test_util/test_harnesses.py:2290: error: Incompatible types in assignment (expression has type "tuple[ndarray[tuple[Any, ...], dtype[unsignedinteger[_8Bit]]] | ndarray[tuple[Any, ...], dtype[unsignedinteger[_32Bit]]] | int, ...]", variable has type "list[list[int]]")  [assignment]
+ jax/_src/internal_test_util/test_harnesses.py:2290: error: Incompatible types in assignment (expression has type "tuple[ndarray[tuple[()], dtype[unsignedinteger[_8Bit]]] | ndarray[tuple[()], dtype[unsignedinteger[_32Bit]]] | int, ...]", variable has type "list[list[int]]")  [assignment]

colour (https://github.com/colour-science/colour)
+ colour/geometry/section.py:118: error: Tuple index out of range  [misc]
- colour/utilities/tests/test_array.py:1680: error: Incompatible types in assignment (expression has type "ndarray[tuple[Any, ...], dtype[generic[Any]]]", variable has type "int")  [assignment]
- colour/utilities/tests/test_array.py:1697: error: Incompatible types in assignment (expression has type "ndarray[tuple[Any, ...], dtype[generic[Any]]]", variable has type "int")  [assignment]
+ colour/utilities/tests/test_array.py:1680: error: Incompatible types in assignment (expression has type "ndarray[tuple[int], dtype[number[Any, Any] | numpy.bool[builtins.bool] | character[Any] | datetime64[Any] | timedelta64[Any]]]", variable has type "int")  [assignment]
+ colour/utilities/tests/test_array.py:1697: error: Incompatible types in assignment (expression has type "ndarray[tuple[int], dtype[number[Any, Any] | numpy.bool[builtins.bool] | character[Any] | datetime64[Any] | timedelta64[Any]]]", variable has type "int")  [assignment]

pandas (https://github.com/pandas-dev/pandas)
- pandas/tests/groupby/methods/test_rank.py:72: error: List item 0 has incompatible type "list[Timestamp | NaTType]"; expected "ndarray[tuple[Any, ...], dtype[Any]]"  [list-item]
+ pandas/tests/groupby/methods/test_rank.py:72: error: List item 0 has incompatible type "list[Timestamp | NaTType]"; expected "ndarray[tuple[int], dtype[Any]]"  [list-item]
- pandas/tests/groupby/methods/test_rank.py:79: error: List item 1 has incompatible type "list[Timestamp | NaTType]"; expected "ndarray[tuple[Any, ...], dtype[Any]]"  [list-item]
+ pandas/tests/groupby/methods/test_rank.py:79: error: List item 1 has incompatible type "list[Timestamp | NaTType]"; expected "ndarray[tuple[int], dtype[Any]]"  [list-item]
- pandas/tests/groupby/methods/test_rank.py:86: error: List item 2 has incompatible type "list[Timedelta | NaTType]"; expected "ndarray[tuple[Any, ...], dtype[Any]]"  [list-item]
+ pandas/tests/groupby/methods/test_rank.py:86: error: List item 2 has incompatible type "list[Timedelta | NaTType]"; expected "ndarray[tuple[int], dtype[Any]]"  [list-item]
- pandas/tests/groupby/methods/test_rank.py:93: error: List item 3 has incompatible type "list[Period | Any]"; expected "ndarray[tuple[Any, ...], dtype[Any]]"  [list-item]
+ pandas/tests/groupby/methods/test_rank.py:93: error: List item 3 has incompatible type "list[Period | Any]"; expected "ndarray[tuple[int], dtype[Any]]"  [list-item]
- pandas/tests/groupby/methods/test_rank.py:202: error: List item 0 has incompatible type "list[float | Timestamp | NaTType]"; expected "ndarray[tuple[Any, ...], dtype[Any]]"  [list-item]
+ pandas/tests/groupby/methods/test_rank.py:202: error: List item 0 has incompatible type "list[float | Timestamp | NaTType]"; expected "ndarray[tuple[int], dtype[Any]]"  [list-item]
- pandas/tests/groupby/methods/test_rank.py:212: error: List item 1 has incompatible type "list[float | Timestamp | NaTType]"; expected "ndarray[tuple[Any, ...], dtype[Any]]"  [list-item]
+ pandas/tests/groupby/methods/test_rank.py:212: error: List item 1 has incompatible type "list[float | Timestamp | NaTType]"; expected "ndarray[tuple[int], dtype[Any]]"  [list-item]
- pandas/tests/groupby/methods/test_rank.py:222: error: List item 2 has incompatible type "list[float | Timedelta | NaTType]"; expected "ndarray[tuple[Any, ...], dtype[Any]]"  [list-item]
+ pandas/tests/groupby/methods/test_rank.py:222: error: List item 2 has incompatible type "list[float | Timedelta | NaTType]"; expected "ndarray[tuple[int], dtype[Any]]"  [list-item]
- pandas/tests/groupby/methods/test_rank.py:232: error: List item 3 has incompatible type "list[float | Period | Any]"; expected "ndarray[tuple[Any, ...], dtype[Any]]"  [list-item]
+ pandas/tests/groupby/methods/test_rank.py:232: error: List item 3 has incompatible type "list[float | Period | Any]"; expected "ndarray[tuple[int], dtype[Any]]"  [list-item]

xarray (https://github.com/pydata/xarray)
+ xarray/tests/test_backends.py: note: In member "test_roundtrip_character_array" of class "NetCDF4Base":
+ xarray/tests/test_backends.py:1845: error: Incompatible types in assignment (expression has type "ndarray[tuple[int], dtype[Any]]", variable has type "ndarray[tuple[int, int], dtype[Any]]")  [assignment]
+ xarray/tests/test_backends.py: note: At top level:
+ xarray/tests/test_conventions.py: note: In member "test_booltype_array" of class "TestBoolTypeArray":
+ xarray/tests/test_conventions.py:41: error: Incompatible types in assignment (expression has type "ndarray[tuple[int, int], dtype[Any]]", variable has type "ndarray[tuple[int], dtype[Any]]")  [assignment]

pandas-stubs (https://github.com/pandas-dev/pandas-stubs)
+ tests/series/test_truediv.py:192: error: Unused "type: ignore" comment  [unused-ignore]
+ tests/series/test_truediv.py:202: error: Unused "type: ignore" comment  [unused-ignore]
+ tests/series/test_truediv.py:212: error: Unused "type: ignore" comment  [unused-ignore]
+ tests/series/test_truediv.py:220: error: Unused "type: ignore" comment  [unused-ignore]
+ tests/series/test_floordiv.py:166: error: Unused "type: ignore" comment  [unused-ignore]
+ tests/series/test_floordiv.py:178: error: Unused "type: ignore" comment  [unused-ignore]
+ tests/series/timestamp/test_sub.py:115: error: Expression is of type "ndarray[tuple[Any, ...], dtype[Any]]", not "ndarray[tuple[Any, ...], dtype[datetime64[Any]]]"  [assert-type]
+ tests/series/timestamp/test_add.py:126: error: Expression is of type "ndarray[tuple[Any, ...], dtype[Any]]", not "ndarray[tuple[Any, ...], dtype[datetime64[Any]]]"  [assert-type]
+ tests/series/timestamp/test_add.py:132: error: Unused "type: ignore" comment  [unused-ignore]
+ tests/series/timestamp/test_add.py:136: error: Unused "type: ignore" comment  [unused-ignore]
+ tests/series/timedelta/test_sub.py:118: error: Expression is of type "ndarray[tuple[Any, ...], dtype[Any]]", not "ndarray[tuple[Any, ...], dtype[datetime64[Any]]]"  [assert-type]
+ tests/series/timedelta/test_sub.py:122: error: Unused "type: ignore" comment  [unused-ignore]
+ tests/series/timedelta/test_add.py:108: error: Expression is of type "ndarray[tuple[Any, ...], dtype[Any]]", not "ndarray[tuple[Any, ...], dtype[datetime64[Any]]]"  [assert-type]
+ tests/series/str/test_mul.py:146: error: Expression is of type "ndarray[tuple[Any, ...], dtype[Any]]", not "Any"  [assert-type]
+ tests/series/str/test_mul.py:155: error: Unused "type: ignore" comment  [unused-ignore]
+ tests/series/str/test_mul.py:164: error: Unused "type: ignore" comment  [unused-ignore]
+ tests/series/int/test_truediv.py:197: error: Expression is of type "ndarray[tuple[Any, ...], dtype[float64]]", not "Any"  [assert-type]
+ tests/series/int/test_truediv.py:209: error: Unused "type: ignore" comment  [unused-ignore]
+ tests/series/int/test_truediv.py:217: error: Unused "type: ignore" comment  [unused-ignore]
+ tests/series/int/test_truediv.py:229: error: Unused "type: ignore" comment  [unused-ignore]
+ tests/series/int/test_truediv.py:243: error: Unused "type: ignore" comment  [unused-ignore]
+ tests/series/int/test_mul.py:144: error: Expression is of type "ndarray[tuple[Any, ...], dtype[Any]]", not "Any"  [assert-type]
+ tests/series/int/test_mul.py:152: error: Unused "type: ignore" comment  [unused-ignore]
+ tests/series/int/test_mul.py:162: error: Unused "type: ignore" comment  [unused-ignore]
+ tests/series/int/test_floordiv.py:142: error: Expression is of type "ndarray[tuple[Any, ...], dtype[Any]]", not "Any"  [assert-type]
+ tests/series/int/test_floordiv.py:150: error: Unused "type: ignore" comment  [unused-ignore]
+ tests/series/int/test_floordiv.py:158: error: Unused "type: ignore" comment  [unused-ignore]
+ tests/series/float/test_truediv.py:197: error: Expression is of type "ndarray[tuple[Any, ...], dtype[float64]]", not "Any"  [assert-type]
+ tests/series/float/test_truediv.py:209: error: Unused "type: ignore" comment  [unused-ignore]
+ tests/series/float/test_truediv.py:217: error: Unused "type: ignore" comment  [unused-ignore]
+ tests/series/float/test_truediv.py:229: error: Unused "type: ignore" comment  [unused-ignore]
+ tests/series/float/test_truediv.py:243: error: Unused "type: ignore" comment  [unused-ignore]
+ tests/series/float/test_mul.py:144: error: Expression is of type "ndarray[tuple[Any, ...], dtype[Any]]", not "Any"  [assert-type]
+ tests/series/float/test_mul.py:152: error: Unused "type: ignore" comment  [unused-ignore]
+ tests/series/float/test_mul.py:162: error: Unused "type: ignore" comment  [unused-ignore]
+ tests/series/float/test_floordiv.py:141: error: Expression is of type "ndarray[tuple[Any, ...], dtype[Any]]", not "Any"  [assert-type]
+ tests/series/float/test_floordiv.py:149: error: Unused "type: ignore" comment  [unused-ignore]
+ tests/series/float/test_floordiv.py:157: error: Unused "type: ignore" comment  [unused-ignore]
+ tests/series/complex/test_mul.py:158: error: Expression is of type "ndarray[tuple[Any, ...], dtype[Any]]", not "Any"  [assert-type]
+ tests/series/complex/test_mul.py:166: error: Unused "type: ignore" comment  [unused-ignore]
+ tests/series/complex/test_mul.py:182: error: Unused "type: ignore" comment  [unused-ignore]
+ tests/series/bool/test_mul.py:144: error: Expression is of type "ndarray[tuple[Any, ...], dtype[Any]]", not "Any"  [assert-type]
+ tests/series/bool/test_mul.py:152: error: Unused "type: ignore" comment  [unused-ignore]
+ tests/series/bool/test_mul.py:162: error: Unused "type: ignore" comment  [unused-ignore]
+ tests/indexes/str/test_mul.py:110: error: Expression is of type "ndarray[tuple[Any, ...], dtype[Any]]", not "Any"  [assert-type]
+ tests/indexes/int/test_truediv.py:103: error: Expression is of type "ndarray[tuple[Any, ...], dtype[float64]]", not "Any"  [assert-type]
+ tests/indexes/int/test_mul.py:108: error: Expression is of type "ndarray[tuple[Any, ...], dtype[Any]]", not "Any"  [assert-type]
+ tests/indexes/int/test_floordiv.py:102: error: Expression is of type "ndarray[tuple[Any, ...], dtype[Any]]", not "Any"  [assert-type]
+ tests/indexes/float/test_truediv.py:103: error: Expression is of type "ndarray[tuple[Any, ...], dtype[float64]]", not "Any"  [assert-type]
+ tests/indexes/float/test_mul.py:108: error: Expression is of type "ndarray[tuple[Any, ...], dtype[Any]]", not "Any"  [assert-type]
+ tests/indexes/float/test_floordiv.py:102: error: Expression is of type "ndarray[tuple[Any, ...], dtype[Any]]", not "Any"  [assert-type]
+ tests/indexes/complex/test_mul.py:110: error: Expression is of type "ndarray[tuple[Any, ...], dtype[Any]]", not "Any"  [assert-type]
+ tests/indexes/bool/test_mul.py:108: error: Expression is of type "ndarray[tuple[Any, ...], dtype[Any]]", not "Any"  [assert-type]
+ tests/indexes/test_indexes.py:607: error: Expression is of type "IntervalIndex[Interval[int]]", not "IntervalIndex[Interval[Timestamp]]"  [assert-type]
+ tests/arrays/test_categorical.py:63: error: Expression is of type "ndarray[tuple[int], dtype[str_]]", not "ndarray[tuple[Any, ...], dtype[str_]]"  [assert-type]

optuna (https://github.com/optuna/optuna)
  ./optuna/study/_multi_objective.py:182: error: INTERNAL ERROR -- Please try using mypy master on GitHub:

@jorenham

jorenham commented Aug 11, 2026

Copy link
Copy Markdown
Member Author

RE: the new mypy_primer errors:

  • jax: some changed error messages; nothing to worry about
  • colour: The new [misc] error (src) seems to be caused by a wrong cast in colour (calling tolist on a 3d array returns list[list[list[float]]], not list[list[float]]), so mypy is right to complain. The other diffs are harmless changes in error messages
  • pandas: just some harmless changes in error messages
  • xarray: re-assignment of a 1d array to/from a 2d array (src / src); this can be resolved by allowing redefinitions in their mypy config, and is nothing we can do anything about.
  • pandas-stubs: Wow; 31 unused-ignores! There are also 24 new assert-type errors, but almost all of those all involve Any in assert_type, which wouldn't cause issues for end-users, and in most cases shows that mypy is now able to better infer np.array return types. There's one new assert-type error that I don't understand though, where mypy now infers a Interval[int] where it used to infer Interval[Timestamp], even though this PR doesn't touch anything related to (the units of) datetime64, so I think we can ignore this and call it mypy oddity.

tldr; all good

@charris
charris merged commit 69102ff into numpy:main Aug 11, 2026
15 checks passed
@charris

charris commented Aug 11, 2026

Copy link
Copy Markdown
Member

Thanks Joren.

@jorenham
jorenham deleted the typing/array/improve-shape-typing-some-more branch August 11, 2026 22:18
@jorenham jorenham mentioned this pull request Aug 12, 2026
Merged
@jorenham jorenham mentioned this pull request Sep 30, 2026
Merged
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

No one assigned

Labels

41 - Static typing

Projects

None yet

Milestone

2.6.0 Release

Development

See More

Successfully merging this pull request may close these issues.

2 participants

@jorenham @charris