diff --git a/pandas/_testing/asserters.py b/pandas/_testing/asserters.py index 2adc70438cce7..38c4ac421abb3 100644 --- a/pandas/_testing/asserters.py +++ b/pandas/_testing/asserters.py @@ -382,9 +382,13 @@ def _get_ilevel_values(index, level): # skip exact index checking when `check_categorical` is False if check_exact and check_categorical: if not left.equals(right): - diff = ( - np.sum((left._values != right._values).astype(int)) * 100.0 / len(left) - ) + # error: Value of type variable "_Number" of "sum" cannot be + # "Union[ExtensionArray, ndarray, Any]" + thesum = np.sum( + (left._values != right._values).astype(int) + ) # type: ignore[type-var] + # error: Unsupported operand types for * ("ExtensionArray" and "float") + diff = thesum * 100.0 / len(left) # type: ignore[operator] msg = f"{obj} values are different ({np.round(diff, 5)} %)" raise_assert_detail(obj, msg, left, right) else: diff --git a/pandas/_typing.py b/pandas/_typing.py index 3e584774e539a..9ce3adda926cf 100644 --- a/pandas/_typing.py +++ b/pandas/_typing.py @@ -76,11 +76,17 @@ ArrayLike = Union["ExtensionArray", np.ndarray] AnyArrayLike = Union[ArrayLike, "Index", "Series"] - +AnySequenceLike = Union[ + "ExtensionArray", + "Index", + "Series", + Sequence[Any], + np.ndarray, +] # scalars PythonScalar = Union[str, int, float, bool] -DatetimeLikeScalar = TypeVar("DatetimeLikeScalar", "Period", "Timestamp", "Timedelta") +DatetimeLikeScalar = Union["Period", "Timestamp", "Timedelta"] PandasScalar = Union["Period", "Timestamp", "Timedelta", "Interval"] Scalar = Union[PythonScalar, PandasScalar] diff --git a/pandas/core/algorithms.py b/pandas/core/algorithms.py index a888bfabd6f80..ad4bb8bcab3a5 100644 --- a/pandas/core/algorithms.py +++ b/pandas/core/algorithms.py @@ -30,6 +30,7 @@ ) from pandas._typing import ( AnyArrayLike, + AnySequenceLike, ArrayLike, DtypeObj, FrameOrSeriesUnion, @@ -446,7 +447,7 @@ def unique(values): unique1d = unique -def isin(comps: AnyArrayLike, values: AnyArrayLike) -> np.ndarray: +def isin(comps: AnySequenceLike, values: AnySequenceLike) -> np.ndarray: """ Compute the isin boolean array. @@ -482,9 +483,11 @@ def isin(comps: AnyArrayLike, values: AnyArrayLike) -> np.ndarray: comps = _ensure_arraylike(comps) comps = extract_array(comps, extract_numpy=True) if is_extension_array_dtype(comps.dtype): - # error: Incompatible return value type (got "Series", expected "ndarray") - # error: Item "ndarray" of "Union[Any, ndarray]" has no attribute "isin" - return comps.isin(values) # type: ignore[return-value,union-attr] + # error: Argument 1 to "isin" of "ExtensionArray" has incompatible type + # "Union[Any, ExtensionArray, ndarray]"; expected "Sequence[Any]" + # error: Item "ndarray" of "Union[Any, ExtensionArray, ndarray]" has no + # attribute "isin" + return comps.isin(values) # type: ignore[arg-type, union-attr] elif needs_i8_conversion(comps.dtype): # Dispatch to DatetimeLikeArrayMixin.isin diff --git a/pandas/core/array_algos/putmask.py b/pandas/core/array_algos/putmask.py index 3daf1b3ae3902..0666112cec33d 100644 --- a/pandas/core/array_algos/putmask.py +++ b/pandas/core/array_algos/putmask.py @@ -191,7 +191,7 @@ def extract_bool_array(mask: ArrayLike) -> np.ndarray: # We could have BooleanArray, Sparse[bool], ... # Except for BooleanArray, this is equivalent to just # np.asarray(mask, dtype=bool) - mask = mask.to_numpy(dtype=bool, na_value=False) + mask = mask.to_numpy(dtype=np.dtype(bool), na_value=False) mask = np.asarray(mask, dtype=bool) return mask diff --git a/pandas/core/arrays/_mixins.py b/pandas/core/arrays/_mixins.py index f56f9b2735f88..255e5f43335e4 100644 --- a/pandas/core/arrays/_mixins.py +++ b/pandas/core/arrays/_mixins.py @@ -3,11 +3,13 @@ from functools import wraps from typing import ( Any, + List, Optional, Sequence, Type, TypeVar, Union, + overload, ) import numpy as np @@ -253,8 +255,19 @@ def __setitem__(self, key, value): def _validate_setitem_value(self, value): return value + @overload + def __getitem__(self: NDArrayBackedExtensionArrayT, key: int) -> Any: + ... + + @overload def __getitem__( - self: NDArrayBackedExtensionArrayT, key: Union[int, slice, np.ndarray] + self: NDArrayBackedExtensionArrayT, key: Union[slice, np.ndarray, List[Any]] + ) -> NDArrayBackedExtensionArrayT: + ... + + def __getitem__( + self: NDArrayBackedExtensionArrayT, + key: Union[int, slice, np.ndarray, List[Any]], ) -> Union[NDArrayBackedExtensionArrayT, Any]: if lib.is_integer(key): # fast-path diff --git a/pandas/core/arrays/base.py b/pandas/core/arrays/base.py index 68909e30650c7..4fc568d04cb44 100644 --- a/pandas/core/arrays/base.py +++ b/pandas/core/arrays/base.py @@ -14,6 +14,8 @@ Any, Callable, Dict, + Iterator, + List, Optional, Sequence, Tuple, @@ -21,6 +23,7 @@ TypeVar, Union, cast, + overload, ) import numpy as np @@ -29,6 +32,7 @@ from pandas._typing import ( ArrayLike, Dtype, + NpDtype, Shape, ) from pandas.compat import set_function_name @@ -73,6 +77,7 @@ ) if TYPE_CHECKING: + from typing import Literal class ExtensionArraySupportsAnyAll("ExtensionArray"): def any(self, *, skipna: bool = True) -> bool: @@ -293,8 +298,16 @@ def _from_factorized(cls, values, original): # Must be a Sequence # ------------------------------------------------------------------------ + @overload + def __getitem__(self, item: int) -> Any: + ... + + @overload + def __getitem__(self, item: Union[slice, np.ndarray, List[Any]]) -> ExtensionArray: + ... + def __getitem__( - self, item: Union[int, slice, np.ndarray] + self, item: Union[int, slice, np.ndarray, List[Any]] ) -> Union[ExtensionArray, Any]: """ Select a subset of self. @@ -381,7 +394,7 @@ def __len__(self) -> int: """ raise AbstractMethodError(self) - def __iter__(self): + def __iter__(self) -> Iterator[Any]: """ Iterate over elements of the array. """ @@ -391,7 +404,7 @@ def __iter__(self): for i in range(len(self)): yield self[i] - def __contains__(self, item) -> Union[bool, np.bool_]: + def __contains__(self, item: Any) -> Union[bool, np.bool_]: """ Return for `item in self`. """ @@ -430,9 +443,10 @@ def __ne__(self, other: Any) -> ArrayLike: # type: ignore[override] def to_numpy( self, - dtype: Optional[Dtype] = None, + dtype: Optional[NpDtype] = None, copy: bool = False, - na_value=lib.no_default, + na_value: Optional[Any] = lib.no_default, + **kwargs: Any, ) -> np.ndarray: """ Convert to a NumPy ndarray. @@ -459,12 +473,7 @@ def to_numpy( ------- numpy.ndarray """ - # error: Argument "dtype" to "asarray" has incompatible type - # "Union[ExtensionDtype, str, dtype[Any], Type[str], Type[float], Type[int], - # Type[complex], Type[bool], Type[object], None]"; expected "Union[dtype[Any], - # None, type, _SupportsDType, str, Union[Tuple[Any, int], Tuple[Any, Union[int, - # Sequence[int]]], List[Any], _DTypeDict, Tuple[Any, Any]]]" - result = np.asarray(self, dtype=dtype) # type: ignore[arg-type] + result = np.asarray(self, dtype=dtype) if copy or na_value is not lib.no_default: result = result.copy() if na_value is not lib.no_default: @@ -516,8 +525,15 @@ def nbytes(self) -> int: # ------------------------------------------------------------------------ # Additional Methods # ------------------------------------------------------------------------ + @overload + def astype(self, dtype: Type[str], copy: bool = True) -> np.ndarray: + ... + + @overload + def astype(self, dtype: Dtype, copy: bool = True) -> ArrayLike: + ... - def astype(self, dtype, copy=True): + def astype(self, dtype: Dtype, copy: bool = True) -> ArrayLike: """ Cast to a NumPy array with 'dtype'. @@ -551,7 +567,7 @@ def astype(self, dtype, copy=True): ): # allow conversion to StringArrays return dtype.construct_array_type()._from_sequence(self, copy=False) - return np.array(self, dtype=dtype, copy=copy) + return np.array(self, dtype=cast(NpDtype, dtype), copy=copy) def isna(self) -> Union[np.ndarray, ExtensionArraySupportsAnyAll]: """ @@ -597,8 +613,8 @@ def argsort( ascending: bool = True, kind: str = "quicksort", na_position: str = "last", - *args, - **kwargs, + *args: Any, + **kwargs: Any, ) -> np.ndarray: """ Return the indices that would sort this array. @@ -686,7 +702,12 @@ def argmax(self, skipna: bool = True) -> int: raise NotImplementedError return nargminmax(self, "argmax") - def fillna(self, value=None, method=None, limit=None): + def fillna( + self, + value: Optional[Union[Any, ArrayLike]] = None, + method: Optional[Literal["backfill", "bfill", "ffill", "pad"]] = None, + limit: Optional[int] = None, + ) -> ExtensionArray: """ Fill NA/NaN values using the specified method. @@ -735,7 +756,7 @@ def fillna(self, value=None, method=None, limit=None): new_values = self.copy() return new_values - def dropna(self): + def dropna(self) -> ExtensionArrayT: """ Return ExtensionArray without NA values. @@ -800,7 +821,7 @@ def shift(self, periods: int = 1, fill_value: object = None) -> ExtensionArray: b = empty return self._concat_same_type([a, b]) - def unique(self): + def unique(self) -> ExtensionArray: """ Compute the ExtensionArray of unique values. @@ -811,7 +832,12 @@ def unique(self): uniques = unique(self.astype(object)) return self._from_sequence(uniques, dtype=self.dtype) - def searchsorted(self, value, side="left", sorter=None): + def searchsorted( + self, + value: Sequence[Any], + side: Literal["left", "right"] = "left", + sorter: Optional[Sequence[Any]] = None, + ) -> np.ndarray: """ Find indices where elements should be inserted to maintain order. @@ -856,7 +882,7 @@ def searchsorted(self, value, side="left", sorter=None): # 1. Values outside the range of the `data_for_sorting` fixture # 2. Values between the values in the `data_for_sorting` fixture # 3. Missing values. - arr = self.astype(object) + arr = cast(np.ndarray, self.astype(object)) return arr.searchsorted(value, side=side, sorter=sorter) def equals(self, other: object) -> bool: @@ -893,7 +919,7 @@ def equals(self, other: object) -> bool: equal_na = self.isna() & other.isna() # type: ignore[operator] return bool((equal_values | equal_na).all()) - def isin(self, values) -> np.ndarray: + def isin(self, values: Sequence[Any]) -> np.ndarray: """ Pointwise comparison for set containment in the given values. @@ -907,7 +933,7 @@ def isin(self, values) -> np.ndarray: ------- np.ndarray[bool] """ - return isin(np.asarray(self), values) + return isin(self.astype(object), values) def _values_for_factorize(self) -> Tuple[np.ndarray, Any]: """ @@ -931,7 +957,7 @@ def _values_for_factorize(self) -> Tuple[np.ndarray, Any]: The values returned by this method are also used in :func:`pandas.util.hash_pandas_object`. """ - return self.astype(object), np.nan + return cast(np.ndarray, self.astype(object)), np.nan def factorize(self, na_sentinel: int = -1) -> Tuple[np.ndarray, ExtensionArray]: """ @@ -1029,7 +1055,11 @@ def factorize(self, na_sentinel: int = -1) -> Tuple[np.ndarray, ExtensionArray]: @Substitution(klass="ExtensionArray") @Appender(_extension_array_shared_docs["repeat"]) - def repeat(self, repeats, axis=None): + def repeat( + self, + repeats: Union[int, Sequence[int]], + axis: Literal[None] = None, + ) -> ExtensionArray: nv.validate_repeat((), {"axis": axis}) ind = np.arange(len(self)).repeat(repeats) return self.take(ind) @@ -1039,12 +1069,12 @@ def repeat(self, repeats, axis=None): # ------------------------------------------------------------------------ def take( - self: ExtensionArrayT, + self, indices: Sequence[int], *, allow_fill: bool = False, fill_value: Any = None, - ) -> ExtensionArrayT: + ) -> ExtensionArray: """ Take elements from an array. @@ -1133,7 +1163,7 @@ def take(self, indices, allow_fill=False, fill_value=None): # pandas.api.extensions.take raise AbstractMethodError(self) - def copy(self: ExtensionArrayT) -> ExtensionArrayT: + def copy(self) -> ExtensionArray: """ Return a copy of the array. @@ -1213,7 +1243,7 @@ def _formatter(self, boxed: bool = False) -> Callable[[Any], Optional[str]]: # Reshaping # ------------------------------------------------------------------------ - def transpose(self, *axes) -> ExtensionArray: + def transpose(self, *axes: int) -> ExtensionArray: """ Return a transposed view on this array. @@ -1226,7 +1256,9 @@ def transpose(self, *axes) -> ExtensionArray: def T(self) -> ExtensionArray: return self.transpose() - def ravel(self, order="C") -> ExtensionArray: + def ravel( + self, order: Optional[Literal["C", "F", "A", "K"]] = "C" + ) -> ExtensionArray: """ Return a flattened view on this array. @@ -1300,13 +1332,13 @@ def _reduce(self, name: str, *, skipna: bool = True, **kwargs): """ raise TypeError(f"cannot perform {name} with type {self.dtype}") - def __hash__(self): + def __hash__(self) -> int: raise TypeError(f"unhashable type: {repr(type(self).__name__)}") # ------------------------------------------------------------------------ # Non-Optimized Default Methods - def delete(self: ExtensionArrayT, loc) -> ExtensionArrayT: + def delete(self, loc: Union[int, Sequence[int]]) -> ExtensionArray: indexer = np.delete(np.arange(len(self)), loc) return self.take(indexer) diff --git a/pandas/core/arrays/boolean.py b/pandas/core/arrays/boolean.py index 4258279e37551..f2cb696fab07f 100644 --- a/pandas/core/arrays/boolean.py +++ b/pandas/core/arrays/boolean.py @@ -8,6 +8,7 @@ Tuple, Type, Union, + overload, ) import warnings @@ -20,6 +21,7 @@ from pandas._typing import ( ArrayLike, Dtype, + DtypeArg, ) from pandas.compat.numpy import function as nv @@ -296,7 +298,7 @@ def dtype(self) -> BooleanDtype: @classmethod def _from_sequence( - cls, scalars, *, dtype: Optional[Dtype] = None, copy: bool = False + cls, scalars, *, dtype: Optional[DtypeArg] = None, copy: bool = False ) -> BooleanArray: if dtype: assert dtype == "boolean" @@ -379,7 +381,15 @@ def reconstruct(x): def _coerce_to_array(self, value) -> Tuple[np.ndarray, np.ndarray]: return coerce_to_array(value) - def astype(self, dtype, copy: bool = True) -> ArrayLike: + @overload + def astype(self, dtype: Type[str], copy: bool = True) -> np.ndarray: + ... + + @overload + def astype(self, dtype: Dtype, copy: bool = True) -> ArrayLike: + ... + + def astype(self, dtype: Dtype, copy: bool = True) -> ArrayLike: """ Cast to a NumPy array or ExtensionArray with 'dtype'. diff --git a/pandas/core/arrays/categorical.py b/pandas/core/arrays/categorical.py index 0062ed01e957a..24a3568e500fa 100644 --- a/pandas/core/arrays/categorical.py +++ b/pandas/core/arrays/categorical.py @@ -15,6 +15,7 @@ TypeVar, Union, cast, + overload, ) from warnings import warn @@ -482,6 +483,14 @@ def _constructor(self) -> Type[Categorical]: def _from_sequence(cls, scalars, *, dtype: Optional[Dtype] = None, copy=False): return Categorical(scalars, dtype=dtype, copy=copy) + @overload + def astype(self, dtype: Type[str], copy: bool = True) -> np.ndarray: + ... + + @overload + def astype(self, dtype: Dtype, copy: bool = True) -> ArrayLike: + ... + def astype(self, dtype: Dtype, copy: bool = True) -> ArrayLike: """ Coerce this type to another dtype @@ -2444,11 +2453,7 @@ def _str_get_dummies(self, sep="|"): # sep may not be in categories. Just bail on this. from pandas.core.arrays import PandasArray - # error: Argument 1 to "PandasArray" has incompatible type - # "ExtensionArray"; expected "Union[ndarray, PandasArray]" - return PandasArray(self.astype(str))._str_get_dummies( # type: ignore[arg-type] - sep - ) + return PandasArray(self.astype(str))._str_get_dummies(sep) # The Series.cat accessor diff --git a/pandas/core/arrays/datetimelike.py b/pandas/core/arrays/datetimelike.py index 0900688e04374..9830be3a39e06 100644 --- a/pandas/core/arrays/datetimelike.py +++ b/pandas/core/arrays/datetimelike.py @@ -9,6 +9,7 @@ TYPE_CHECKING, Any, Callable, + List, Optional, Sequence, Tuple, @@ -341,7 +342,7 @@ def __array__(self, dtype: Optional[NpDtype] = None) -> np.ndarray: return self._ndarray def __getitem__( - self, key: Union[int, slice, np.ndarray] + self, key: Union[int, slice, np.ndarray, List[Any]] ) -> Union[DatetimeLikeArrayMixin, DTScalarOrNaT]: """ This getitem defers to the underlying array, which by-definition can @@ -563,8 +564,7 @@ def _validate_comparison_value(self, other): raise InvalidComparison(other) if isinstance(other, self._recognized_scalars) or other is NaT: - # error: Too many arguments for "object" - other = self._scalar_type(other) # type: ignore[call-arg] + other = self._scalar_type(other) try: self._check_compatible_with(other) except (TypeError, IncompatibleFrequency) as err: @@ -614,16 +614,14 @@ def _validate_shift_value(self, fill_value): if is_valid_na_for_dtype(fill_value, self.dtype): fill_value = NaT elif isinstance(fill_value, self._recognized_scalars): - # error: Too many arguments for "object" - fill_value = self._scalar_type(fill_value) # type: ignore[call-arg] + fill_value = self._scalar_type(fill_value) else: # only warn if we're not going to raise if self._scalar_type is Period and lib.is_integer(fill_value): # kludge for #31971 since Period(integer) tries to cast to str new_fill = Period._from_ordinal(fill_value, freq=self.freq) else: - # error: Too many arguments for "object" - new_fill = self._scalar_type(fill_value) # type: ignore[call-arg] + new_fill = self._scalar_type(fill_value) # stacklevel here is chosen to be correct when called from # DataFrame.shift or Series.shift @@ -684,8 +682,7 @@ def _validate_scalar( raise TypeError(msg) elif isinstance(value, self._recognized_scalars): - # error: Too many arguments for "object" - value = self._scalar_type(value) # type: ignore[call-arg] + value = self._scalar_type(value) else: msg = self._validation_error_message(value, allow_listlike) diff --git a/pandas/core/arrays/floating.py b/pandas/core/arrays/floating.py index fdd358a1b3856..8119906c1c5c0 100644 --- a/pandas/core/arrays/floating.py +++ b/pandas/core/arrays/floating.py @@ -5,6 +5,7 @@ Optional, Tuple, Type, + overload, ) import warnings @@ -16,6 +17,7 @@ ) from pandas._typing import ( ArrayLike, + Dtype, DtypeObj, ) from pandas.compat.numpy import function as nv @@ -277,7 +279,15 @@ def _from_sequence_of_strings( def _coerce_to_array(self, value) -> Tuple[np.ndarray, np.ndarray]: return coerce_to_array(value, dtype=self.dtype) - def astype(self, dtype, copy: bool = True) -> ArrayLike: + @overload + def astype(self, dtype: Type[str], copy: bool = True) -> np.ndarray: + ... + + @overload + def astype(self, dtype: Dtype, copy: bool = True) -> ArrayLike: + ... + + def astype(self, dtype: Dtype, copy: bool = True) -> ArrayLike: """ Cast to a NumPy array or ExtensionArray with 'dtype'. diff --git a/pandas/core/arrays/integer.py b/pandas/core/arrays/integer.py index ae44acf06591f..72c67d25895a8 100644 --- a/pandas/core/arrays/integer.py +++ b/pandas/core/arrays/integer.py @@ -6,6 +6,7 @@ Optional, Tuple, Type, + overload, ) import warnings @@ -342,6 +343,14 @@ def _from_sequence_of_strings( def _coerce_to_array(self, value) -> Tuple[np.ndarray, np.ndarray]: return coerce_to_array(value, dtype=self.dtype) + @overload + def astype(self, dtype: Type[str], copy: bool = True) -> np.ndarray: + ... + + @overload + def astype(self, dtype: Dtype, copy: bool = True) -> ArrayLike: + ... + def astype(self, dtype, copy: bool = True) -> ArrayLike: """ Cast to a NumPy array or ExtensionArray with 'dtype'. diff --git a/pandas/core/arrays/interval.py b/pandas/core/arrays/interval.py index 1a626f327b97c..e8df10d66fe3d 100644 --- a/pandas/core/arrays/interval.py +++ b/pandas/core/arrays/interval.py @@ -11,6 +11,7 @@ Sequence, Type, TypeVar, + Union, cast, ) @@ -1497,7 +1498,9 @@ def delete(self: IntervalArrayT, loc) -> IntervalArrayT: return self._shallow_copy(left=new_left, right=new_right) @Appender(_extension_array_shared_docs["repeat"] % _shared_docs_kwargs) - def repeat(self: IntervalArrayT, repeats: int, axis=None) -> IntervalArrayT: + def repeat( + self: IntervalArrayT, repeats: Union[int, Sequence[int]], axis=None + ) -> IntervalArrayT: nv.validate_repeat((), {"axis": axis}) left_repeat = self.left.repeat(repeats) right_repeat = self.right.repeat(repeats) diff --git a/pandas/core/arrays/masked.py b/pandas/core/arrays/masked.py index 31d58d9d89d49..0ea067fef887f 100644 --- a/pandas/core/arrays/masked.py +++ b/pandas/core/arrays/masked.py @@ -3,12 +3,14 @@ from typing import ( TYPE_CHECKING, Any, + List, Optional, Sequence, Tuple, Type, TypeVar, Union, + overload, ) import numpy as np @@ -138,7 +140,7 @@ def dtype(self) -> BaseMaskedDtype: raise AbstractMethodError(self) def __getitem__( - self, item: Union[int, slice, np.ndarray] + self, item: Union[int, slice, np.ndarray, List[Any]] ) -> Union[BaseMaskedArray, Any]: if is_integer(item): if self._mask[item]: @@ -212,14 +214,12 @@ def __len__(self) -> int: def __invert__(self: BaseMaskedArrayT) -> BaseMaskedArrayT: return type(self)(~self._data, self._mask.copy()) - # error: Argument 1 of "to_numpy" is incompatible with supertype "ExtensionArray"; - # supertype defines the argument type as "Union[ExtensionDtype, str, dtype[Any], - # Type[str], Type[float], Type[int], Type[complex], Type[bool], Type[object], None]" - def to_numpy( # type: ignore[override] + def to_numpy( self, dtype: Optional[NpDtype] = None, copy: bool = False, na_value: Scalar = lib.no_default, + **kwargs: Any, ) -> np.ndarray: """ Convert to a NumPy Array. @@ -305,6 +305,14 @@ def to_numpy( # type: ignore[override] data = self._data.astype(dtype, copy=copy) return data + @overload + def astype(self, dtype: Type[str], copy: bool = True) -> np.ndarray: + ... + + @overload + def astype(self, dtype: Dtype, copy: bool = True) -> ArrayLike: + ... + def astype(self, dtype: Dtype, copy: bool = True) -> ArrayLike: dtype = pandas_dtype(dtype) diff --git a/pandas/core/arrays/numpy_.py b/pandas/core/arrays/numpy_.py index 5ef3c24726924..0a26d0bb77820 100644 --- a/pandas/core/arrays/numpy_.py +++ b/pandas/core/arrays/numpy_.py @@ -2,6 +2,7 @@ import numbers from typing import ( + Any, Optional, Tuple, Union, @@ -346,14 +347,12 @@ def skew( # ------------------------------------------------------------------------ # Additional Methods - # error: Argument 1 of "to_numpy" is incompatible with supertype "ExtensionArray"; - # supertype defines the argument type as "Union[ExtensionDtype, str, dtype[Any], - # Type[str], Type[float], Type[int], Type[complex], Type[bool], Type[object], None]" - def to_numpy( # type: ignore[override] + def to_numpy( self, dtype: Optional[NpDtype] = None, copy: bool = False, na_value=lib.no_default, + **kwargs: Any, ) -> np.ndarray: result = np.asarray(self._ndarray, dtype=dtype) diff --git a/pandas/core/arrays/string_arrow.py b/pandas/core/arrays/string_arrow.py index 6f7badd3c2cd2..e7475d254d590 100644 --- a/pandas/core/arrays/string_arrow.py +++ b/pandas/core/arrays/string_arrow.py @@ -18,7 +18,7 @@ missing as libmissing, ) from pandas._typing import ( - Dtype, + DtypeArg, NpDtype, ) from pandas.util._decorators import doc @@ -221,7 +221,7 @@ def _chk_pyarrow_available(cls) -> None: raise ImportError(msg) @classmethod - def _from_sequence(cls, scalars, dtype: Optional[Dtype] = None, copy=False): + def _from_sequence(cls, scalars, dtype: Optional[DtypeArg] = None, copy=False): cls._chk_pyarrow_available() # convert non-na-likes to str, and nan-likes to ArrowStringDtype.na_value scalars = lib.ensure_string_array(scalars, copy=False) @@ -229,7 +229,7 @@ def _from_sequence(cls, scalars, dtype: Optional[Dtype] = None, copy=False): @classmethod def _from_sequence_of_strings( - cls, strings, dtype: Optional[Dtype] = None, copy=False + cls, strings, dtype: Optional[DtypeArg] = None, copy=False ): return cls._from_sequence(strings, dtype=dtype, copy=copy) @@ -248,14 +248,12 @@ def __arrow_array__(self, type=None): """Convert myself to a pyarrow Array or ChunkedArray.""" return self._data - # error: Argument 1 of "to_numpy" is incompatible with supertype "ExtensionArray"; - # supertype defines the argument type as "Union[ExtensionDtype, str, dtype[Any], - # Type[str], Type[float], Type[int], Type[complex], Type[bool], Type[object], None]" - def to_numpy( # type: ignore[override] + def to_numpy( self, dtype: Optional[NpDtype] = None, copy: bool = False, na_value=lib.no_default, + **kwargs: Any, ) -> np.ndarray: """ Convert to a NumPy ndarray. diff --git a/pandas/core/base.py b/pandas/core/base.py index 56ec2597314b2..7b552c05312c0 100644 --- a/pandas/core/base.py +++ b/pandas/core/base.py @@ -20,9 +20,10 @@ import pandas._libs.lib as lib from pandas._typing import ( - Dtype, + ArrayLike, DtypeObj, IndexLabel, + NpDtype, ) from pandas.compat import PYPY from pandas.compat.numpy import function as nv @@ -507,7 +508,7 @@ def array(self) -> ExtensionArray: def to_numpy( self, - dtype: Optional[Dtype] = None, + dtype: Optional[NpDtype] = None, copy: bool = False, na_value=lib.no_default, **kwargs, @@ -607,22 +608,14 @@ def to_numpy( dtype='datetime64[ns]') """ if is_extension_array_dtype(self.dtype): - # error: Too many arguments for "to_numpy" of "ExtensionArray" - return self.array.to_numpy( # type: ignore[call-arg] - dtype, copy=copy, na_value=na_value, **kwargs - ) + return self.array.to_numpy(dtype, copy=copy, na_value=na_value, **kwargs) elif kwargs: bad_keys = list(kwargs.keys())[0] raise TypeError( f"to_numpy() got an unexpected keyword argument '{bad_keys}'" ) - # error: Argument "dtype" to "asarray" has incompatible type - # "Union[ExtensionDtype, str, dtype[Any], Type[str], Type[float], Type[int], - # Type[complex], Type[bool], Type[object], None]"; expected "Union[dtype[Any], - # None, type, _SupportsDType, str, Union[Tuple[Any, int], Tuple[Any, Union[int, - # Sequence[int]]], List[Any], _DTypeDict, Tuple[Any, Any]]]" - result = np.asarray(self._values, dtype=dtype) # type: ignore[arg-type] + result = np.asarray(self._values, dtype=dtype) # TODO(GH-24345): Avoid potential double copy if copy or na_value is not lib.no_default: result = result.copy() @@ -1073,7 +1066,7 @@ def unique(self): values = self._values if not isinstance(values, np.ndarray): - result = values.unique() + result: ArrayLike = values.unique() if self.dtype.kind in ["m", "M"] and isinstance(self, ABCSeries): # GH#31182 Series._values returns EA, unpack for backward-compat if getattr(self.dtype, "tz", None) is None: diff --git a/pandas/core/construction.py b/pandas/core/construction.py index 78a7f1890b5de..b5a17e1ef882e 100644 --- a/pandas/core/construction.py +++ b/pandas/core/construction.py @@ -33,7 +33,7 @@ from pandas.core.dtypes.base import ( ExtensionDtype, - registry, + _registry as registry, ) from pandas.core.dtypes.cast import ( construct_1d_arraylike_from_scalar, diff --git a/pandas/core/dtypes/base.py b/pandas/core/dtypes/base.py index d83405803753a..d7b1b05f5d5a9 100644 --- a/pandas/core/dtypes/base.py +++ b/pandas/core/dtypes/base.py @@ -12,6 +12,7 @@ Tuple, Type, Union, + cast, ) import numpy as np @@ -211,7 +212,7 @@ def construct_array_type(cls) -> Type[ExtensionArray]: raise NotImplementedError @classmethod - def construct_from_string(cls, string: str): + def construct_from_string(cls, string: str) -> ExtensionDtype: r""" Construct this type from a string. @@ -388,7 +389,7 @@ def register_extension_dtype(cls: Type[ExtensionDtype]) -> Type[ExtensionDtype]: ... class MyExtensionDtype(ExtensionDtype): ... name = "myextension" """ - registry.register(cls) + _registry.register(cls) return cls @@ -422,9 +423,7 @@ def register(self, dtype: Type[ExtensionDtype]) -> None: self.dtypes.append(dtype) - def find( - self, dtype: Union[Type[ExtensionDtype], str] - ) -> Optional[Type[ExtensionDtype]]: + def find(self, dtype: Union[Type[ExtensionDtype], str]) -> Optional[ExtensionDtype]: """ Parameters ---------- @@ -439,7 +438,7 @@ def find( if not isinstance(dtype, type): dtype_type = type(dtype) if issubclass(dtype_type, ExtensionDtype): - return dtype + return cast(ExtensionDtype, dtype) return None @@ -452,4 +451,4 @@ def find( return None -registry = Registry() +_registry = Registry() diff --git a/pandas/core/dtypes/common.py b/pandas/core/dtypes/common.py index 7a2d6468f1b63..d087c6ba50096 100644 --- a/pandas/core/dtypes/common.py +++ b/pandas/core/dtypes/common.py @@ -3,9 +3,11 @@ """ from typing import ( + TYPE_CHECKING, Any, Callable, Union, + cast, ) import warnings @@ -23,7 +25,7 @@ Optional, ) -from pandas.core.dtypes.base import registry +from pandas.core.dtypes.base import _registry as registry from pandas.core.dtypes.dtypes import ( CategoricalDtype, DatetimeTZDtype, @@ -58,6 +60,9 @@ is_sequence, ) +if TYPE_CHECKING: + from pandas.core.arrays.base import ExtensionArray + POSSIBLY_CAST_DTYPES = { np.dtype(t).name for t in [ @@ -128,7 +133,7 @@ def ensure_str(value: Union[bytes, Any]) -> str: return value -def ensure_int_or_float(arr: ArrayLike, copy: bool = False) -> np.ndarray: +def ensure_int_or_float(arr: ArrayLike, copy: bool = False) -> ArrayLike: """ Ensure that an dtype array of some integer dtype has an int64 dtype if possible. @@ -155,22 +160,18 @@ def ensure_int_or_float(arr: ArrayLike, copy: bool = False) -> np.ndarray: will remain unchanged. """ # TODO: GH27506 potential bug with ExtensionArrays - try: - # error: Unexpected keyword argument "casting" for "astype" - return arr.astype("int64", copy=copy, casting="safe") # type: ignore[call-arg] - except TypeError: - pass - try: - # error: Unexpected keyword argument "casting" for "astype" - return arr.astype("uint64", copy=copy, casting="safe") # type: ignore[call-arg] - except TypeError: - if is_extension_array_dtype(arr.dtype): - # pandas/core/dtypes/common.py:168: error: Item "ndarray" of - # "Union[ExtensionArray, ndarray]" has no attribute "to_numpy" [union-attr] - return arr.to_numpy( # type: ignore[union-attr] - dtype="float64", na_value=np.nan - ) - return arr.astype("float64", copy=copy) + if is_extension_array_dtype(arr.dtype): + return cast("ExtensionArray", arr).to_numpy(dtype="float64", na_value=np.nan) + else: + assert isinstance(arr, np.ndarray) # For typing + try: + return arr.astype("int64", copy=copy, casting="safe") + except TypeError: + pass + try: + return arr.astype("uint64", copy=copy, casting="safe") + except TypeError: + return arr.astype("float64", copy=copy) def ensure_python_int(value: Union[int, np.integer]) -> int: @@ -1848,9 +1849,7 @@ def pandas_dtype(dtype) -> DtypeObj: # registered extension types result = registry.find(dtype) if result is not None: - # error: Incompatible return value type (got "Type[ExtensionDtype]", - # expected "Union[dtype, ExtensionDtype]") - return result # type: ignore[return-value] + return result # try a numpy dtype # raise a consistent TypeError if failed diff --git a/pandas/core/groupby/groupby.py b/pandas/core/groupby/groupby.py index 979c7aa990184..7c18d4b74c19c 100644 --- a/pandas/core/groupby/groupby.py +++ b/pandas/core/groupby/groupby.py @@ -2275,12 +2275,12 @@ def pre_processor(vals: ArrayLike) -> Tuple[np.ndarray, Optional[np.dtype]]: inference: Optional[np.dtype] = None if is_integer_dtype(vals.dtype): if isinstance(vals, ExtensionArray): - out = vals.to_numpy(dtype=float, na_value=np.nan) + out = vals.to_numpy(dtype=np.dtype(float), na_value=np.nan) else: out = vals inference = np.dtype(np.int64) elif is_bool_dtype(vals.dtype) and isinstance(vals, ExtensionArray): - out = vals.to_numpy(dtype=float, na_value=np.nan) + out = vals.to_numpy(dtype=np.dtype(float), na_value=np.nan) elif is_datetime64_dtype(vals.dtype): inference = np.dtype("datetime64[ns]") out = np.asarray(vals).astype(float) diff --git a/pandas/core/indexes/base.py b/pandas/core/indexes/base.py index 26d25645b02c6..83908ee6174d7 100644 --- a/pandas/core/indexes/base.py +++ b/pandas/core/indexes/base.py @@ -4577,11 +4577,7 @@ def __getitem__(self, key): result = getitem(key) if not is_scalar(result): - # error: Argument 1 to "ndim" has incompatible type "Union[ExtensionArray, - # Any]"; expected "Union[Union[int, float, complex, str, bytes, generic], - # Sequence[Union[int, float, complex, str, bytes, generic]], - # Sequence[Sequence[Any]], _SupportsArray]" - if np.ndim(result) > 1: # type: ignore[arg-type] + if np.ndim(result) > 1: deprecate_ndim_indexing(result) return result # NB: Using _constructor._simple_new would break if MultiIndex diff --git a/pandas/core/internals/blocks.py b/pandas/core/internals/blocks.py index b1a552cff2274..b6fc61f014b06 100644 --- a/pandas/core/internals/blocks.py +++ b/pandas/core/internals/blocks.py @@ -239,7 +239,9 @@ def get_values(self, dtype: Optional[DtypeObj] = None) -> np.ndarray: this is often overridden to handle to_dense like operations """ if dtype == _dtype_obj: - return self.values.astype(_dtype_obj) + # error: Incompatible return value type (got "Union[ndarray, + # ExtensionArray]", expected "ndarray") + return self.values.astype(_dtype_obj) # type: ignore[return-value] # error: Incompatible return value type (got "Union[ndarray, ExtensionArray]", # expected "ndarray") return self.values # type: ignore[return-value] @@ -1468,9 +1470,7 @@ def iget(self, col): elif isinstance(col, slice): if col != slice(None): raise NotImplementedError(col) - # error: Invalid index type "List[Any]" for "ExtensionArray"; expected - # type "Union[int, slice, ndarray]" - return self.values[[loc]] # type: ignore[index] + return self.values[[loc]] return self.values[loc] else: if col != 0: diff --git a/pandas/core/internals/concat.py b/pandas/core/internals/concat.py index b82ab807562f4..d6d7d13523699 100644 --- a/pandas/core/internals/concat.py +++ b/pandas/core/internals/concat.py @@ -425,6 +425,7 @@ def _concatenate_join_units( elif any(isinstance(t, ExtensionArray) for t in to_concat): # concatting with at least one EA means we are concatting a single column # the non-EA values are 2D arrays with shape (1, n) + to_concat = [t if isinstance(t, ExtensionArray) else t[0, :] for t in to_concat] concat_values = concat_compat(to_concat, axis=0, ea_compat_axis=True) concat_values = ensure_block_shape(concat_values, 2) diff --git a/pandas/core/internals/managers.py b/pandas/core/internals/managers.py index 9c21fcf957ecd..7c87e80547d72 100644 --- a/pandas/core/internals/managers.py +++ b/pandas/core/internals/managers.py @@ -3,6 +3,7 @@ from collections import defaultdict import itertools from typing import ( + TYPE_CHECKING, Any, Callable, DefaultDict, @@ -14,6 +15,7 @@ Tuple, TypeVar, Union, + cast, ) import warnings @@ -27,6 +29,7 @@ ArrayLike, Dtype, DtypeObj, + NpDtype, Shape, ) from pandas.errors import PerformanceWarning @@ -79,6 +82,9 @@ operate_blockwise, ) +if TYPE_CHECKING: + from pandas.core.arrays.base import ExtensionArray + # TODO: flexible with index=None and/or items=None T = TypeVar("T", bound="BlockManager") @@ -845,7 +851,7 @@ def copy_func(ax): def as_array( self, transpose: bool = False, - dtype: Optional[Dtype] = None, + dtype: Optional[NpDtype] = None, copy: bool = False, na_value=lib.no_default, ) -> np.ndarray: @@ -890,12 +896,7 @@ def as_array( else: arr = np.asarray(blk.get_values()) if dtype: - # error: Argument 1 to "astype" of "_ArrayOrScalarCommon" has - # incompatible type "Union[ExtensionDtype, str, dtype[Any], - # Type[object]]"; expected "Union[dtype[Any], None, type, - # _SupportsDType, str, Union[Tuple[Any, int], Tuple[Any, Union[int, - # Sequence[int]]], List[Any], _DTypeDict, Tuple[Any, Any]]]" - arr = arr.astype(dtype, copy=False) # type: ignore[arg-type] + arr = arr.astype(dtype, copy=False) else: arr = self._interleave(dtype=dtype, na_value=na_value) # The underlying data was copied within _interleave @@ -928,25 +929,15 @@ def _interleave( elif is_dtype_equal(dtype, str): dtype = "object" - # error: Argument "dtype" to "empty" has incompatible type - # "Union[ExtensionDtype, str, dtype[Any], Type[object], None]"; expected - # "Union[dtype[Any], None, type, _SupportsDType, str, Union[Tuple[Any, int], - # Tuple[Any, Union[int, Sequence[int]]], List[Any], _DTypeDict, Tuple[Any, - # Any]]]" - result = np.empty(self.shape, dtype=dtype) # type: ignore[arg-type] - + result = np.empty(self.shape, dtype=cast(NpDtype, dtype)) itemmask = np.zeros(self.shape[0]) for blk in self.blocks: rl = blk.mgr_locs if blk.is_extension: # Avoid implicit conversion of extension blocks to object - - # error: Item "ndarray" of "Union[ndarray, ExtensionArray]" has no - # attribute "to_numpy" - arr = blk.values.to_numpy( # type: ignore[union-attr] - dtype=dtype, na_value=na_value - ) + blk_values = cast("ExtensionArray", blk.values) + arr = blk_values.to_numpy(dtype=cast(NpDtype, dtype), na_value=na_value) else: # error: Argument 1 to "get_values" of "Block" has incompatible type # "Union[ExtensionDtype, str, dtype[Any], Type[object], None]"; expected diff --git a/pandas/core/internals/ops.py b/pandas/core/internals/ops.py index 88e70723517e3..46eb7da4b6725 100644 --- a/pandas/core/internals/ops.py +++ b/pandas/core/internals/ops.py @@ -108,28 +108,28 @@ def _get_same_shape_values( # TODO(EA2D): with 2D EAs only this first clause would be needed if not (left_ea or right_ea): - # error: Invalid index type "Tuple[Any, slice]" for "Union[ndarray, - # ExtensionArray]"; expected type "Union[int, slice, ndarray]" - lvals = lvals[rblk.mgr_locs.indexer, :] # type: ignore[index] + # error: No overload variant of "__getitem__" of + # "ExtensionArray" matches argument type "Tuple[Any, slice]" + lvals = lvals[rblk.mgr_locs.indexer, :] # type: ignore[call-overload] assert lvals.shape == rvals.shape, (lvals.shape, rvals.shape) elif left_ea and right_ea: assert lvals.shape == rvals.shape, (lvals.shape, rvals.shape) elif right_ea: # lvals are 2D, rvals are 1D - # error: Invalid index type "Tuple[Any, slice]" for "Union[ndarray, - # ExtensionArray]"; expected type "Union[int, slice, ndarray]" - lvals = lvals[rblk.mgr_locs.indexer, :] # type: ignore[index] + # error: No overload variant of "__getitem__" of + # "ExtensionArray" matches argument type "Tuple[Any, slice]"" + lvals = lvals[rblk.mgr_locs.indexer, :] # type: ignore[call-overload] assert lvals.shape[0] == 1, lvals.shape - # error: Invalid index type "Tuple[int, slice]" for "Union[Any, - # ExtensionArray]"; expected type "Union[int, slice, ndarray]" - lvals = lvals[0, :] # type: ignore[index] + # error: No overload variant of "__getitem__" of + # "ExtensionArray" matches argument type "Tuple[int, slice]" + lvals = lvals[0, :] # type: ignore[call-overload] else: # lvals are 1D, rvals are 2D assert rvals.shape[0] == 1, rvals.shape - # error: Invalid index type "Tuple[int, slice]" for "Union[ndarray, - # ExtensionArray]"; expected type "Union[int, slice, ndarray]" - rvals = rvals[0, :] # type: ignore[index] + # error: No overload variant of "__getitem__" of + # "ExtensionArray" matches argument type "Tuple[int, slice]" + rvals = rvals[0, :] # type: ignore[call-overload] return lvals, rvals diff --git a/pandas/core/missing.py b/pandas/core/missing.py index 48b2084319292..bd6f645e7f424 100644 --- a/pandas/core/missing.py +++ b/pandas/core/missing.py @@ -310,11 +310,7 @@ def interpolate_1d( if method in NP_METHODS: # np.interp requires sorted X values, #21037 - # error: Argument 1 to "argsort" has incompatible type "Union[ExtensionArray, - # Any]"; expected "Union[Union[int, float, complex, str, bytes, generic], - # Sequence[Union[int, float, complex, str, bytes, generic]], - # Sequence[Sequence[Any]], _SupportsArray]" - indexer = np.argsort(inds[valid]) # type: ignore[arg-type] + indexer = np.argsort(inds[valid]) result[invalid] = np.interp( inds[invalid], inds[valid][indexer], yvalues[valid][indexer] ) diff --git a/pandas/core/tools/datetimes.py b/pandas/core/tools/datetimes.py index 9822356d11d7c..7f34728972b63 100644 --- a/pandas/core/tools/datetimes.py +++ b/pandas/core/tools/datetimes.py @@ -1075,9 +1075,7 @@ def calc_with_mask(carg, mask): # string with NaN-like try: - # error: Argument 2 to "isin" has incompatible type "List[Any]"; expected - # "Union[Union[ExtensionArray, ndarray], Index, Series]" - mask = ~algorithms.isin(arg, list(nat_strings)) # type: ignore[arg-type] + mask = ~algorithms.isin(arg, list(nat_strings)) return calc_with_mask(arg, mask) except (ValueError, OverflowError, TypeError): pass diff --git a/pandas/io/formats/format.py b/pandas/io/formats/format.py index a768ec8ad4eb3..38f7631fbf3f9 100644 --- a/pandas/io/formats/format.py +++ b/pandas/io/formats/format.py @@ -1797,17 +1797,13 @@ def get_format_timedelta64( If box, then show the return in quotes """ - values_int = values.view(np.int64) + values_int = cast(np.ndarray, values.view(np.int64)) consider_values = values_int != iNaT one_day_nanos = 86400 * 10 ** 9 even_days = ( - # error: Unsupported operand types for % ("ExtensionArray" and "int") - np.logical_and( - consider_values, values_int % one_day_nanos != 0 # type: ignore[operator] - ).sum() - == 0 + np.logical_and(consider_values, values_int % one_day_nanos != 0).sum() == 0 ) if even_days: diff --git a/pandas/io/parsers/base_parser.py b/pandas/io/parsers/base_parser.py index a011a789bf17c..a8a44dacd41cf 100644 --- a/pandas/io/parsers/base_parser.py +++ b/pandas/io/parsers/base_parser.py @@ -531,11 +531,7 @@ def _convert_to_ndarrays( try: values = lib.map_infer(values, conv_f) except ValueError: - # error: Argument 2 to "isin" has incompatible type "List[Any]"; - # expected "Union[Union[ExtensionArray, ndarray], Index, Series]" - mask = algorithms.isin( - values, list(na_values) # type: ignore[arg-type] - ).view(np.uint8) + mask = algorithms.isin(values, list(na_values)).view(np.uint8) values = lib.map_infer_mask(values, conv_f, mask) cvals, na_count = self._infer_types( @@ -661,9 +657,7 @@ def _infer_types(self, values, na_values, try_num_bool=True): """ na_count = 0 if issubclass(values.dtype.type, (np.number, np.bool_)): - # error: Argument 2 to "isin" has incompatible type "List[Any]"; expected - # "Union[Union[ExtensionArray, ndarray], Index, Series]" - mask = algorithms.isin(values, list(na_values)) # type: ignore[arg-type] + mask = algorithms.isin(values, list(na_values)) # error: Incompatible types in assignment (expression has type # "number[Any]", variable has type "int") na_count = mask.sum() # type: ignore[assignment] diff --git a/pandas/tests/arrays/test_array.py b/pandas/tests/arrays/test_array.py index 5d2b7c43f6765..bfe588883d9f3 100644 --- a/pandas/tests/arrays/test_array.py +++ b/pandas/tests/arrays/test_array.py @@ -5,7 +5,7 @@ import pytest import pytz -from pandas.core.dtypes.base import registry +from pandas.core.dtypes.base import _registry as registry import pandas as pd import pandas._testing as tm diff --git a/pandas/tests/arrays/test_period.py b/pandas/tests/arrays/test_period.py index e7f3e8c659316..2592a0263c585 100644 --- a/pandas/tests/arrays/test_period.py +++ b/pandas/tests/arrays/test_period.py @@ -4,7 +4,7 @@ from pandas._libs.tslibs import iNaT from pandas._libs.tslibs.period import IncompatibleFrequency -from pandas.core.dtypes.base import registry +from pandas.core.dtypes.base import _registry as registry from pandas.core.dtypes.dtypes import PeriodDtype import pandas as pd diff --git a/pandas/tests/base/test_conversion.py b/pandas/tests/base/test_conversion.py index 7045a0abbeb81..1d0971749ec6b 100644 --- a/pandas/tests/base/test_conversion.py +++ b/pandas/tests/base/test_conversion.py @@ -426,11 +426,6 @@ def test_to_numpy_kwargs_raises(): with pytest.raises(TypeError, match=msg): s.to_numpy(foo=True) - # extension - s = Series([1, 2, 3], dtype="Int64") - with pytest.raises(TypeError, match=msg): - s.to_numpy(foo=True) - @pytest.mark.parametrize( "data", diff --git a/pandas/tests/dtypes/test_dtypes.py b/pandas/tests/dtypes/test_dtypes.py index ca311768dc2d9..51a7969162abf 100644 --- a/pandas/tests/dtypes/test_dtypes.py +++ b/pandas/tests/dtypes/test_dtypes.py @@ -4,7 +4,7 @@ import pytest import pytz -from pandas.core.dtypes.base import registry +from pandas.core.dtypes.base import _registry as registry from pandas.core.dtypes.common import ( is_bool_dtype, is_categorical, diff --git a/pandas/tests/extension/decimal/array.py b/pandas/tests/extension/decimal/array.py index 58e5dc34d59d5..316f31ed879c6 100644 --- a/pandas/tests/extension/decimal/array.py +++ b/pandas/tests/extension/decimal/array.py @@ -5,6 +5,7 @@ import random import sys from typing import ( + Any, Type, Union, ) @@ -103,9 +104,10 @@ def _from_factorized(cls, values, original): _HANDLED_TYPES = (decimal.Decimal, numbers.Number, np.ndarray) def to_numpy( - self, dtype=None, copy: bool = False, na_value=no_default, decimals=None + self, dtype=None, copy: bool = False, na_value=no_default, **kwargs: Any ) -> np.ndarray: result = np.asarray(self, dtype=dtype) + decimals = kwargs.get("decimals", None) if decimals is not None: result = np.asarray([round(x, decimals) for x in result]) return result diff --git a/pandas/tests/frame/indexing/test_setitem.py b/pandas/tests/frame/indexing/test_setitem.py index 661df8a792c65..bc954d26edbeb 100644 --- a/pandas/tests/frame/indexing/test_setitem.py +++ b/pandas/tests/frame/indexing/test_setitem.py @@ -5,7 +5,7 @@ import pandas.util._test_decorators as td -from pandas.core.dtypes.base import registry as ea_registry +from pandas.core.dtypes.base import _registry as ea_registry from pandas.core.dtypes.common import ( is_categorical_dtype, is_interval_dtype,