TYP: pandas/core/frame.py

pandas/core/frame.py

@@ -22,12 +22,14 @@
     TYPE_CHECKING,
     Any,
     AnyStr,
+    Callable,
     Dict,
     FrozenSet,
     Hashable,
     Iterable,
     Iterator,
     List,
+    Mapping,
     Optional,
     Sequence,
     Set,
@@ -48,6 +50,7 @@
 from pandas._libs.lib import no_default
 from pandas._typing import (
     AggFuncType,
+    AnyArrayLike,
     ArrayLike,
     Axes,
     Axis,
@@ -62,6 +65,7 @@
     Label,
     Level,
     Renamer,
+    Scalar,
     StorageOptions,
     ValueKeyFunc,
 )
@@ -1069,7 +1073,9 @@ def iterrows(self) -> Iterable[Tuple[Label, Series]]:
             s = klass(v, index=columns, name=k)
             yield k, s
 
-    def itertuples(self, index: bool = True, name: Optional[str] = "Pandas"):
+    def itertuples(
+        self, index: bool = True, name: Optional[str] = "Pandas"
+    ) -> Iterable[Tuple]:
         """
         Iterate over DataFrame rows as namedtuples.
 
@@ -1159,7 +1165,7 @@ def __len__(self) -> int:
         """
         return len(self.index)
 
-    def dot(self, other):
+    def dot(self, other: Union[AnyArrayLike, FrameOrSeriesUnion]) -> FrameOrSeriesUnion:
         """
         Compute the matrix multiplication between the DataFrame and other.
 
@@ -1269,13 +1275,17 @@ def dot(self, other):
         else:  # pragma: no cover
             raise TypeError(f"unsupported type: {type(other)}")
 
-    def __matmul__(self, other):
+    def __matmul__(
+        self, other: Union[AnyArrayLike, FrameOrSeriesUnion]
+    ) -> FrameOrSeriesUnion:
         """
         Matrix multiplication using binary `@` operator in Python>=3.5.
         """
         return self.dot(other)
 
-    def __rmatmul__(self, other):
+    def __rmatmul__(
+        self, other: Union[AnyArrayLike, FrameOrSeriesUnion]
+    ) -> FrameOrSeriesUnion:
         """
         Matrix multiplication using binary `@` operator in Python>=3.5.
         """
@@ -1292,7 +1302,13 @@ def __rmatmul__(self, other):
     # IO methods (to / from other formats)
 
     @classmethod
-    def from_dict(cls, data, orient="columns", dtype=None, columns=None) -> DataFrame:
+    def from_dict(
+        cls,
+        data,
+        orient: str = "columns",
+        dtype: Optional[Dtype] = None,
+        columns: Optional[List[Label]] = None,
+    ) -> DataFrame:
         """
         Construct DataFrame from dict of array-like or dicts.
 
@@ -1371,7 +1387,10 @@ def from_dict(cls, data, orient="columns", dtype=None, columns=None) -> DataFram
         return cls(data, index=index, columns=columns, dtype=dtype)
 
     def to_numpy(
-        self, dtype=None, copy: bool = False, na_value=lib.no_default
+        self,
+        dtype: Optional[Dtype] = None,
+        copy: bool = False,
+        na_value: Scalar = lib.no_default,
     ) -> np.ndarray:
         """
         Convert the DataFrame to a NumPy array.
@@ -1438,7 +1457,7 @@ def to_numpy(
 
         return result
 
-    def to_dict(self, orient: str = "dict", into=dict):
+    def to_dict(self, orient: str = "dict", into=dict) -> Union[Dict, List, Mapping]:
         """
         Convert the DataFrame to a dictionary.
 
@@ -1727,7 +1746,7 @@ def from_records(
         exclude=None,
         columns=None,
         coerce_float: bool = False,
-        nrows=None,
+        nrows: Optional[bool] = None,
     ) -> DataFrame:
         """
         Convert structured or record ndarray to DataFrame.
@@ -2710,7 +2729,7 @@ def info(
             show_counts=show_counts,
         )
 
-    def memory_usage(self, index=True, deep=False) -> Series:
+    def memory_usage(self, index: bool = True, deep: bool = False) -> Series:
         """
         Return the memory usage of each column in bytes.
 
@@ -2933,7 +2952,7 @@ def T(self) -> DataFrame:
     # ----------------------------------------------------------------------
     # Indexing Methods
 
-    def _ixs(self, i: int, axis: int = 0):
+    def _ixs(self, i: int, axis: Axis = 0):
         """
         Parameters
         ----------
@@ -3110,7 +3129,7 @@ def _getitem_multilevel(self, key):
             # loc is neither a slice nor ndarray, so must be an int
             return self._ixs(loc, axis=1)
 
-    def _get_value(self, index, col, takeable: bool = False):
+    def _get_value(self, index, col, takeable: bool = False) -> Scalar:
         """
         Quickly retrieve single value at passed column and index.
 
@@ -3264,7 +3283,7 @@ def _iset_item(self, loc: int, value):
         if len(self):
             self._check_setitem_copy()
 
-    def _set_item(self, key, value):
+    def _set_item(self, key, value) -> None:
         """
         Add series to DataFrame in specified column.
 
@@ -3289,7 +3308,9 @@ def _set_item(self, key, value):
 
         self._set_item_mgr(key, value)
 
-    def _set_value(self, index, col, value, takeable: bool = False):
+    def _set_value(
+        self, index: int, col, value: Scalar, takeable: bool = False
+    ) -> None:
         """
         Put single value at passed column and index.
 
