diff --git a/pandas/core/internals/managers.py b/pandas/core/internals/managers.py index 69ceb95985140..d4f9c15a9f73f 100644 --- a/pandas/core/internals/managers.py +++ b/pandas/core/internals/managers.py @@ -8,6 +8,7 @@ import numpy as np from pandas._libs import Timedelta, Timestamp, internals as libinternals, lib +from pandas._typing import DtypeObj from pandas.util._validators import validate_bool_kwarg from pandas.core.dtypes.cast import ( @@ -847,7 +848,7 @@ def to_dict(self, copy: bool = True): return {dtype: self.combine(blocks, copy=copy) for dtype, blocks in bd.items()} - def fast_xs(self, loc): + def fast_xs(self, loc: int): """ get a cross sectional for a given location in the items ; handle dups @@ -883,12 +884,12 @@ def fast_xs(self, loc): for i, rl in enumerate(blk.mgr_locs): result[rl] = blk.iget((i, loc)) - if is_extension_array_dtype(dtype): + if isinstance(dtype, ExtensionDtype): result = dtype.construct_array_type()._from_sequence(result, dtype=dtype) return result - def consolidate(self): + def consolidate(self) -> "BlockManager": """ Join together blocks having same dtype @@ -940,7 +941,7 @@ def get(self, item): new_axis=self.items[indexer], indexer=indexer, axis=0, allow_dups=True ) - def iget(self, i): + def iget(self, i: int) -> "SingleBlockManager": """ Return the data as a SingleBlockManager. """ @@ -1377,7 +1378,7 @@ def canonicalize(block): block.equals(oblock) for block, oblock in zip(self_blocks, other_blocks) ) - def unstack(self, unstacker_func, fill_value): + def unstack(self, unstacker_func, fill_value) -> "BlockManager": """ Return a BlockManager with all blocks unstacked.. @@ -1396,8 +1397,8 @@ def unstack(self, unstacker_func, fill_value): dummy = unstacker_func(np.empty((0, 0)), value_columns=self.items) new_columns = dummy.get_new_columns() new_index = dummy.get_new_index() - new_blocks = [] - columns_mask = [] + new_blocks: List[Block] = [] + columns_mask: List[np.ndarray] = [] for blk in self.blocks: blocks, mask = blk._unstack( @@ -1478,7 +1479,7 @@ def _post_setstate(self): pass @property - def _block(self): + def _block(self) -> Block: return self.blocks[0] @property @@ -1495,14 +1496,14 @@ def _blklocs(self): """ compat with BlockManager """ return None - def get_slice(self, slobj, axis=0): + def get_slice(self, slobj: slice, axis: int = 0) -> "SingleBlockManager": if axis >= self.ndim: raise IndexError("Requested axis not found in manager") - return type(self)(self._block._slice(slobj), self.index[slobj], fastpath=True,) + return type(self)(self._block._slice(slobj), self.index[slobj], fastpath=True) @property - def index(self): + def index(self) -> Index: return self.axes[0] @property @@ -1516,7 +1517,7 @@ def array_dtype(self): def get_dtype_counts(self): return {self.dtype.name: 1} - def get_dtypes(self): + def get_dtypes(self) -> np.ndarray: return np.array([self._block.dtype]) def external_values(self): @@ -1527,7 +1528,7 @@ def internal_values(self): """The array that Series._values returns""" return self._block.internal_values() - def get_values(self): + def get_values(self) -> np.ndarray: """ return a dense type view """ return np.array(self._block.to_dense(), copy=False) @@ -1535,7 +1536,7 @@ def get_values(self): def _can_hold_na(self) -> bool: return self._block._can_hold_na - def is_consolidated(self): + def is_consolidated(self) -> bool: return True def _consolidate_check(self): @@ -1813,9 +1814,7 @@ def _shape_compat(x): return stacked, placement -def _interleaved_dtype( - blocks: List[Block], -) -> Optional[Union[np.dtype, ExtensionDtype]]: +def _interleaved_dtype(blocks: Sequence[Block]) -> Optional[DtypeObj]: """ Find the common dtype for `blocks`. @@ -1825,7 +1824,7 @@ def _interleaved_dtype( Returns ------- - dtype : Optional[Union[np.dtype, ExtensionDtype]] + dtype : np.dtype, ExtensionDtype, or None None is returned when `blocks` is empty. """ if not len(blocks): diff --git a/pandas/core/ops/__init__.py b/pandas/core/ops/__init__.py index b74dea686a89f..1ef3889703341 100644 --- a/pandas/core/ops/__init__.py +++ b/pandas/core/ops/__init__.py @@ -5,22 +5,17 @@ """ import datetime import operator -from typing import TYPE_CHECKING, Optional, Set, Tuple, Union +from typing import TYPE_CHECKING, Optional, Set, Tuple import numpy as np from pandas._libs import Timedelta, Timestamp, lib from pandas._libs.ops_dispatch import maybe_dispatch_ufunc_to_dunder_op # noqa:F401 -from pandas._typing import Level +from pandas._typing import ArrayLike, Level from pandas.util._decorators import Appender from pandas.core.dtypes.common import is_list_like, is_timedelta64_dtype -from pandas.core.dtypes.generic import ( - ABCDataFrame, - ABCExtensionArray, - ABCIndexClass, - ABCSeries, -) +from pandas.core.dtypes.generic import ABCDataFrame, ABCIndexClass, ABCSeries from pandas.core.dtypes.missing import isna from pandas.core.construction import extract_array @@ -451,10 +446,7 @@ def _align_method_SERIES(left, right, align_asobject=False): def _construct_result( - left: ABCSeries, - result: Union[np.ndarray, ABCExtensionArray], - index: ABCIndexClass, - name, + left: ABCSeries, result: ArrayLike, index: ABCIndexClass, name, ): """ Construct an appropriately-labelled Series from the result of an op. diff --git a/pandas/core/ops/array_ops.py b/pandas/core/ops/array_ops.py index 37a4a6eddaebe..10e3f32de3958 100644 --- a/pandas/core/ops/array_ops.py +++ b/pandas/core/ops/array_ops.py @@ -4,11 +4,12 @@ """ from functools import partial import operator -from typing import Any, Optional, Union +from typing import Any, Optional import numpy as np from pandas._libs import Timedelta, Timestamp, lib, ops as libops +from pandas._typing import ArrayLike from pandas.core.dtypes.cast import ( construct_1d_object_array_from_listlike, @@ -155,9 +156,7 @@ def na_arithmetic_op(left, right, op, str_rep: str): return missing.dispatch_fill_zeros(op, left, right, result) -def arithmetic_op( - left: Union[np.ndarray, ABCExtensionArray], right: Any, op, str_rep: str -): +def arithmetic_op(left: ArrayLike, right: Any, op, str_rep: str): """ Evaluate an arithmetic operation `+`, `-`, `*`, `/`, `//`, `%`, `**`, ... @@ -200,9 +199,7 @@ def arithmetic_op( return res_values -def comparison_op( - left: Union[np.ndarray, ABCExtensionArray], right: Any, op -) -> Union[np.ndarray, ABCExtensionArray]: +def comparison_op(left: ArrayLike, right: Any, op) -> ArrayLike: """ Evaluate a comparison operation `=`, `!=`, `>=`, `>`, `<=`, or `<`. @@ -215,7 +212,7 @@ def comparison_op( Returns ------- - ndarrray or ExtensionArray + ndarray or ExtensionArray """ # NB: We assume extract_array has already been called on left and right lvalues = left @@ -302,9 +299,7 @@ def na_logical_op(x: np.ndarray, y, op): return result.reshape(x.shape) -def logical_op( - left: Union[np.ndarray, ABCExtensionArray], right: Any, op -) -> Union[np.ndarray, ABCExtensionArray]: +def logical_op(left: ArrayLike, right: Any, op) -> ArrayLike: """ Evaluate a logical operation `|`, `&`, or `^`.