TYP: concat_compat (#52315)

Author: jbrockmendel

pandas/core/algorithms.py

@@ -1666,7 +1666,11 @@ def union_with_duplicates(
             lvals = lvals._values
         if isinstance(rvals, ABCIndex):
             rvals = rvals._values
-        unique_vals = unique(concat_compat([lvals, rvals]))
+        # error: List item 0 has incompatible type "Union[ExtensionArray,
+        # ndarray[Any, Any], Index]"; expected "Union[ExtensionArray,
+        # ndarray[Any, Any]]"
+        combined = concat_compat([lvals, rvals])  # type: ignore[list-item]
+        unique_vals = unique(combined)
         unique_vals = ensure_wrapped_if_datetimelike(unique_vals)
     repeats = final_count.reindex(unique_vals).values
     return np.repeat(unique_vals, repeats)

pandas/core/dtypes/cast.py

@@ -10,6 +10,7 @@
     TYPE_CHECKING,
     Any,
     Literal,
+    Sequence,
     Sized,
     TypeVar,
     cast,
@@ -1317,7 +1318,7 @@ def find_result_type(left: ArrayLike, right: Any) -> DtypeObj:
 
 
 def common_dtype_categorical_compat(
-    objs: list[Index | ArrayLike], dtype: DtypeObj
+    objs: Sequence[Index | ArrayLike], dtype: DtypeObj
 ) -> DtypeObj:
     """
     Update the result of find_common_type to account for NAs in a Categorical.

pandas/core/dtypes/concat.py

@@ -3,7 +3,11 @@
 """
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
+from typing import (
+    TYPE_CHECKING,
+    Sequence,
+    cast,
+)
 
 import numpy as np
 
@@ -24,12 +28,20 @@
 )
 
 if TYPE_CHECKING:
-    from pandas._typing import AxisInt
+    from pandas._typing import (
+        ArrayLike,
+        AxisInt,
+    )
 
-    from pandas.core.arrays import Categorical
+    from pandas.core.arrays import (
+        Categorical,
+        ExtensionArray,
+    )
 
 
-def concat_compat(to_concat, axis: AxisInt = 0, ea_compat_axis: bool = False):
+def concat_compat(
+    to_concat: Sequence[ArrayLike], axis: AxisInt = 0, ea_compat_axis: bool = False
+) -> ArrayLike:
     """
     provide concatenation of an array of arrays each of which is a single
     'normalized' dtypes (in that for example, if it's object, then it is a
@@ -38,7 +50,7 @@ def concat_compat(to_concat, axis: AxisInt = 0, ea_compat_axis: bool = False):
 
     Parameters
     ----------
-    to_concat : array of arrays
+    to_concat : sequence of arrays
     axis : axis to provide concatenation
     ea_compat_axis : bool, default False
         For ExtensionArray compat, behave as if axis == 1 when determining
@@ -93,10 +105,12 @@ def is_nonempty(x) -> bool:
 
         if isinstance(to_concat[0], ABCExtensionArray):
             # TODO: what about EA-backed Index?
+            to_concat_eas = cast("Sequence[ExtensionArray]", to_concat)
             cls = type(to_concat[0])
-            return cls._concat_same_type(to_concat)
+            return cls._concat_same_type(to_concat_eas)
         else:
-            return np.concatenate(to_concat)
+            to_concat_arrs = cast("Sequence[np.ndarray]", to_concat)
+            return np.concatenate(to_concat_arrs)
 
     elif all_empty:
         # we have all empties, but may need to coerce the result dtype to
@@ -111,7 +125,10 @@ def is_nonempty(x) -> bool:
                 to_concat = [x.astype("object") for x in to_concat]
                 kinds = {"o"}
 
-    result = np.concatenate(to_concat, axis=axis)
+    # error: Argument 1 to "concatenate" has incompatible type
+    # "Sequence[Union[ExtensionArray, ndarray[Any, Any]]]"; expected
+    # "Union[_SupportsArray[dtype[Any]], _NestedSequence[_SupportsArray[dtype[Any]]]]"
+    result: np.ndarray = np.concatenate(to_concat, axis=axis)  # type: ignore[arg-type]
     if "b" in kinds and result.dtype.kind in ["i", "u", "f"]:
         # GH#39817 cast to object instead of casting bools to numeric
         result = result.astype(object, copy=False)
@@ -283,21 +300,21 @@ def _maybe_unwrap(x):
     return Categorical(new_codes, categories=categories, ordered=ordered, fastpath=True)
 
 
-def _concatenate_2d(to_concat, axis: AxisInt):
+def _concatenate_2d(to_concat: Sequence[np.ndarray], axis: AxisInt) -> np.ndarray:
     # coerce to 2d if needed & concatenate
     if axis == 1:
         to_concat = [np.atleast_2d(x) for x in to_concat]
     return np.concatenate(to_concat, axis=axis)
 
 
-def _concat_datetime(to_concat, axis: AxisInt = 0):
+def _concat_datetime(to_concat: Sequence[ArrayLike], axis: AxisInt = 0) -> ArrayLike:
     """
     provide concatenation of an datetimelike array of arrays each of which is a
     single M8[ns], datetime64[ns, tz] or m8[ns] dtype
 
     Parameters
     ----------
-    to_concat : array of arrays
+    to_concat : sequence of arrays
     axis : axis to provide concatenation
 
     Returns
@@ -316,5 +333,10 @@ def _concat_datetime(to_concat, axis: AxisInt = 0):
         #  in Timestamp/Timedelta
         return _concatenate_2d([x.astype(object) for x in to_concat], axis=axis)
 
-    result = type(to_concat[0])._concat_same_type(to_concat, axis=axis)
+    # error: Unexpected keyword argument "axis" for "_concat_same_type" of
+    # "ExtensionArray"
+    to_concat_eas = cast("list[ExtensionArray]", to_concat)
+    result = type(to_concat_eas[0])._concat_same_type(  # type: ignore[call-arg]
+        to_concat_eas, axis=axis
+    )
     return result