TYP: blocks (#39543)

Author: jbrockmendel

pandas/core/internals/blocks.py

@@ -239,7 +239,7 @@ def array_values(self) -> ExtensionArray:
         """
         return PandasArray(self.values)
 
-    def get_values(self, dtype: Optional[Dtype] = None):
+    def get_values(self, dtype: Optional[DtypeObj] = None) -> np.ndarray:
         """
         return an internal format, currently just the ndarray
         this is often overridden to handle to_dense like operations
@@ -282,7 +282,7 @@ def make_block(self, values, placement=None) -> Block:
 
         return make_block(values, placement=placement, ndim=self.ndim)
 
-    def make_block_same_class(self, values, placement=None, ndim=None):
+    def make_block_same_class(self, values, placement=None, ndim=None) -> Block:
         """ Wrap given values in a block of same type as self. """
         if placement is None:
             placement = self.mgr_locs
@@ -318,7 +318,7 @@ def _slice(self, slicer):
 
         return self.values[slicer]
 
-    def getitem_block(self, slicer, new_mgr_locs=None):
+    def getitem_block(self, slicer, new_mgr_locs=None) -> Block:
         """
         Perform __getitem__-like, return result as block.
 
@@ -338,11 +338,11 @@ def getitem_block(self, slicer, new_mgr_locs=None):
         return type(self)._simple_new(new_values, new_mgr_locs, self.ndim)
 
     @property
-    def shape(self):
+    def shape(self) -> Shape:
         return self.values.shape
 
     @property
-    def dtype(self):
+    def dtype(self) -> DtypeObj:
         return self.values.dtype
 
     def iget(self, i):
@@ -1063,7 +1063,7 @@ def f(mask, val, idx):
             new_blocks = self.split_and_operate(mask, f, True)
             return new_blocks
 
-    def coerce_to_target_dtype(self, other):
+    def coerce_to_target_dtype(self, other) -> Block:
         """
         coerce the current block to a dtype compat for other
         we will return a block, possibly object, and not raise
@@ -1091,13 +1091,13 @@ def interpolate(
         coerce: bool = False,
         downcast: Optional[str] = None,
         **kwargs,
-    ):
+    ) -> List[Block]:
 
         inplace = validate_bool_kwarg(inplace, "inplace")
 
         if not self._can_hold_na:
             # If there are no NAs, then interpolate is a no-op
-            return self if inplace else self.copy()
+            return [self] if inplace else [self.copy()]
 
         # a fill na type method
         try:
@@ -1219,7 +1219,9 @@ def func(yvalues: np.ndarray) -> np.ndarray:
         blocks = [self.make_block_same_class(interp_values)]
         return self._maybe_downcast(blocks, downcast)
 
-    def take_nd(self, indexer, axis: int, new_mgr_locs=None, fill_value=lib.no_default):
+    def take_nd(
+        self, indexer, axis: int, new_mgr_locs=None, fill_value=lib.no_default
+    ) -> Block:
         """
         Take values according to indexer and return them as a block.bb
 
@@ -1256,7 +1258,7 @@ def diff(self, n: int, axis: int = 1) -> List[Block]:
         new_values = algos.diff(self.values, n, axis=axis, stacklevel=7)
         return [self.make_block(values=new_values)]
 
-    def shift(self, periods: int, axis: int = 0, fill_value=None):
+    def shift(self, periods: int, axis: int = 0, fill_value: Any = None) -> List[Block]:
         """ shift the block by periods, possibly upcast """
         # convert integer to float if necessary. need to do a lot more than
         # that, handle boolean etc also
@@ -1369,7 +1371,7 @@ def _unstack(self, unstacker, fill_value, new_placement):
         blocks = [make_block(new_values, placement=new_placement)]
         return blocks, mask
 
-    def quantile(self, qs, interpolation="linear", axis: int = 0):
+    def quantile(self, qs, interpolation="linear", axis: int = 0) -> Block:
         """
         compute the quantiles of the
 
@@ -1521,7 +1523,7 @@ def __init__(self, values, placement, ndim: int):
             raise AssertionError("block.size != values.size")
 
     @property
-    def shape(self):
+    def shape(self) -> Shape:
         # TODO(EA2D): override unnecessary with 2D EAs
         if self.ndim == 1:
             return (len(self.values),)
@@ -1647,7 +1649,7 @@ def setitem(self, indexer, value):
         self.values[indexer] = value
         return self
 
-    def get_values(self, dtype: Optional[Dtype] = None):
+    def get_values(self, dtype: Optional[DtypeObj] = None) -> np.ndarray:
         # ExtensionArrays must be iterable, so this works.
         # TODO(EA2D): reshape not needed with 2D EAs
         return np.asarray(self.values).reshape(self.shape)
@@ -1669,7 +1671,7 @@ def to_native_types(self, na_rep="nan", quoting=None, **kwargs):
 
     def take_nd(
         self, indexer, axis: int = 0, new_mgr_locs=None, fill_value=lib.no_default
-    ):
+    ) -> Block:
         """
         Take values according to indexer and return them as a block.
         """
