Skip to content

Commit

Permalink
simplify
Browse files Browse the repository at this point in the history
  • Loading branch information
MarcoGorelli committed Sep 30, 2024
1 parent 9094d5c commit fd69d7b
Show file tree
Hide file tree
Showing 2 changed files with 5 additions and 7 deletions.
6 changes: 1 addition & 5 deletions narwhals/_pandas_like/series.py
Original file line number Diff line number Diff line change
Expand Up @@ -619,11 +619,7 @@ def quantile(

def zip_with(self: Self, mask: Any, other: Any) -> PandasLikeSeries:
ser = self._native_series
mask = validate_column_comparand(ser.index, mask)
if isinstance(mask, str) or not isinstance(
mask, (self.__native_namespace__().Series, Sequence)
):
mask = [mask]
mask = validate_column_comparand(ser.index, mask, broadcast_length_one=False)
other = validate_column_comparand(ser.index, other)
res = ser.where(mask, other)
return self._from_native_series(res)
Expand Down
6 changes: 4 additions & 2 deletions narwhals/_pandas_like/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,9 @@
}


def validate_column_comparand(index: Any, other: Any) -> Any:
def validate_column_comparand(
index: Any, other: Any, *, broadcast_length_one: bool = True
) -> Any:
"""Validate RHS of binary operation.
If the comparison isn't supported, return `NotImplemented` so that the
Expand All @@ -50,7 +52,7 @@ def validate_column_comparand(index: Any, other: Any) -> Any:
if isinstance(other, PandasLikeDataFrame):
return NotImplemented
if isinstance(other, PandasLikeSeries):
if other.len() == 1:
if other.len() == 1 and broadcast_length_one:
# broadcast
return other.item()
if other._native_series.index is not index:
Expand Down

0 comments on commit fd69d7b

Please sign in to comment.