Skip to content

CLN: Replace FrameOrSeriesUnion annotation by DataFrame | Series #41955

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 11 commits into from
Jul 4, 2021
6 changes: 0 additions & 6 deletions pandas/_typing.py
Original file line number Diff line number Diff line change
Expand Up @@ -101,12 +101,6 @@
]
Timezone = Union[str, tzinfo]

# FrameOrSeriesUnion means either a DataFrame or a Series. E.g.
# `def func(a: FrameOrSeriesUnion) -> FrameOrSeriesUnion: ...` means that if a Series
# is passed in, either a Series or DataFrame is returned, and if a DataFrame is passed
# in, either a DataFrame or a Series is returned.
FrameOrSeriesUnion = Union["DataFrame", "Series"]

# FrameOrSeries is stricter and ensures that the same subclass of NDFrame always is
# used. E.g. `def func(a: FrameOrSeries) -> FrameOrSeries: ...` means that if a
# Series is passed into a function, a Series is always returned and if a DataFrame is
Expand Down
3 changes: 1 addition & 2 deletions pandas/core/algorithms.py
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,6 @@
AnyArrayLike,
ArrayLike,
DtypeObj,
FrameOrSeriesUnion,
Scalar,
)
from pandas.util._decorators import doc
Expand Down Expand Up @@ -1211,7 +1210,7 @@ def __init__(self, obj, n: int, keep: str):
if self.keep not in ("first", "last", "all"):
raise ValueError('keep must be either "first", "last" or "all"')

def compute(self, method: str) -> FrameOrSeriesUnion:
def compute(self, method: str) -> DataFrame | Series:
raise NotImplementedError

def nlargest(self):
Expand Down
39 changes: 19 additions & 20 deletions pandas/core/apply.py
Original file line number Diff line number Diff line change
Expand Up @@ -25,7 +25,6 @@
AggObjType,
Axis,
FrameOrSeries,
FrameOrSeriesUnion,
)
from pandas.util._decorators import cache_readonly

Expand Down Expand Up @@ -137,10 +136,10 @@ def f(x):
self.f: AggFuncType = f

@abc.abstractmethod
def apply(self) -> FrameOrSeriesUnion:
def apply(self) -> DataFrame | Series:
pass

def agg(self) -> FrameOrSeriesUnion | None:
def agg(self) -> DataFrame | Series | None:
"""
Provide an implementation for the aggregators.

Expand Down Expand Up @@ -171,7 +170,7 @@ def agg(self) -> FrameOrSeriesUnion | None:
# caller can react
return None

def transform(self) -> FrameOrSeriesUnion:
def transform(self) -> DataFrame | Series:
"""
Transform a DataFrame or Series.

Expand Down Expand Up @@ -252,7 +251,7 @@ def transform_dict_like(self, func):

func = self.normalize_dictlike_arg("transform", obj, func)

results: dict[Hashable, FrameOrSeriesUnion] = {}
results: dict[Hashable, DataFrame | Series] = {}
failed_names = []
all_type_errors = True
for name, how in func.items():
Expand Down Expand Up @@ -283,7 +282,7 @@ def transform_dict_like(self, func):
)
return concat(results, axis=1)

def transform_str_or_callable(self, func) -> FrameOrSeriesUnion:
def transform_str_or_callable(self, func) -> DataFrame | Series:
"""
Compute transform in the case of a string or callable func
"""
Expand All @@ -305,7 +304,7 @@ def transform_str_or_callable(self, func) -> FrameOrSeriesUnion:
except Exception:
return func(obj, *args, **kwargs)

def agg_list_like(self) -> FrameOrSeriesUnion:
def agg_list_like(self) -> DataFrame | Series:
"""
Compute aggregation in the case of a list-like argument.

Expand Down Expand Up @@ -399,7 +398,7 @@ def agg_list_like(self) -> FrameOrSeriesUnion:
)
return concatenated.reindex(full_ordered_index, copy=False)

def agg_dict_like(self) -> FrameOrSeriesUnion:
def agg_dict_like(self) -> DataFrame | Series:
"""
Compute aggregation in the case of a dict-like argument.

