__init__.py

"""
Arithmetic operations for PandasObjects

This is not a public API.
"""
from __future__ import annotations

import operator
from typing import (
    TYPE_CHECKING,
    cast,
)

import numpy as np

from pandas._libs.ops_dispatch import maybe_dispatch_ufunc_to_dunder_op
from pandas._typing import (
    Axis,
    AxisInt,
    Level,
)
from pandas.util._decorators import Appender

from pandas.core.dtypes.common import (
    is_array_like,
    is_list_like,
)
from pandas.core.dtypes.generic import (
    ABCDataFrame,
    ABCSeries,
)
from pandas.core.dtypes.missing import isna

from pandas.core import (
    algorithms,
    roperator,
)
from pandas.core.ops.array_ops import (
    arithmetic_op,
    comp_method_OBJECT_ARRAY,
    comparison_op,
    get_array_op,
    logical_op,
    maybe_prepare_scalar_for_op,
)
from pandas.core.ops.common import (
    get_op_result_name,
    unpack_zerodim_and_defer,
)
from pandas.core.ops.docstrings import (
    _flex_comp_doc_FRAME,
    _op_descriptions,
    make_flex_doc,
)
from pandas.core.ops.invalid import invalid_comparison
from pandas.core.ops.mask_ops import (
    kleene_and,
    kleene_or,
    kleene_xor,
)
from pandas.core.ops.methods import add_flex_arithmetic_methods
from pandas.core.roperator import (
    radd,
    rand_,
    rdiv,
    rdivmod,
    rfloordiv,
    rmod,
    rmul,
    ror_,
    rpow,
    rsub,
    rtruediv,
    rxor,
)

if TYPE_CHECKING:
    from pandas import (
        DataFrame,
        Series,
    )

# -----------------------------------------------------------------------------
# constants
ARITHMETIC_BINOPS: set[str] = {
    "add",
    "sub",
    "mul",
    "pow",
    "mod",
    "floordiv",
    "truediv",
    "divmod",
    "radd",
    "rsub",
    "rmul",
    "rpow",
    "rmod",
    "rfloordiv",
    "rtruediv",
    "rdivmod",
}


COMPARISON_BINOPS: set[str] = {"eq", "ne", "lt", "gt", "le", "ge"}


# -----------------------------------------------------------------------------
# Masking NA values and fallbacks for operations numpy does not support


def fill_binop(left, right, fill_value):
    """
    If a non-None fill_value is given, replace null entries in left and right
    with this value, but only in positions where _one_ of left/right is null,
    not both.

    Parameters
    ----------
    left : array-like
    right : array-like
    fill_value : object

    Returns
    -------
    left : array-like
    right : array-like

    Notes
    -----
    Makes copies if fill_value is not None and NAs are present.
    """
    if fill_value is not None:
        left_mask = isna(left)
        right_mask = isna(right)

        # one but not both
        mask = left_mask ^ right_mask

        if left_mask.any():
            # Avoid making a copy if we can
            left = left.copy()
            left[left_mask & mask] = fill_value

        if right_mask.any():
            # Avoid making a copy if we can
            right = right.copy()
            right[right_mask & mask] = fill_value

    return left, right


# -----------------------------------------------------------------------------
# Series


def align_method_SERIES(left: Series, right, align_asobject: bool = False):
    """align lhs and rhs Series"""
    # ToDo: Different from align_method_FRAME, list, tuple and ndarray
    # are not coerced here
    # because Series has inconsistencies described in #13637

    if isinstance(right, ABCSeries):
        # avoid repeated alignment
        if not left.index.equals(right.index):
            if align_asobject:
                # to keep original value's dtype for bool ops
                left = left.astype(object)
                right = right.astype(object)

            left, right = left.align(right, copy=False)

    return left, right


def flex_method_SERIES(op):
    name = op.__name__.strip("_")
    doc = make_flex_doc(name, "series")

    @Appender(doc)
    def flex_wrapper(self, other, level=None, fill_value=None, axis: Axis = 0):
        # validate axis
        if axis is not None:
            self._get_axis_number(axis)

        res_name = get_op_result_name(self, other)

        if isinstance(other, ABCSeries):
            return self._binop(other, op, level=level, fill_value=fill_value)
        elif isinstance(other, (np.ndarray, list, tuple)):
            if len(other) != len(self):
                raise ValueError("Lengths must be equal")
            other = self._constructor(other, self.index)
            result = self._binop(other, op, level=level, fill_value=fill_value)
            result.name = res_name
            return result
        else:
            if fill_value is not None:
                self = self.fillna(fill_value)

            return op(self, other)

    flex_wrapper.__name__ = name
    return flex_wrapper


# -----------------------------------------------------------------------------
# DataFrame


