annotations (#32193)

Author: jbrockmendel

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,15 +1528,15 @@ 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)
 
     @property
     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):

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.

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 `^`.