Expand Down Expand Up @@ -467,7 +466,7 @@ def agg_dict_like(self) -> FrameOrSeriesUnion:

return result

def apply_str(self) -> FrameOrSeriesUnion:
def apply_str(self) -> DataFrame | Series:
"""
Compute apply in case of a string.

Expand All @@ -492,7 +491,7 @@ def apply_str(self) -> FrameOrSeriesUnion:
raise ValueError(f"Operation {f} does not support axis=1")
return self._try_aggregate_string_function(obj, f, *self.args, **self.kwargs)

def apply_multiple(self) -> FrameOrSeriesUnion:
def apply_multiple(self) -> DataFrame | Series:
"""
Compute apply in case of a list-like or dict-like.

Expand All @@ -504,7 +503,7 @@ def apply_multiple(self) -> FrameOrSeriesUnion:
return self.obj.aggregate(self.f, self.axis, *self.args, **self.kwargs)

def normalize_dictlike_arg(
self, how: str, obj: FrameOrSeriesUnion, func: AggFuncTypeDict
self, how: str, obj: DataFrame | Series, func: AggFuncTypeDict
) -> AggFuncTypeDict:
"""
Handler for dict-like argument.
Expand Down Expand Up @@ -617,7 +616,7 @@ def series_generator(self) -> Iterator[Series]:
@abc.abstractmethod
def wrap_results_for_axis(
self, results: ResType, res_index: Index
) -> FrameOrSeriesUnion:
) -> DataFrame | Series:
pass

# ---------------------------------------------------------------
Expand All @@ -638,7 +637,7 @@ def values(self):
def dtypes(self) -> Series:
return self.obj.dtypes

def apply(self) -> FrameOrSeriesUnion:
def apply(self) -> DataFrame | Series:
"""compute the results"""
# dispatch to agg
if is_list_like(self.f):
Expand Down Expand Up @@ -812,7 +811,7 @@ def apply_series_generator(self) -> tuple[ResType, Index]:

return results, res_index

def wrap_results(self, results: ResType, res_index: Index) -> FrameOrSeriesUnion:
def wrap_results(self, results: ResType, res_index: Index) -> DataFrame | Series:
from pandas import Series

# see if we can infer the results
Expand All @@ -835,7 +834,7 @@ def wrap_results(self, results: ResType, res_index: Index) -> FrameOrSeriesUnion

return result

def apply_str(self) -> FrameOrSeriesUnion:
def apply_str(self) -> DataFrame | Series:
# Caller is responsible for checking isinstance(self.f, str)
# TODO: GH#39993 - Avoid special-casing by replacing with lambda
if self.f == "size":
Expand Down Expand Up @@ -866,7 +865,7 @@ def result_columns(self) -> Index:

def wrap_results_for_axis(
self, results: ResType, res_index: Index
) -> FrameOrSeriesUnion:
) -> DataFrame | Series:
"""return the results for the rows"""

if self.result_type == "reduce":
Expand Down Expand Up @@ -949,9 +948,9 @@ def result_columns(self) -> Index:

def wrap_results_for_axis(
self, results: ResType, res_index: Index
) -> FrameOrSeriesUnion:
) -> DataFrame | Series:
"""return the results for the columns"""
result: FrameOrSeriesUnion
result: DataFrame | Series

# we have requested to expand
if self.result_type == "expand":
Expand Down Expand Up @@ -1005,7 +1004,7 @@ def __init__(
kwargs=kwargs,
)

def apply(self) -> FrameOrSeriesUnion:
def apply(self) -> DataFrame | Series:
obj = self.obj

if len(obj) == 0:
Expand Down Expand Up @@ -1056,7 +1055,7 @@ def apply_empty_result(self) -> Series:
obj, method="apply"
)

def apply_standard(self) -> FrameOrSeriesUnion:
def apply_standard(self) -> DataFrame | Series:
f = self.f
obj = self.obj

Expand Down
5 changes: 2 additions & 3 deletions pandas/core/describe.py
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,6 @@
from pandas._libs.tslibs import Timestamp
from pandas._typing import (
FrameOrSeries,
FrameOrSeriesUnion,
Hashable,
)
from pandas.util._validators import validate_percentile
Expand Down Expand Up @@ -107,12 +106,12 @@ class NDFrameDescriberAbstract(ABC):
Whether to treat datetime dtypes as numeric.
"""

