Function _wrap_expr_method_operator

dask/dataframe/dask_expr/_collection.py:197–289

Add method operators to Series or DataFrame like DataFrame.add. _wrap_expr_method_operator("add", DataFrame)

(name, class_)

def _wrap_expr_method_operator(name, class_):
    """
    Add method operators to Series or DataFrame like DataFrame.add.
    _wrap_expr_method_operator("add", DataFrame)
    """
    if class_ == DataFrame:

        def method(self, other, axis="columns", level=None, fill_value=None):
            if level is not None:
                raise NotImplementedError("level must be None")

            axis = _validate_axis(axis)

            if (
                is_dataframe_like(other)
                or is_series_like(other)
                and axis in (0, "index")
            ) and not is_dask_collection(other):
                other = self._create_alignable_frame(other)

            if axis in (1, "columns"):
                if isinstance(other, Series):
                    msg = f"Unable to {name} dd.Series with axis=1"
                    raise ValueError(msg)

            frame = self
            if isinstance(other, FrameBase) and not expr.are_co_aligned(
                self.expr, other.expr
            ):
                return new_collection(
                    expr.MethodOperatorAlign(
                        op=name,
                        frame=frame,
                        other=other,
                        axis=axis,
                        level=level,
                        fill_value=fill_value,
                    )
                )

            return new_collection(
                expr.MethodOperator(
                    name=name,
                    left=frame,
                    right=other,
                    axis=axis,
                    level=level,
                    fill_value=fill_value,
                )
            )

    elif class_ == Series:

        def method(self, other, level=None, fill_value=None, axis=0):  # type: ignore[misc]
            if level is not None:
                raise NotImplementedError("level must be None")

            axis = _validate_axis(axis)