def align_method_FRAME(
    left, right, axis, flex: bool | None = False, level: Level = None
):
    """
    Convert rhs to meet lhs dims if input is list, tuple or np.ndarray.

    Parameters
    ----------
    left : DataFrame
    right : Any
    axis : int, str, or None
    flex : bool or None, default False
        Whether this is a flex op, in which case we reindex.
        None indicates not to check for alignment.
    level : int or level name, default None

    Returns
    -------
    left : DataFrame
    right : Any
    """

    def to_series(right):
        msg = "Unable to coerce to Series, length must be {req_len}: given {given_len}"

        # pass dtype to avoid doing inference, which would break consistency
        #  with Index/Series ops
        dtype = None
        if getattr(right, "dtype", None) == object:
            # can't pass right.dtype unconditionally as that would break on e.g.
            #  datetime64[h] ndarray
            dtype = object

        if axis is not None and left._get_axis_name(axis) == "index":
            if len(left.index) != len(right):
                raise ValueError(
                    msg.format(req_len=len(left.index), given_len=len(right))
                )
            right = left._constructor_sliced(right, index=left.index, dtype=dtype)
        else:
            if len(left.columns) != len(right):
                raise ValueError(
                    msg.format(req_len=len(left.columns), given_len=len(right))
                )
            right = left._constructor_sliced(right, index=left.columns, dtype=dtype)
        return right

    if isinstance(right, np.ndarray):
        if right.ndim == 1:
            right = to_series(right)

        elif right.ndim == 2:
            # We need to pass dtype=right.dtype to retain object dtype
            #  otherwise we lose consistency with Index and array ops
            dtype = None
            if getattr(right, "dtype", None) == object:
                # can't pass right.dtype unconditionally as that would break on e.g.
                #  datetime64[h] ndarray
                dtype = object

            if right.shape == left.shape:
                right = left._constructor(
                    right, index=left.index, columns=left.columns, dtype=dtype
                )

            elif right.shape[0] == left.shape[0] and right.shape[1] == 1:
                # Broadcast across columns
                right = np.broadcast_to(right, left.shape)
                right = left._constructor(
                    right, index=left.index, columns=left.columns, dtype=dtype
                )

            elif right.shape[1] == left.shape[1] and right.shape[0] == 1:
                # Broadcast along rows
                right = to_series(right[0, :])

            else:
                raise ValueError(
                    "Unable to coerce to DataFrame, shape "
                    f"must be {left.shape}: given {right.shape}"
                )

        elif right.ndim > 2:
            raise ValueError(
                "Unable to coerce to Series/DataFrame, "
                f"dimension must be <= 2: {right.shape}"
            )

    elif is_list_like(right) and not isinstance(right, (ABCSeries, ABCDataFrame)):
        # GH 36702. Raise when attempting arithmetic with list of array-like.
        if any(is_array_like(el) for el in right):
            raise ValueError(
                f"Unable to coerce list of {type(right[0])} to Series/DataFrame"
            )
        # GH17901
        right = to_series(right)

    if flex is not None and isinstance(right, ABCDataFrame):
        if not left._indexed_same(right):
            if flex:
                left, right = left.align(right, join="outer", level=level, copy=False)
            else:
                raise ValueError(
                    "Can only compare identically-labeled (both index and columns) "
                    "DataFrame objects"
                )
    elif isinstance(right, ABCSeries):
        # axis=1 is default for DataFrame-with-Series op
        axis = left._get_axis_number(axis) if axis is not None else 1

        if not flex:
            if not left.axes[axis].equals(right.index):
                raise ValueError(
                    "Operands are not aligned. Do "
                    "`left, right = left.align(right, axis=1, copy=False)` "
                    "before operating."
                )

        left, right = left.align(
            right, join="outer", axis=axis, level=level, copy=False
        )
        right = _maybe_align_series_as_frame(left, right, axis)

    return left, right


def should_reindex_frame_op(
    left: DataFrame, right, op, axis: int, fill_value, level
) -> bool:
    """
    Check if this is an operation between DataFrames that will need to reindex.
    """
    assert isinstance(left, ABCDataFrame)

    if op is operator.pow or op is roperator.rpow:
        # GH#32685 pow has special semantics for operating with null values
        return False

    if not isinstance(right, ABCDataFrame):
        return False

    if fill_value is None and level is None and axis == 1:
        # TODO: any other cases we should handle here?

        # Intersection is always unique so we have to check the unique columns
        left_uniques = left.columns.unique()
        right_uniques = right.columns.unique()
        cols = left_uniques.intersection(right_uniques)
        if len(cols) and not (
            len(cols) == len(left_uniques) and len(cols) == len(right_uniques)
        ):
            # TODO: is there a shortcut available when len(cols) == 0?
            return True

    return False