def __init__(self, obj: FrameOrSeriesUnion, datetime_is_numeric: bool):
def __init__(self, obj: DataFrame | Series, datetime_is_numeric: bool):
self.obj = obj
self.datetime_is_numeric = datetime_is_numeric

@abstractmethod
def describe(self, percentiles: Sequence[float]) -> FrameOrSeriesUnion:
def describe(self, percentiles: Sequence[float]) -> DataFrame | Series:
"""Do describe either series or dataframe.

Parameters
Expand Down
23 changes: 11 additions & 12 deletions pandas/core/frame.py
Original file line number Diff line number Diff line change
Expand Up @@ -58,7 +58,6 @@
FillnaOptions,
FloatFormatType,
FormattersType,
FrameOrSeriesUnion,
Frequency,
IndexKeyFunc,
IndexLabel,
Expand Down Expand Up @@ -1362,7 +1361,7 @@ def dot(self, other: Series) -> Series:
def dot(self, other: DataFrame | Index | ArrayLike) -> DataFrame:
...

def dot(self, other: AnyArrayLike | FrameOrSeriesUnion) -> FrameOrSeriesUnion:
def dot(self, other: AnyArrayLike | DataFrame | Series) -> DataFrame | Series:
Copy link
Member

Choose a reason for hiding this comment

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

A Series is also AnyArrayLike

Suggested change
def dot(self, other: AnyArrayLike | DataFrame | Series) -> DataFrame | Series:
def dot(self, other: AnyArrayLike | DataFrame) -> DataFrame | Series:

Copy link
Member

@MarcoGorelli MarcoGorelli Jul 3, 2021

Choose a reason for hiding this comment

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

ooh, we should add a check for this

EDIT

#42359

Copy link
Member

Choose a reason for hiding this comment

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

ok. we can merge this and let the check fix this

"""
Compute the matrix multiplication between the DataFrame and other.

Expand Down Expand Up @@ -1478,13 +1477,13 @@ def __matmul__(self, other: Series) -> Series:

@overload
def __matmul__(
self, other: AnyArrayLike | FrameOrSeriesUnion
) -> FrameOrSeriesUnion:
self, other: AnyArrayLike | DataFrame | Series
Copy link
Member

Choose a reason for hiding this comment

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

same (and already included in previous overload anyway)

) -> DataFrame | Series:
...

def __matmul__(
self, other: AnyArrayLike | FrameOrSeriesUnion
) -> FrameOrSeriesUnion:
self, other: AnyArrayLike | DataFrame | Series
Copy link
Member

Choose a reason for hiding this comment

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