@@ -3322,7 +3343,7 @@ def _set_value(self, index, col, value, takeable: bool = False):
                 self.loc[index, col] = value
             self._item_cache.pop(col, None)
 
-    def _ensure_valid_index(self, value):
+    def _ensure_valid_index(self, value) -> None:
         """
         Ensure that if we don't have an index, that we can create one from the
         passed value.
@@ -3357,7 +3378,7 @@ def _box_col_values(self, values, loc: int) -> Series:
     # ----------------------------------------------------------------------
     # Unsorted
 
-    def query(self, expr: str, inplace: bool = False, **kwargs):
+    def query(self, expr: str, inplace: bool = False, **kwargs) -> Optional[DataFrame]:
         """
         Query the columns of a DataFrame with a boolean expression.
 
@@ -3518,10 +3539,13 @@ def query(self, expr: str, inplace: bool = False, **kwargs):
 
         if inplace:
             self._update_inplace(result)
+            return None
         else:
             return result
 
-    def eval(self, expr: str, inplace: bool = False, **kwargs):
+    def eval(
+        self, expr: str, inplace: bool = False, **kwargs
+    ) -> Optional[Union[AnyArrayLike, DataFrame, Scalar]]:
         """
         Evaluate a string describing operations on DataFrame columns.
 
@@ -3946,15 +3970,17 @@ def _sanitize_column(self, value):
         return value
 
     @property
-    def _series(self):
+    def _series(self) -> "Dict[int, Series]":
         return {
             item: Series(
                 self._mgr.iget(idx), index=self.index, name=item, fastpath=True
             )
             for idx, item in enumerate(self.columns)
         }
 
-    def lookup(self, row_labels, col_labels) -> np.ndarray:
+    def lookup(
+        self, row_labels: Sequence[Label], col_labels: Sequence[Label]
+    ) -> np.ndarray:
         """
         Label-based "fancy indexing" function for DataFrame.
         Given equal-length arrays of row and column labels, return an
@@ -6012,7 +6038,7 @@ def _arith_method(self, other, op):
 
     _logical_method = _arith_method
 
-    def _dispatch_frame_op(self, right, func, axis: Optional[int] = None):
+    def _dispatch_frame_op(self, right, func: Callable, axis: Optional[int] = None):
         """
         Evaluate the frame operation func(left, right) by evaluating
         column-by-column, dispatching to the Series implementation.
@@ -6093,7 +6119,7 @@ def _arith_op(left, right):
         new_data = self._dispatch_frame_op(other, _arith_op)
         return new_data
 
-    def _construct_result(self, result) -> DataFrame:
+    def _construct_result(self, result: DataFrame) -> DataFrame:
         """
         Wrap the result of an arithmetic, comparison, or logical operation.
 
@@ -7323,7 +7349,7 @@ def explode(
 
         return result
 
-    def unstack(self, level=-1, fill_value=None):
+    def unstack(self, level: Level = -1, fill_value=None):
         """
         Pivot a level of the (necessarily hierarchical) index labels.
 
@@ -7394,7 +7420,7 @@ def melt(
         var_name=None,
         value_name="value",
         col_level: Optional[Level] = None,
-        ignore_index=True,
+        ignore_index: bool = True,
     ) -> DataFrame:
 
         return melt(
@@ -7802,7 +7828,7 @@ def apply(
         )
         return op.get_result()
 
-    def applymap(self, func, na_action: Optional[str] = None) -> DataFrame:
+    def applymap(self, func: Callable, na_action: Optional[str] = None) -> DataFrame:
         """
         Apply a function to a Dataframe elementwise.
 
@@ -8253,7 +8279,9 @@ def merge(
             validate=validate,
         )
 
-    def round(self, decimals=0, *args, **kwargs) -> DataFrame:
+    def round(
+        self, decimals: Union[int, Dict, Series] = 0, *args, **kwargs
+    ) -> DataFrame:
         """
         Round a DataFrame to a variable number of decimal places.
 
@@ -8367,7 +8395,9 @@ def _series_round(s, decimals):
     # ----------------------------------------------------------------------
     # Statistical methods, etc.
 
-    def corr(self, method="pearson", min_periods=1) -> DataFrame:
+    def corr(
+        self, method: Union[str, Callable] = "pearson", min_periods: int = 1
+    ) -> DataFrame:
         """
         Compute pairwise correlation of columns, excluding NA/null values.
 
@@ -9331,7 +9361,11 @@ def quantile(
         return result
 
     def to_timestamp(
-        self, freq=None, how: str = "start", axis: Axis = 0, copy: bool = True
+        self,
+        freq: Optional[str] = None,
+        how: str = "start",
+        axis: Axis = 0,
+        copy: bool = True,
     ) -> DataFrame:
         """
         Cast to DatetimeIndex of timestamps, at *beginning* of period.
@@ -9364,7 +9398,9 @@ def to_timestamp(
         setattr(new_obj, axis_name, new_ax)
         return new_obj
 
-    def to_period(self, freq=None, axis: Axis = 0, copy: bool = True) -> DataFrame:
+    def to_period(
+        self, freq: Optional[str] = None, axis: Axis = 0, copy: bool = True
+    ) -> DataFrame:
         """
         Convert DataFrame from DatetimeIndex to PeriodIndex.
 
@@ -9396,7 +9432,7 @@ def to_period(self, freq=None, axis: Axis = 0, copy: bool = True) -> DataFrame:
         setattr(new_obj, axis_name, new_ax)
         return new_obj
 
-    def isin(self, values) -> DataFrame:
+    def isin(self, values: Union[FrameOrSeriesUnion, Dict, Iterable]) -> DataFrame:
         """
         Whether each element in the DataFrame is contained in values.