125 lines
2.9 KiB
Python
125 lines
2.9 KiB
Python
"""
|
|
Boilerplate functions used in defining binary operations.
|
|
"""
|
|
from functools import wraps
|
|
from typing import Callable
|
|
|
|
from pandas._libs.lib import item_from_zerodim
|
|
from pandas._typing import F
|
|
|
|
from pandas.core.dtypes.generic import ABCDataFrame, ABCIndexClass, ABCSeries
|
|
|
|
|
|
def unpack_zerodim_and_defer(name: str) -> Callable[[F], F]:
|
|
"""
|
|
Boilerplate for pandas conventions in arithmetic and comparison methods.
|
|
|
|
Parameters
|
|
----------
|
|
name : str
|
|
|
|
Returns
|
|
-------
|
|
decorator
|
|
"""
|
|
|
|
def wrapper(method: F) -> F:
|
|
return _unpack_zerodim_and_defer(method, name)
|
|
|
|
return wrapper
|
|
|
|
|
|
def _unpack_zerodim_and_defer(method, name: str):
|
|
"""
|
|
Boilerplate for pandas conventions in arithmetic and comparison methods.
|
|
|
|
Ensure method returns NotImplemented when operating against "senior"
|
|
classes. Ensure zero-dimensional ndarrays are always unpacked.
|
|
|
|
Parameters
|
|
----------
|
|
method : binary method
|
|
name : str
|
|
|
|
Returns
|
|
-------
|
|
method
|
|
"""
|
|
is_cmp = name.strip("__") in {"eq", "ne", "lt", "le", "gt", "ge"}
|
|
|
|
@wraps(method)
|
|
def new_method(self, other):
|
|
|
|
if is_cmp and isinstance(self, ABCIndexClass) and isinstance(other, ABCSeries):
|
|
# For comparison ops, Index does *not* defer to Series
|
|
pass
|
|
else:
|
|
for cls in [ABCDataFrame, ABCSeries, ABCIndexClass]:
|
|
if isinstance(self, cls):
|
|
break
|
|
if isinstance(other, cls):
|
|
return NotImplemented
|
|
|
|
other = item_from_zerodim(other)
|
|
|
|
return method(self, other)
|
|
|
|
return new_method
|
|
|
|
|
|
def get_op_result_name(left, right):
|
|
"""
|
|
Find the appropriate name to pin to an operation result. This result
|
|
should always be either an Index or a Series.
|
|
|
|
Parameters
|
|
----------
|
|
left : {Series, Index}
|
|
right : object
|
|
|
|
Returns
|
|
-------
|
|
name : object
|
|
Usually a string
|
|
"""
|
|
if isinstance(right, (ABCSeries, ABCIndexClass)):
|
|
name = _maybe_match_name(left, right)
|
|
else:
|
|
name = left.name
|
|
return name
|
|
|
|
|
|
def _maybe_match_name(a, b):
|
|
"""
|
|
Try to find a name to attach to the result of an operation between
|
|
a and b. If only one of these has a `name` attribute, return that
|
|
name. Otherwise return a consensus name if they match of None if
|
|
they have different names.
|
|
|
|
Parameters
|
|
----------
|
|
a : object
|
|
b : object
|
|
|
|
Returns
|
|
-------
|
|
name : str or None
|
|
|
|
See Also
|
|
--------
|
|
pandas.core.common.consensus_name_attr
|
|
"""
|
|
a_has = hasattr(a, "name")
|
|
b_has = hasattr(b, "name")
|
|
if a_has and b_has:
|
|
if a.name == b.name:
|
|
return a.name
|
|
else:
|
|
# TODO: what if they both have np.nan for their names?
|
|
return None
|
|
elif a_has:
|
|
return a.name
|
|
elif b_has:
|
|
return b.name
|
|
return None
|