TYP: pandas/core/frame.py (easy: Axis/Level) (pandas-dev#38441)

* typing (Axis/Level)

* review comment

* replace List[Axis]->Sequence[Axis] in reorder_levels

Author: arw2019

pandas/core/frame.py

@@ -4050,7 +4050,7 @@ def _reindex_index(
         new_index,
         method,
         copy: bool,
-        level,
+        level: Level,
         fill_value=np.nan,
         limit=None,
         tolerance=None,
@@ -4070,7 +4070,7 @@ def _reindex_columns(
         new_columns,
         method,
         copy: bool,
-        level,
+        level: Level,
         fill_value=None,
         limit=None,
         tolerance=None,
@@ -4110,14 +4110,14 @@ def align(
         self,
         other,
         join: str = "outer",
-        axis=None,
-        level=None,
+        axis: Optional[Axis] = None,
+        level: Optional[Level] = None,
         copy: bool = True,
         fill_value=None,
         method: Optional[str] = None,
         limit=None,
-        fill_axis=0,
-        broadcast_axis=None,
+        fill_axis: Axis = 0,
+        broadcast_axis: Optional[Axis] = None,
     ) -> DataFrame:
         return super().align(
             other,
@@ -4198,10 +4198,10 @@ def reindex(self, *args, **kwargs) -> DataFrame:
     def drop(
         self,
         labels=None,
-        axis=0,
+        axis: Axis = 0,
         index=None,
         columns=None,
-        level=None,
+        level: Optional[Level] = None,
         inplace: bool = False,
         errors: str = "raise",
     ):
@@ -4474,7 +4474,7 @@ def fillna(
         self,
         value=None,
         method: Optional[str] = None,
-        axis=None,
+        axis: Optional[Axis] = None,
         inplace: bool = False,
         limit=None,
         downcast=None,
@@ -4587,7 +4587,7 @@ def _replace_columnwise(
 
     @doc(NDFrame.shift, klass=_shared_doc_kwargs["klass"])
     def shift(
-        self, periods=1, freq=None, axis=0, fill_value=lib.no_default
+        self, periods=1, freq=None, axis: Axis = 0, fill_value=lib.no_default
     ) -> DataFrame:
         axis = self._get_axis_number(axis)
 
@@ -5074,7 +5074,12 @@ def notnull(self) -> DataFrame:
         return ~self.isna()
 
     def dropna(
-        self, axis=0, how: str = "any", thresh=None, subset=None, inplace: bool = False
+        self,
+        axis: Axis = 0,
+        how: str = "any",
+        thresh=None,
+        subset=None,
+        inplace: bool = False,
     ):
         """
         Remove missing values.
@@ -5454,7 +5459,7 @@ def f(vals):
     def sort_values(  # type: ignore[override]
         self,
         by,
-        axis=0,
+        axis: Axis = 0,
         ascending=True,
         inplace: bool = False,
         kind: str = "quicksort",
@@ -5514,8 +5519,8 @@ def sort_values(  # type: ignore[override]
 
     def sort_index(
         self,
-        axis=0,
-        level=None,
+        axis: Axis = 0,
+        level: Optional[Level] = None,
         ascending: bool = True,
         inplace: bool = False,
         kind: str = "quicksort",
@@ -5932,7 +5937,7 @@ def nsmallest(self, n, columns, keep: str = "first") -> DataFrame:
             self, n=n, keep=keep, columns=columns
         ).nsmallest()
 
-    def swaplevel(self, i=-2, j=-1, axis=0) -> DataFrame:
+    def swaplevel(self, i: Axis = -2, j: Axis = -1, axis: Axis = 0) -> DataFrame:
         """
         Swap levels i and j in a MultiIndex on a particular axis.
 
@@ -5963,7 +5968,7 @@ def swaplevel(self, i=-2, j=-1, axis=0) -> DataFrame:
             result.columns = result.columns.swaplevel(i, j)
         return result
 
-    def reorder_levels(self, order, axis=0) -> DataFrame:
+    def reorder_levels(self, order: Sequence[Axis], axis: Axis = 0) -> DataFrame:
         """
         Rearrange index levels using input order. May not drop or duplicate levels.
 
@@ -6726,8 +6731,8 @@ def update(
     def groupby(
         self,
         by=None,
-        axis=0,
-        level=None,
+        axis: Axis = 0,
+        level: Optional[Level] = None,
         as_index: bool = True,
         sort: bool = True,
         group_keys: bool = True,
@@ -7080,7 +7085,7 @@ def pivot_table(
             observed=observed,
         )
 
-    def stack(self, level=-1, dropna: bool = True):
+    def stack(self, level: Level = -1, dropna: bool = True):
         """
         Stack the prescribed level(s) from columns to index.
 
@@ -7399,7 +7404,7 @@ def melt(
         value_vars=None,
         var_name=None,
         value_name="value",
-        col_level=None,
+        col_level: Optional[Level] = None,
         ignore_index=True,
     ) -> DataFrame:
 
@@ -7607,7 +7612,7 @@ def _gotitem(
         see_also=_agg_summary_and_see_also_doc,
         examples=_agg_examples_doc,
     )
-    def aggregate(self, func=None, axis=0, *args, **kwargs):
+    def aggregate(self, func=None, axis: Axis = 0, *args, **kwargs):
         axis = self._get_axis_number(axis)
 
         relabeling, func, columns, order = reconstruct_func(func, **kwargs)
@@ -7638,7 +7643,7 @@ def aggregate(self, func=None, axis=0, *args, **kwargs):
 
         return result
 
-    def _aggregate(self, arg, axis=0, *args, **kwargs):
+    def _aggregate(self, arg, axis: Axis = 0, *args, **kwargs):
         if axis == 1:
             # NDFrame.aggregate returns a tuple, and we need to transpose
             # only result
@@ -7661,7 +7666,9 @@ def transform(
         assert isinstance(result, DataFrame)
         return result
 
-    def apply(self, func, axis=0, raw: bool = False, result_type=None, args=(), **kwds):
+    def apply(
+        self, func, axis: Axis = 0, raw: bool = False, result_type=None, args=(), **kwds
+    ):
         """
         Apply a function along an axis of the DataFrame.
 
@@ -8578,7 +8585,7 @@ def cov(
 
         return self._constructor(base_cov, index=idx, columns=cols)
 
-    def corrwith(self, other, axis=0, drop=False, method="pearson") -> Series:
+    def corrwith(self, other, axis: Axis = 0, drop=False, method="pearson") -> Series:
         """
         Compute pairwise correlation.
 
@@ -8674,7 +8681,9 @@ def c(x):
     # ----------------------------------------------------------------------
     # ndarray-like stats methods
 
-    def count(self, axis=0, level=None, numeric_only: bool = False):
+    def count(
+        self, axis: Axis = 0, level: Optional[Level] = None, numeric_only: bool = False
+    ):
         """
         Count non-NA cells for each column or row.
 
@@ -8778,7 +8787,7 @@ def count(self, axis=0, level=None, numeric_only: bool = False):
 
         return result.astype("int64")
 
-    def _count_level(self, level, axis=0, numeric_only=False):
+    def _count_level(self, level: Level, axis: Axis = 0, numeric_only=False):
         if numeric_only:
             frame = self._get_numeric_data()
         else:
@@ -8828,7 +8837,7 @@ def _reduce(
         op,
         name: str,
         *,
-        axis=0,
+        axis: Axis = 0,
         skipna: bool = True,
         numeric_only: Optional[bool] = None,
         filter_type=None,
@@ -8936,7 +8945,7 @@ def _get_data() -> DataFrame:
         result = self._constructor_sliced(result, index=labels)
         return result
 
-    def nunique(self, axis=0, dropna: bool = True) -> Series:
+    def nunique(self, axis: Axis = 0, dropna: bool = True) -> Series:
         """
         Count distinct observations over requested axis.
 
@@ -8976,7 +8985,7 @@ def nunique(self, axis=0, dropna: bool = True) -> Series:
         """
         return self.apply(Series.nunique, axis=axis, dropna=dropna)
 
-    def idxmin(self, axis=0, skipna: bool = True) -> Series:
+    def idxmin(self, axis: Axis = 0, skipna: bool = True) -> Series:
         """
         Return index of first occurrence of minimum over requested axis.
 
@@ -9053,7 +9062,7 @@ def idxmin(self, axis=0, skipna: bool = True) -> Series:
         result = [index[i] if i >= 0 else np.nan for i in indices]
         return self._constructor_sliced(result, index=self._get_agg_axis(axis))
 
-    def idxmax(self, axis=0, skipna: bool = True) -> Series:
+    def idxmax(self, axis: Axis = 0, skipna: bool = True) -> Series:
         """
         Return index of first occurrence of maximum over requested axis.
 
@@ -9142,7 +9151,7 @@ def _get_agg_axis(self, axis_num: int) -> Index:
             raise ValueError(f"Axis must be 0 or 1 (got {repr(axis_num)})")
 
     def mode(
-        self, axis=0, numeric_only: bool = False, dropna: bool = True
+        self, axis: Axis = 0, numeric_only: bool = False, dropna: bool = True
     ) -> DataFrame:
         """
         Get the mode(s) of each element along the selected axis.
@@ -9231,7 +9240,11 @@ def f(s):
         return data.apply(f, axis=axis)
 
     def quantile(
-        self, q=0.5, axis=0, numeric_only: bool = True, interpolation: str = "linear"
+        self,
+        q=0.5,
+        axis: Axis = 0,
+        numeric_only: bool = True,
+        interpolation: str = "linear",
     ):
         """
         Return values at the given quantile over requested axis.