@@ -1733,7 +1735,9 @@ def _slice(self, slicer):
 
         return self.values[slicer]
 
-    def fillna(self, value, limit=None, inplace=False, downcast=None):
+    def fillna(
+        self, value, limit=None, inplace: bool = False, downcast=None
+    ) -> List[Block]:
         values = self.values if inplace else self.values.copy()
         values = values.fillna(value=value, limit=limit)
         return [
@@ -1765,9 +1769,7 @@ def diff(self, n: int, axis: int = 1) -> List[Block]:
             axis = 0
         return super().diff(n, axis)
 
-    def shift(
-        self, periods: int, axis: int = 0, fill_value: Any = None
-    ) -> List[ExtensionBlock]:
+    def shift(self, periods: int, axis: int = 0, fill_value: Any = None) -> List[Block]:
         """
         Shift the block by `periods`.
 
@@ -1947,7 +1949,7 @@ def _holder(self):
     def fill_value(self):
         return np.datetime64("NaT", "ns")
 
-    def get_values(self, dtype: Optional[Dtype] = None):
+    def get_values(self, dtype: Optional[DtypeObj] = None) -> np.ndarray:
         """
         return object dtype as boxed values, such as Timestamps/Timedelta
         """
@@ -1996,11 +1998,11 @@ def diff(self, n: int, axis: int = 0) -> List[Block]:
             TimeDeltaBlock(new_values, placement=self.mgr_locs.indexer, ndim=self.ndim)
         ]
 
-    def shift(self, periods, axis=0, fill_value=None):
+    def shift(self, periods: int, axis: int = 0, fill_value: Any = None) -> List[Block]:
         # TODO(EA2D) this is unnecessary if these blocks are backed by 2D EAs
         values = self.array_values()
         new_values = values.shift(periods, fill_value=fill_value, axis=axis)
-        return self.make_block_same_class(new_values)
+        return [self.make_block_same_class(new_values)]
 
     def to_native_types(self, na_rep="NaT", **kwargs):
         """ convert to our native types format """
@@ -2118,7 +2120,7 @@ def is_view(self) -> bool:
         # check the ndarray values of the DatetimeIndex values
         return self.values._data.base is not None
 
-    def get_values(self, dtype: Optional[Dtype] = None):
+    def get_values(self, dtype: Optional[DtypeObj] = None) -> np.ndarray:
         """
         Returns an ndarray of values.
 
@@ -2157,7 +2159,9 @@ def external_values(self):
             return self.values._data
         return np.asarray(self.values.astype("datetime64[ns]", copy=False))
 
-    def fillna(self, value, limit=None, inplace=False, downcast=None):
+    def fillna(
+        self, value, limit=None, inplace: bool = False, downcast=None
+    ) -> List[Block]:
         # We support filling a DatetimeTZ with a `value` whose timezone
         # is different by coercing to object.
         if self._can_hold_element(value):
@@ -2168,7 +2172,7 @@ def fillna(self, value, limit=None, inplace=False, downcast=None):
             value, limit=limit, inplace=inplace, downcast=downcast
         )
 
-    def quantile(self, qs, interpolation="linear", axis=0):
+    def quantile(self, qs, interpolation="linear", axis: int = 0) -> Block:
         naive = self.values.view("M8[ns]")
 
         # TODO(EA2D): kludge for 2D block with 1D values
@@ -2228,7 +2232,9 @@ def _maybe_coerce_values(self, values):
     def _holder(self):
         return TimedeltaArray
 
-    def fillna(self, value, **kwargs):
+    def fillna(
+        self, value, limit=None, inplace: bool = False, downcast=None
+    ) -> List[Block]:
         # TODO(EA2D): if we operated on array_values, TDA.fillna would handle
         #  raising here.
         if is_integer(value):
@@ -2238,7 +2244,7 @@ def fillna(self, value, **kwargs):
                 "longer supported.  To obtain the old behavior, pass "
                 "`pd.Timedelta(seconds=n)` instead."
             )
-        return super().fillna(value, **kwargs)
+        return super().fillna(value, limit=limit, inplace=inplace, downcast=downcast)
 
 
 class ObjectBlock(Block):
@@ -2450,7 +2456,9 @@ def get_block_type(values, dtype: Optional[Dtype] = None):
     return cls
 
 
-def make_block(values, placement, klass=None, ndim=None, dtype: Optional[Dtype] = None):
+def make_block(
+    values, placement, klass=None, ndim=None, dtype: Optional[Dtype] = None
+) -> Block:
     # Ensure that we don't allow PandasArray / PandasDtype in internals.
     # For now, blocks should be backed by ndarrays when possible.
     if isinstance(values, ABCPandasArray):
@@ -2477,7 +2485,7 @@ def make_block(values, placement, klass=None, ndim=None, dtype: Optional[Dtype]
 # -----------------------------------------------------------------
 
 
-def extend_blocks(result, blocks=None):
+def extend_blocks(result, blocks=None) -> List[Block]:
     """ return a new extended blocks, given the result """
     if blocks is None:
         blocks = []