From 87adf67f2700f0e2274e7a696e2c4261d5090439 Mon Sep 17 00:00:00 2001 From: Brock Date: Tue, 20 Jun 2023 11:14:17 -0700 Subject: [PATCH] TYP: reshape.merge --- pandas/core/groupby/generic.py | 7 ++++++- pandas/core/reshape/merge.py | 29 ++++++++++++++--------------- 2 files changed, 20 insertions(+), 16 deletions(-) diff --git a/pandas/core/groupby/generic.py b/pandas/core/groupby/generic.py index cecb9a84c62dd..43854c5849481 100644 --- a/pandas/core/groupby/generic.py +++ b/pandas/core/groupby/generic.py @@ -829,7 +829,12 @@ def value_counts( right = [diff.cumsum() - 1, codes[-1]] - _, idx = get_join_indexers(left, right, sort=False, how="left") + # error: Argument 1 to "get_join_indexers" has incompatible type + # "List[ndarray[Any, Any]]"; expected "List[Union[Union[ExtensionArray, + # ndarray[Any, Any]], Index, Series]] + _, idx = get_join_indexers( + left, right, sort=False, how="left" # type: ignore[arg-type] + ) out = np.where(idx != -1, out[idx], 0) if sort: diff --git a/pandas/core/reshape/merge.py b/pandas/core/reshape/merge.py index e8420983c65e7..cbcdc3813c24d 100644 --- a/pandas/core/reshape/merge.py +++ b/pandas/core/reshape/merge.py @@ -54,7 +54,6 @@ ensure_object, is_bool, is_bool_dtype, - is_extension_array_dtype, is_float_dtype, is_integer, is_integer_dtype, @@ -1611,8 +1610,8 @@ def _validate(self, validate: str) -> None: def get_join_indexers( - left_keys, - right_keys, + left_keys: list[AnyArrayLike], + right_keys: list[AnyArrayLike], sort: bool = False, how: MergeHow | Literal["asof"] = "inner", **kwargs, @@ -1621,8 +1620,8 @@ def get_join_indexers( Parameters ---------- - left_keys : ndarray, Index, Series - right_keys : ndarray, Index, Series + left_keys : list[ndarray, ExtensionArray, Index, Series] + right_keys : list[ndarray, ExtensionArray, Index, Series] sort : bool, default False how : {'inner', 'outer', 'left', 'right'}, default 'inner' @@ -2062,11 +2061,11 @@ def _get_merge_keys( def _get_join_indexers(self) -> tuple[npt.NDArray[np.intp], npt.NDArray[np.intp]]: """return the join indexers""" - def flip(xs) -> np.ndarray: + def flip(xs: list[AnyArrayLike]) -> np.ndarray: """unlike np.transpose, this returns an array of tuples""" - def injection(obj): - if not is_extension_array_dtype(obj): + def injection(obj: AnyArrayLike): + if not isinstance(obj.dtype, ExtensionDtype): # ndarray return obj obj = extract_array(obj) @@ -2213,7 +2212,7 @@ def injection(obj): def _get_multiindex_indexer( - join_keys, index: MultiIndex, sort: bool + join_keys: list[AnyArrayLike], index: MultiIndex, sort: bool ) -> tuple[npt.NDArray[np.intp], npt.NDArray[np.intp]]: # left & right join labels and num. of levels at each location mapped = ( @@ -2250,7 +2249,7 @@ def _get_multiindex_indexer( def _get_single_indexer( - join_key, index: Index, sort: bool = False + join_key: AnyArrayLike, index: Index, sort: bool = False ) -> tuple[npt.NDArray[np.intp], npt.NDArray[np.intp]]: left_key, right_key, count = _factorize_keys(join_key, index._values, sort=sort) @@ -2295,7 +2294,7 @@ def _get_no_sort_one_missing_indexer( def _left_join_on_index( - left_ax: Index, right_ax: Index, join_keys, sort: bool = False + left_ax: Index, right_ax: Index, join_keys: list[AnyArrayLike], sort: bool = False ) -> tuple[Index, npt.NDArray[np.intp] | None, npt.NDArray[np.intp]]: if isinstance(right_ax, MultiIndex): left_indexer, right_indexer = _get_multiindex_indexer( @@ -2316,8 +2315,8 @@ def _left_join_on_index( def _factorize_keys( - lk: ArrayLike, - rk: ArrayLike, + lk: AnyArrayLike, + rk: AnyArrayLike, sort: bool = True, how: MergeHow | Literal["asof"] = "inner", ) -> tuple[npt.NDArray[np.intp], npt.NDArray[np.intp], int]: @@ -2328,9 +2327,9 @@ def _factorize_keys( Parameters ---------- - lk : array-like + lk : ndarray, ExtensionArray, Index, or Series Left key. - rk : array-like + rk : ndarray, ExtensionArray, Index, or Series Right key. sort : bool, defaults to True If True, the encoding is done such that the unique elements in the