def frame_arith_method_with_reindex(left: DataFrame, right: DataFrame, op) -> DataFrame:
    """
    For DataFrame-with-DataFrame operations that require reindexing,
    operate only on shared columns, then reindex.

    Parameters
    ----------
    left : DataFrame
    right : DataFrame
    op : binary operator

    Returns
    -------
    DataFrame
    """
    # GH#31623, only operate on shared columns
    cols, lcols, rcols = left.columns.join(
        right.columns, how="inner", level=None, return_indexers=True
    )

    new_left = left.iloc[:, lcols]
    new_right = right.iloc[:, rcols]
    result = op(new_left, new_right)

    # Do the join on the columns instead of using align_method_FRAME
    #  to avoid constructing two potentially large/sparse DataFrames
    join_columns, _, _ = left.columns.join(
        right.columns, how="outer", level=None, return_indexers=True
    )

    if result.columns.has_duplicates:
        # Avoid reindexing with a duplicate axis.
        # https://github.com/pandas-dev/pandas/issues/35194
        indexer, _ = result.columns.get_indexer_non_unique(join_columns)
        indexer = algorithms.unique1d(indexer)
        result = result._reindex_with_indexers(
            {1: [join_columns, indexer]}, allow_dups=True
        )
    else:
        result = result.reindex(join_columns, axis=1)

    return result


def _maybe_align_series_as_frame(frame: DataFrame, series: Series, axis: AxisInt):
    """
    If the Series operand is not EA-dtype, we can broadcast to 2D and operate
    blockwise.
    """
    rvalues = series._values
    if not isinstance(rvalues, np.ndarray):
        # TODO(EA2D): no need to special-case with 2D EAs
        if rvalues.dtype in ("datetime64[ns]", "timedelta64[ns]"):
            # We can losslessly+cheaply cast to ndarray
            rvalues = np.asarray(rvalues)
        else:
            return series

    if axis == 0:
        rvalues = rvalues.reshape(-1, 1)
    else:
        rvalues = rvalues.reshape(1, -1)

    rvalues = np.broadcast_to(rvalues, frame.shape)
    # pass dtype to avoid doing inference
    return type(frame)(
        rvalues, index=frame.index, columns=frame.columns, dtype=rvalues.dtype
    )


def flex_arith_method_FRAME(op):
    op_name = op.__name__.strip("_")

    na_op = get_array_op(op)
    doc = make_flex_doc(op_name, "dataframe")

    @Appender(doc)
    def f(self, other, axis: Axis = "columns", level=None, fill_value=None):
        axis = self._get_axis_number(axis) if axis is not None else 1
        axis = cast(int, axis)

        if should_reindex_frame_op(self, other, op, axis, fill_value, level):
            return frame_arith_method_with_reindex(self, other, op)

        if isinstance(other, ABCSeries) and fill_value is not None:
            # TODO: We could allow this in cases where we end up going
            #  through the DataFrame path
            raise NotImplementedError(f"fill_value {fill_value} not supported.")

        other = maybe_prepare_scalar_for_op(other, self.shape)
        self, other = align_method_FRAME(self, other, axis, flex=True, level=level)

        if isinstance(other, ABCDataFrame):
            # Another DataFrame
            new_data = self._combine_frame(other, na_op, fill_value)

        elif isinstance(other, ABCSeries):
            new_data = self._dispatch_frame_op(other, op, axis=axis)
        else:
            # in this case we always have `np.ndim(other) == 0`
            if fill_value is not None:
                self = self.fillna(fill_value)

            new_data = self._dispatch_frame_op(other, op)

        return self._construct_result(new_data)

    f.__name__ = op_name

    return f


def flex_comp_method_FRAME(op):
    op_name = op.__name__.strip("_")

    doc = _flex_comp_doc_FRAME.format(
        op_name=op_name, desc=_op_descriptions[op_name]["desc"]
    )

    @Appender(doc)
    def f(self, other, axis: Axis = "columns", level=None):
        axis = self._get_axis_number(axis) if axis is not None else 1

        self, other = align_method_FRAME(self, other, axis, flex=True, level=level)

        new_data = self._dispatch_frame_op(other, op, axis=axis)
        return self._construct_result(new_data)

    f.__name__ = op_name

    return f


__all__ = [
    "add_flex_arithmetic_methods",
    "align_method_FRAME",
    "align_method_SERIES",
    "ARITHMETIC_BINOPS",
    "arithmetic_op",
    "COMPARISON_BINOPS",
    "comparison_op",
    "comp_method_OBJECT_ARRAY",
    "fill_binop",
    "flex_arith_method_FRAME",
    "flex_comp_method_FRAME",
    "flex_method_SERIES",
    "frame_arith_method_with_reindex",
    "invalid_comparison",
    "kleene_and",
    "kleene_or",
    "kleene_xor",
    "logical_op",
    "maybe_dispatch_ufunc_to_dunder_op",
    "radd",
    "rand_",
    "rdiv",
    "rdivmod",
    "rfloordiv",
    "rmod",
    "rmul",
    "ror_",
    "rpow",
    "rsub",
    "rtruediv",
    "rxor",
    "should_reindex_frame_op",
    "unpack_zerodim_and_defer",
]