Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 4 commits into from
Dec 13, 2020
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
77 changes: 45 additions & 32 deletions pandas/core/frame.py
Original file line number Diff line number Diff line change
Expand Up @@ -4050,7 +4050,7 @@ def _reindex_index(
new_index,
method,
copy: bool,
level,
level: Level,
fill_value=np.nan,
limit=None,
tolerance=None,
Expand All @@ -4070,7 +4070,7 @@ def _reindex_columns(
new_columns,
method,
copy: bool,
level,
level: Level,
fill_value=None,
limit=None,
tolerance=None,
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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",
):
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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)

Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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",
Expand Down Expand Up @@ -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",
Expand Down Expand Up @@ -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.

Expand Down Expand Up @@ -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:
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This has to be Sequence not List because with List mypy complains

pandas/core/generic.py:8689: error: Argument 1 to "reorder_levels" of "DataFrame" has incompatible type "List[int]"; expected "List[Union[str, int]]"  [arg-type]
pandas/core/generic.py:8689: note: "List" is invariant -- see http://mypy.readthedocs.io/en/latest/common_issues.html#variance
pandas/core/generic.py:8689: note: Consider using "Sequence" instead, which is covariant
pandas/core/generic.py:8691: error: Argument 1 to "reorder_levels" of "DataFrame" has incompatible type "List[int]"; expected "List[Union[str, int]]"  [arg-type]
pandas/core/generic.py:8691: note: "List" is invariant -- see http://mypy.readthedocs.io/en/latest/common_issues.html#variance
pandas/core/generic.py:8691: note: Consider using "Sequence" instead, which is covariant

"""
Rearrange index levels using input order. May not drop or duplicate levels.

Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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.

Expand Down Expand Up @@ -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:

Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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
Expand All @@ -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.

Expand Down Expand Up @@ -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.

Expand Down Expand Up @@ -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.

Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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.

Expand Down Expand Up @@ -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.

Expand Down Expand Up @@ -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.

Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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.
Expand Down