Skip to content

Commit b00a17e

Browse files
TYP: _concat_same_type method of EA (#37817)
1 parent fee6675 commit b00a17e

File tree

7 files changed

+57
-19
lines changed

7 files changed

+57
-19
lines changed

pandas/core/arrays/_mixins.py

+6-2
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
from typing import Any, Optional, Sequence, TypeVar
1+
from typing import Any, Optional, Sequence, Type, TypeVar
22

33
import numpy as np
44

@@ -170,7 +170,11 @@ def unique(self: NDArrayBackedExtensionArrayT) -> NDArrayBackedExtensionArrayT:
170170

171171
@classmethod
172172
@doc(ExtensionArray._concat_same_type)
173-
def _concat_same_type(cls, to_concat, axis: int = 0):
173+
def _concat_same_type(
174+
cls: Type[NDArrayBackedExtensionArrayT],
175+
to_concat: Sequence[NDArrayBackedExtensionArrayT],
176+
axis: int = 0,
177+
) -> NDArrayBackedExtensionArrayT:
174178
dtypes = {str(x.dtype) for x in to_concat}
175179
if len(dtypes) != 1:
176180
raise ValueError("to_concat must have the same dtype (tz)", dtypes)

pandas/core/arrays/base.py

+15-4
Original file line numberDiff line numberDiff line change
@@ -7,12 +7,23 @@
77
without warning.
88
"""
99
import operator
10-
from typing import Any, Callable, Dict, Optional, Sequence, Tuple, Union, cast
10+
from typing import (
11+
Any,
12+
Callable,
13+
Dict,
14+
Optional,
15+
Sequence,
16+
Tuple,
17+
Type,
18+
TypeVar,
19+
Union,
20+
cast,
21+
)
1122

1223
import numpy as np
1324

1425
from pandas._libs import lib
15-
from pandas._typing import ArrayLike, Shape, TypeVar
26+
from pandas._typing import ArrayLike, Shape
1627
from pandas.compat import set_function_name
1728
from pandas.compat.numpy import function as nv
1829
from pandas.errors import AbstractMethodError
@@ -1132,8 +1143,8 @@ def ravel(self, order="C") -> "ExtensionArray":
11321143

11331144
@classmethod
11341145
def _concat_same_type(
1135-
cls, to_concat: Sequence["ExtensionArray"]
1136-
) -> "ExtensionArray":
1146+
cls: Type[ExtensionArrayT], to_concat: Sequence[ExtensionArrayT]
1147+
) -> ExtensionArrayT:
11371148
"""
11381149
Concatenate multiple array of this dtype.
11391150

pandas/core/arrays/categorical.py

+6-2
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22
from functools import partial
33
import operator
44
from shutil import get_terminal_size
5-
from typing import Dict, Hashable, List, Type, Union, cast
5+
from typing import Dict, Hashable, List, Sequence, Type, TypeVar, Union, cast
66
from warnings import warn
77

88
import numpy as np
@@ -56,6 +56,8 @@
5656

5757
from pandas.io.formats import console
5858

59+
CategoricalT = TypeVar("CategoricalT", bound="Categorical")
60+
5961

6062
def _cat_compare_op(op):
6163
opname = f"__{op.__name__}__"
@@ -2080,7 +2082,9 @@ def equals(self, other: object) -> bool:
20802082
return False
20812083

20822084
@classmethod
2083-
def _concat_same_type(self, to_concat):
2085+
def _concat_same_type(
2086+
cls: Type[CategoricalT], to_concat: Sequence[CategoricalT], axis: int = 0
2087+
) -> CategoricalT:
20842088
from pandas.core.dtypes.concat import union_categoricals
20852089

20862090
return union_categoricals(to_concat)

pandas/core/arrays/datetimelike.py

+5-1
Original file line numberDiff line numberDiff line change
@@ -375,7 +375,11 @@ def view(self, dtype=None):
375375
# ExtensionArray Interface
376376

377377
@classmethod
378-
def _concat_same_type(cls, to_concat, axis: int = 0):
378+
def _concat_same_type(
379+
cls: Type[DatetimeLikeArrayT],
380+
to_concat: Sequence[DatetimeLikeArrayT],
381+
axis: int = 0,
382+
) -> DatetimeLikeArrayT:
379383
new_obj = super()._concat_same_type(to_concat, axis)
380384

381385
obj = to_concat[0]

pandas/core/arrays/interval.py

+17-6
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
import operator
22
from operator import le, lt
33
import textwrap
4-
from typing import TYPE_CHECKING, Optional, Tuple, TypeVar, Union, cast
4+
from typing import TYPE_CHECKING, Optional, Sequence, Tuple, Type, TypeVar, Union, cast
55

66
import numpy as np
77

@@ -724,7 +724,9 @@ def equals(self, other) -> bool:
724724
)
725725

