Merge pull request #21106 from BvB93/concat_backport

TYP,MAINT: Explicitly allow sequences of array-likes in `np.concatenate`

numpy/core/multiarray.pyi

@@ -85,6 +85,8 @@ from numpy.typing import (
     _TD64Like_co,
 )
 
+_T_co = TypeVar("_T_co", covariant=True)
+_T_contra = TypeVar("_T_contra", contravariant=True)
 _SCT = TypeVar("_SCT", bound=generic)
 _ArrayType = TypeVar("_ArrayType", bound=NDArray[Any])
 
@@ -121,6 +123,10 @@ _RollKind = L[  # `raise` is deliberately excluded
     "modifiedpreceding",
 ]
 
+class _SupportsLenAndGetItem(Protocol[_T_contra, _T_co]):
+    def __len__(self) -> int: ...
+    def __getitem__(self, __key: _T_contra) -> _T_co: ...
+
 __all__: List[str]
 
 ALLOW_THREADS: Final[int]  # 0 or 1 (system-specific)
@@ -308,6 +314,7 @@ def ravel_multi_index(
     order: _OrderCF = ...,
 ) -> NDArray[intp]: ...
 
+# NOTE: Allow any sequence of array-like objects
 @overload
 def concatenate(  # type: ignore[misc]
     arrays: _ArrayLike[_SCT],
@@ -320,7 +327,7 @@ def concatenate(  # type: ignore[misc]
 ) -> NDArray[_SCT]: ...
 @overload
 def concatenate(  # type: ignore[misc]
-    arrays: ArrayLike,
+    arrays: _SupportsLenAndGetItem[int, ArrayLike],
     /,
     axis: Optional[SupportsIndex] = ...,
     out: None = ...,
@@ -330,7 +337,7 @@ def concatenate(  # type: ignore[misc]
 ) -> NDArray[Any]: ...
 @overload
 def concatenate(  # type: ignore[misc]
-    arrays: ArrayLike,
+    arrays: _SupportsLenAndGetItem[int, ArrayLike],
     /,
     axis: Optional[SupportsIndex] = ...,
     out: None = ...,
@@ -340,7 +347,7 @@ def concatenate(  # type: ignore[misc]
 ) -> NDArray[_SCT]: ...
 @overload
 def concatenate(  # type: ignore[misc]
-    arrays: ArrayLike,
+    arrays: _SupportsLenAndGetItem[int, ArrayLike],
     /,
     axis: Optional[SupportsIndex] = ...,
     out: None = ...,
@@ -350,7 +357,7 @@ def concatenate(  # type: ignore[misc]
 ) -> NDArray[Any]: ...
 @overload
 def concatenate(
-    arrays: ArrayLike,
+    arrays: _SupportsLenAndGetItem[int, ArrayLike],
     /,
     axis: Optional[SupportsIndex] = ...,
     out: _ArrayType = ...,

numpy/typing/tests/data/fail/false_positives.pyi

@@ -0,0 +1,11 @@
+import numpy as np
+import numpy.typing as npt
+
+AR_f8: npt.NDArray[np.float64]
+
+# NOTE: Mypy bug presumably due to the special-casing of heterogeneous tuples;
+# xref numpy/numpy#20901
+#
+# The expected output should be no different than, e.g., when using a
+# list instead of a tuple
+np.concatenate(([1], AR_f8))  # E: Argument 1 to "concatenate" has incompatible type

numpy/typing/tests/data/reveal/array_constructors.pyi

@@ -38,6 +38,11 @@ reveal_type(np.empty([1, 5, 6], dtype=np.int64))  # E: ndarray[Any, dtype[{int64
 reveal_type(np.empty([1, 5, 6], dtype='c16'))  # E: ndarray[Any, dtype[Any]]
 
 reveal_type(np.concatenate(A))  # E: ndarray[Any, dtype[{float64}]]
+reveal_type(np.concatenate([A, A]))  # E: Any
+reveal_type(np.concatenate([[1], A]))  # E: ndarray[Any, dtype[Any]]
+reveal_type(np.concatenate([[1], [1]]))  # E: ndarray[Any, dtype[Any]]
+reveal_type(np.concatenate((A, A)))  # E: ndarray[Any, dtype[{float64}]]
+reveal_type(np.concatenate(([1], [1])))  # E: ndarray[Any, dtype[Any]]
 reveal_type(np.concatenate([1, 1.0]))  # E: ndarray[Any, dtype[Any]]
 reveal_type(np.concatenate(A, dtype=np.int64))  # E: ndarray[Any, dtype[{int64}]]
 reveal_type(np.concatenate(A, dtype='c16'))  # E: ndarray[Any, dtype[Any]]

numpy/typing/tests/data/reveal/false_positives.pyi

@@ -0,0 +1,10 @@
+from typing import Any
+import numpy.typing as npt
+
+AR_Any: npt.NDArray[Any]
+
+# Mypy bug where overload ambiguity is ignored for `Any`-parametrized types;
+# xref numpy/numpy#20099 and python/mypy#11347
+#
+# The expected output would be something akin to `ndarray[Any, dtype[Any]]`
+reveal_type(AR_Any + 2)  # E: ndarray[Any, dtype[signedinteger[Any]]]