same

) -> DataFrame | Series:
"""
Matrix multiplication using binary `@` operator in Python>=3.5.
"""
Expand Down Expand Up @@ -8401,8 +8400,8 @@ def _gotitem(
self,
key: IndexLabel,
ndim: int,
subset: FrameOrSeriesUnion | None = None,
) -> FrameOrSeriesUnion:
subset: DataFrame | Series | None = None,
) -> DataFrame | Series:
"""
Sub-classes to define. Return a sliced object.

Expand Down Expand Up @@ -8931,7 +8930,7 @@ def append(

def join(
self,
other: FrameOrSeriesUnion,
other: DataFrame | Series,
on: IndexLabel | None = None,
how: str = "left",
lsuffix: str = "",
Expand Down Expand Up @@ -9061,7 +9060,7 @@ def join(

def _join_compat(
self,
other: FrameOrSeriesUnion,
other: DataFrame | Series,
on: IndexLabel | None = None,
how: str = "left",
lsuffix: str = "",
Expand Down Expand Up @@ -9131,7 +9130,7 @@ def _join_compat(
@Appender(_merge_doc, indents=2)
def merge(
self,
right: FrameOrSeriesUnion,
right: DataFrame | Series,
how: str = "inner",
on: IndexLabel | None = None,
left_on: IndexLabel | None = None,
Expand Down Expand Up @@ -10724,7 +10723,7 @@ def _from_nested_dict(data) -> collections.defaultdict:
return new_data


def _reindex_for_setitem(value: FrameOrSeriesUnion, index: Index) -> ArrayLike:
def _reindex_for_setitem(value: DataFrame | Series, index: Index) -> ArrayLike:
# reindex if necessary

if value.index.equals(index) or not len(index):
Expand Down
7 changes: 3 additions & 4 deletions pandas/core/groupby/generic.py
Original file line number Diff line number Diff line change
Expand Up @@ -33,7 +33,6 @@
from pandas._typing import (
ArrayLike,
FrameOrSeries,
FrameOrSeriesUnion,
Manager2D,
)
from pandas.util._decorators import (
Expand Down Expand Up @@ -296,7 +295,7 @@ def _aggregate_multiple_funcs(self, arg) -> DataFrame:

arg = zip(columns, arg)

results: dict[base.OutputKey, FrameOrSeriesUnion] = {}
results: dict[base.OutputKey, DataFrame | Series] = {}
for idx, (name, func) in enumerate(arg):

key = base.OutputKey(label=name, position=idx)
Expand Down Expand Up @@ -422,7 +421,7 @@ def _wrap_applied_output(
keys: Index,
values: list[Any] | None,
not_indexed_same: bool = False,
) -> FrameOrSeriesUnion:
) -> DataFrame | Series:
"""
Wrap the output of SeriesGroupBy.apply into the expected result.

Expand Down Expand Up @@ -1191,7 +1190,7 @@ def _wrap_applied_output_series(
not_indexed_same: bool,
first_not_none,
key_index,
) -> FrameOrSeriesUnion:
) -> DataFrame | Series:
# this is to silence a DeprecationWarning
# TODO: Remove when default dtype of empty Series is object
kwargs = first_not_none._construct_axes_dict()
Expand Down
9 changes: 4 additions & 5 deletions pandas/core/groupby/groupby.py
Original file line number Diff line number Diff line change
Expand Up @@ -45,7 +45,6 @@ class providing the base-class of operations.
ArrayLike,
F,
FrameOrSeries,
FrameOrSeriesUnion,
IndexLabel,
Scalar,
T,
Expand Down Expand Up @@ -728,7 +727,7 @@ def pipe(
plot = property(GroupByPlot)

@final
def get_group(self, name, obj=None) -> FrameOrSeriesUnion:
def get_group(self, name, obj=None) -> DataFrame | Series:
"""
Construct DataFrame from group with provided name.

Expand Down Expand Up @@ -1267,8 +1266,8 @@ def f(g):

@final
def _python_apply_general(
self, f: F, data: FrameOrSeriesUnion
) -> FrameOrSeriesUnion:
self, f: F, data: DataFrame | Series
) -> DataFrame | Series:
"""
Apply function f in python space

Expand Down Expand Up @@ -1786,7 +1785,7 @@ def sem(self, ddof: int = 1):
@final
@Substitution(name="groupby")
@Appender(_common_see_also)
def size(self) -> FrameOrSeriesUnion:
def size(self) -> DataFrame | Series:
"""
Compute group sizes.

Expand Down
Loading