From 6fed3782771956f0e9689bfb803d1d24b3852842 Mon Sep 17 00:00:00 2001 From: Simon Hawkins Date: Fri, 13 Nov 2020 18:34:05 +0000 Subject: [PATCH 1/2] TYP: _concat_same_type method of EA --- pandas/core/arrays/_mixins.py | 4 ++-- pandas/core/arrays/base.py | 19 ++++++++++++++++--- pandas/core/arrays/categorical.py | 8 ++++++-- pandas/core/arrays/datetimelike.py | 6 +++++- pandas/core/arrays/interval.py | 25 +++++++++++++++++++------ pandas/core/arrays/masked.py | 6 ++++-- pandas/core/arrays/sparse/array.py | 7 +++++-- 7 files changed, 57 insertions(+), 18 deletions(-) diff --git a/pandas/core/arrays/_mixins.py b/pandas/core/arrays/_mixins.py index ddcf225d3585f..69cba1eb08098 100644 --- a/pandas/core/arrays/_mixins.py +++ b/pandas/core/arrays/_mixins.py @@ -1,4 +1,4 @@ -from typing import Any, Optional, Sequence, TypeVar +from typing import Any, Optional, Sequence, Type, TypeVar import numpy as np @@ -160,7 +160,7 @@ def unique(self: _T) -> _T: @classmethod @doc(ExtensionArray._concat_same_type) - def _concat_same_type(cls, to_concat, axis: int = 0): + def _concat_same_type(cls: Type[_T], to_concat: Sequence[_T], axis: int = 0) -> _T: dtypes = {str(x.dtype) for x in to_concat} if len(dtypes) != 1: raise ValueError("to_concat must have the same dtype (tz)", dtypes) diff --git a/pandas/core/arrays/base.py b/pandas/core/arrays/base.py index afbddc53804ac..07d0c7b4f2497 100644 --- a/pandas/core/arrays/base.py +++ b/pandas/core/arrays/base.py @@ -7,7 +7,18 @@ without warning. """ import operator -from typing import Any, Callable, Dict, Optional, Sequence, Tuple, Union, cast +from typing import ( + Any, + Callable, + Dict, + Optional, + Sequence, + Tuple, + Type, + TypeVar, + Union, + cast, +) import numpy as np @@ -37,6 +48,8 @@ _extension_array_shared_docs: Dict[str, str] = dict() +ExtensionArrayT = TypeVar("ExtensionArrayT", bound="ExtensionArray") + class ExtensionArray: """ @@ -1130,8 +1143,8 @@ def ravel(self, order="C") -> "ExtensionArray": @classmethod def _concat_same_type( - cls, to_concat: Sequence["ExtensionArray"] - ) -> "ExtensionArray": + cls: Type[ExtensionArrayT], to_concat: Sequence[ExtensionArrayT] + ) -> ExtensionArrayT: """ Concatenate multiple array of this dtype. diff --git a/pandas/core/arrays/categorical.py b/pandas/core/arrays/categorical.py index 970163df908ec..67818e6cf8fae 100644 --- a/pandas/core/arrays/categorical.py +++ b/pandas/core/arrays/categorical.py @@ -2,7 +2,7 @@ from functools import partial import operator from shutil import get_terminal_size -from typing import Dict, Hashable, List, Type, Union, cast +from typing import Dict, Hashable, List, Sequence, Type, TypeVar, Union, cast from warnings import warn import numpy as np @@ -56,6 +56,8 @@ from pandas.io.formats import console +CategoricalT = TypeVar("CategoricalT", bound="Categorical") + def _cat_compare_op(op): opname = f"__{op.__name__}__" @@ -2080,7 +2082,9 @@ def equals(self, other: object) -> bool: return False @classmethod - def _concat_same_type(self, to_concat): + def _concat_same_type( + cls: Type[CategoricalT], to_concat: Sequence[CategoricalT], axis: int = 0 + ) -> CategoricalT: from pandas.core.dtypes.concat import union_categoricals return union_categoricals(to_concat) diff --git a/pandas/core/arrays/datetimelike.py b/pandas/core/arrays/datetimelike.py index aa386e8a264cd..0ce32fcd822e0 100644 --- a/pandas/core/arrays/datetimelike.py +++ b/pandas/core/arrays/datetimelike.py @@ -375,7 +375,11 @@ def view(self, dtype=None): # ExtensionArray Interface @classmethod - def _concat_same_type(cls, to_concat, axis: int = 0): + def _concat_same_type( + cls: Type[DatetimeLikeArrayT], + to_concat: Sequence[DatetimeLikeArrayT], + axis: int = 0, + ) -> DatetimeLikeArrayT: new_obj = super()._concat_same_type(to_concat, axis) obj = to_concat[0] diff --git a/pandas/core/arrays/interval.py b/pandas/core/arrays/interval.py index 977e4abff4287..0fe05b3a6abc8 100644 --- a/pandas/core/arrays/interval.py +++ b/pandas/core/arrays/interval.py @@ -1,7 +1,7 @@ import operator from operator import le, lt import textwrap -from typing import TYPE_CHECKING, Optional, Tuple, Union, cast +from typing import TYPE_CHECKING, Optional, Sequence, Tuple, Type, TypeVar, Union, cast import numpy as np @@ -56,6 +56,8 @@ from pandas import Index from pandas.core.arrays import DatetimeArray, TimedeltaArray +IntervalArrayT = TypeVar("IntervalArrayT", bound="IntervalArray") + _interval_shared_docs = {} _shared_docs_kwargs = dict( @@ -722,7 +724,9 @@ def equals(self, other) -> bool: ) @classmethod - def _concat_same_type(cls, to_concat): + def _concat_same_type( + cls: Type[IntervalArrayT], to_concat: Sequence[IntervalArrayT] + ) -> IntervalArrayT: """ Concatenate multiple IntervalArray @@ -1468,10 +1472,19 @@ def _get_combined_data( axis=1, ) else: - left = cast(Union["DatetimeArray", "TimedeltaArray"], left) - right = cast(Union["DatetimeArray", "TimedeltaArray"], right) - combined = type(left)._concat_same_type( - [left.reshape(-1, 1), right.reshape(-1, 1)], + # error: Item "type" of "Union[Type[Index], Type[ExtensionArray]]" has + # no attribute "_concat_same_type" [union-attr] + + # error: Unexpected keyword argument "axis" for "_concat_same_type" of + # "ExtensionArray" [call-arg] + + # error: Item "Index" of "Union[Index, ExtensionArray]" has no + # attribute "reshape" [union-attr] + + # error: Item "ExtensionArray" of "Union[Index, ExtensionArray]" has no + # attribute "reshape" [union-attr] + combined = type(left)._concat_same_type( # type: ignore[union-attr,call-arg] + [left.reshape(-1, 1), right.reshape(-1, 1)], # type: ignore[union-attr] axis=1, ) return combined diff --git a/pandas/core/arrays/masked.py b/pandas/core/arrays/masked.py index 9cc4cc72e4c8e..a4b88427ceb05 100644 --- a/pandas/core/arrays/masked.py +++ b/pandas/core/arrays/masked.py @@ -1,4 +1,4 @@ -from typing import TYPE_CHECKING, Optional, Tuple, Type, TypeVar +from typing import TYPE_CHECKING, Optional, Sequence, Tuple, Type, TypeVar import numpy as np @@ -261,7 +261,9 @@ def nbytes(self) -> int: return self._data.nbytes + self._mask.nbytes @classmethod - def _concat_same_type(cls: Type[BaseMaskedArrayT], to_concat) -> BaseMaskedArrayT: + def _concat_same_type( + cls: Type[BaseMaskedArrayT], to_concat: Sequence[BaseMaskedArrayT] + ) -> BaseMaskedArrayT: data = np.concatenate([x._data for x in to_concat]) mask = np.concatenate([x._mask for x in to_concat]) return cls(data, mask) diff --git a/pandas/core/arrays/sparse/array.py b/pandas/core/arrays/sparse/array.py index d976526955ac2..77ff8560fe17a 100644 --- a/pandas/core/arrays/sparse/array.py +++ b/pandas/core/arrays/sparse/array.py @@ -4,7 +4,7 @@ from collections import abc import numbers import operator -from typing import Any, Callable, Union +from typing import Any, Callable, Sequence, Type, TypeVar, Union import warnings import numpy as np @@ -56,6 +56,7 @@ # ---------------------------------------------------------------------------- # Array +SparseArrayT = TypeVar("SparseArrayT", bound="SparseArray") _sparray_doc_kwargs = dict(klass="SparseArray") @@ -942,7 +943,9 @@ def copy(self): return self._simple_new(values, self.sp_index, self.dtype) @classmethod - def _concat_same_type(cls, to_concat): + def _concat_same_type( + cls: Type[SparseArrayT], to_concat: Sequence[SparseArrayT] + ) -> SparseArrayT: fill_value = to_concat[0].fill_value values = [] From b86f489286f51300260ec0dc41508903b8a95ba7 Mon Sep 17 00:00:00 2001 From: Simon Hawkins Date: Sat, 14 Nov 2020 10:05:03 +0000 Subject: [PATCH 2/2] post merge fixup --- pandas/core/arrays/_mixins.py | 6 +++++- pandas/core/arrays/base.py | 2 +- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/pandas/core/arrays/_mixins.py b/pandas/core/arrays/_mixins.py index 29d51a797cc57..7eaadecbd6491 100644 --- a/pandas/core/arrays/_mixins.py +++ b/pandas/core/arrays/_mixins.py @@ -170,7 +170,11 @@ def unique(self: NDArrayBackedExtensionArrayT) -> NDArrayBackedExtensionArrayT: @classmethod @doc(ExtensionArray._concat_same_type) - def _concat_same_type(cls: Type[_T], to_concat: Sequence[_T], axis: int = 0) -> _T: + def _concat_same_type( + cls: Type[NDArrayBackedExtensionArrayT], + to_concat: Sequence[NDArrayBackedExtensionArrayT], + axis: int = 0, + ) -> NDArrayBackedExtensionArrayT: dtypes = {str(x.dtype) for x in to_concat} if len(dtypes) != 1: raise ValueError("to_concat must have the same dtype (tz)", dtypes) diff --git a/pandas/core/arrays/base.py b/pandas/core/arrays/base.py index 65f76e7f004e1..0968545a6b8a4 100644 --- a/pandas/core/arrays/base.py +++ b/pandas/core/arrays/base.py @@ -23,7 +23,7 @@ import numpy as np from pandas._libs import lib -from pandas._typing import ArrayLike, Shape, TypeVar +from pandas._typing import ArrayLike, Shape from pandas.compat import set_function_name from pandas.compat.numpy import function as nv from pandas.errors import AbstractMethodError