726726
@classmethod
727-
def _concat_same_type(cls, to_concat):
727+
def _concat_same_type(
728+
cls: Type[IntervalArrayT], to_concat: Sequence[IntervalArrayT]
729+
) -> IntervalArrayT:
728730
"""
729731
Concatenate multiple IntervalArray
730732
@@ -1470,10 +1472,19 @@ def _get_combined_data(
14701472
axis=1,
14711473
)
14721474
else:
1473-
left = cast(Union["DatetimeArray", "TimedeltaArray"], left)
1474-
right = cast(Union["DatetimeArray", "TimedeltaArray"], right)
1475-
combined = type(left)._concat_same_type(
1476-
[left.reshape(-1, 1), right.reshape(-1, 1)],
1475+
# error: Item "type" of "Union[Type[Index], Type[ExtensionArray]]" has
1476+
# no attribute "_concat_same_type" [union-attr]
1477+
1478+
# error: Unexpected keyword argument "axis" for "_concat_same_type" of
1479+
# "ExtensionArray" [call-arg]
1480+
1481+
# error: Item "Index" of "Union[Index, ExtensionArray]" has no
1482+
# attribute "reshape" [union-attr]
1483+
1484+
# error: Item "ExtensionArray" of "Union[Index, ExtensionArray]" has no
1485+
# attribute "reshape" [union-attr]
1486+
combined = type(left)._concat_same_type( # type: ignore[union-attr,call-arg]
1487+
[left.reshape(-1, 1), right.reshape(-1, 1)], # type: ignore[union-attr]
14771488
axis=1,
14781489
)
14791490
return combined

pandas/core/arrays/masked.py

+4-2
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
from typing import TYPE_CHECKING, Optional, Tuple, Type, TypeVar
1+
from typing import TYPE_CHECKING, Optional, Sequence, Tuple, Type, TypeVar
22

33
import numpy as np
44

@@ -261,7 +261,9 @@ def nbytes(self) -> int:
261261
return self._data.nbytes + self._mask.nbytes
262262

263263
@classmethod
264-
def _concat_same_type(cls: Type[BaseMaskedArrayT], to_concat) -> BaseMaskedArrayT:
264+
def _concat_same_type(
265+
cls: Type[BaseMaskedArrayT], to_concat: Sequence[BaseMaskedArrayT]
266+
) -> BaseMaskedArrayT:
265267
data = np.concatenate([x._data for x in to_concat])
266268
mask = np.concatenate([x._mask for x in to_concat])
267269
return cls(data, mask)

pandas/core/arrays/sparse/array.py

+4-2
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@
44
from collections import abc
55
import numbers
66
import operator
7-
from typing import Any, Callable, Type, TypeVar, Union
7+
from typing import Any, Callable, Sequence, Type, TypeVar, Union
88
import warnings
99

1010
import numpy as np
@@ -946,7 +946,9 @@ def copy(self: SparseArrayT) -> SparseArrayT:
946946
return self._simple_new(values, self.sp_index, self.dtype)
947947

948948
@classmethod
949-
def _concat_same_type(cls, to_concat):
949+
def _concat_same_type(
950+
cls: Type[SparseArrayT], to_concat: Sequence[SparseArrayT]
951+
) -> SparseArrayT:
950952
fill_value = to_concat[0].fill_value
951953

952954
values = []

0 commit comments

Comments
 (0)