Repository: narwhals-dev/narwhals
Pull request #1102 (2024-09-30 18:54:26)

chore: Refactor validate_column_comparand
<!--
# Thanks for contributing a pull request! 
## Please make sure you see our contribution guidelines: https://github.com/narwhals-dev/narwhals/blob/main/CONTRIBUTING.md
-->

## What type of PR is this? (check all applicable)

- [ ] 💾 Refactor
- [ ] ✨ Feature
- [ ] 🐛 Bug Fix
- [ ] 🔧 Optimization
- [ ] 📝 Documentation
- [ ] ✅ Test
- [ ] 🐳 Other

## Related issues 

- Related issue # 
- Closes #

## Checklist

- [ ] Code follows style guide (ruff)
- [ ] Tests added 
- [ ] Documented the changes

## If you have comments or can explain your changes, please do so below.




Patch:
diff --git a/docs/api-reference/narwhals.md b/docs/api-reference/narwhals.md
index d678b873..2700c48c 100644
--- a/docs/api-reference/narwhals.md
+++ b/docs/api-reference/narwhals.md
@@ -22,6 +22,7 @@ Here are the top-level functions available in Narwhals.
         - maybe_align_index
         - maybe_convert_dtypes
         - maybe_get_index
+        - maybe_reset_index
         - maybe_set_index
         - mean
         - mean_horizontal
diff --git a/narwhals/__init__.py b/narwhals/__init__.py
index 2b571f0e..e5ee71a1 100644
--- a/narwhals/__init__.py
+++ b/narwhals/__init__.py
@@ -55,6 +55,7 @@ from narwhals.utils import is_ordered_categorical
 from narwhals.utils import maybe_align_index
 from narwhals.utils import maybe_convert_dtypes
 from narwhals.utils import maybe_get_index
+from narwhals.utils import maybe_reset_index
 from narwhals.utils import maybe_set_index
 
 __version__ = "1.8.4"
@@ -72,6 +73,7 @@ __all__ = [
     "maybe_align_index",
     "maybe_convert_dtypes",
     "maybe_get_index",
+    "maybe_reset_index",
     "maybe_set_index",
     "get_native_namespace",
     "all",
diff --git a/narwhals/_arrow/utils.py b/narwhals/_arrow/utils.py
index e34e949d..bb5f3f7f 100644
--- a/narwhals/_arrow/utils.py
+++ b/narwhals/_arrow/utils.py
@@ -58,7 +58,7 @@ def native_to_narwhals_dtype(dtype: Any, dtypes: DTypes) -> DType:
     if pa.types.is_struct(dtype):
         return dtypes.Struct()
     if pa.types.is_list(dtype) or pa.types.is_large_list(dtype):
-        return dtypes.List()
+        return dtypes.List(native_to_narwhals_dtype(dtype.value_type, dtypes))
     if pa.types.is_fixed_size_list(dtype):
         return dtypes.Array()
     return dtypes.Unknown()  # pragma: no cover
diff --git a/narwhals/_duckdb/dataframe.py b/narwhals/_duckdb/dataframe.py
index 5e8cb73d..95567283 100644
--- a/narwhals/_duckdb/dataframe.py
+++ b/narwhals/_duckdb/dataframe.py
@@ -50,8 +50,8 @@ def map_duckdb_dtype_to_narwhals_dtype(duckdb_dtype: Any, dtypes: DTypes) -> DTy
         return dtypes.Duration()
     if duckdb_dtype.startswith("STRUCT"):
         return dtypes.Struct()
-    if re.match(r"\w+\[\]", duckdb_dtype):
-        return dtypes.List()
+    if match_ := re.match(r"(.*)\[\]$", duckdb_dtype):
+        return dtypes.List(map_duckdb_dtype_to_narwhals_dtype(match_.group(1), dtypes))
     if re.match(r"\w+\[\d+\]", duckdb_dtype):
         return dtypes.Array()
     return dtypes.Unknown()
diff --git a/narwhals/_ibis/dataframe.py b/narwhals/_ibis/dataframe.py
index 6f53e277..8ee01e78 100644
--- a/narwhals/_ibis/dataframe.py
+++ b/narwhals/_ibis/dataframe.py
@@ -43,7 +43,9 @@ def map_ibis_dtype_to_narwhals_dtype(ibis_dtype: Any, dtypes: DTypes) -> DType:
     if ibis_dtype.is_timestamp():
         return dtypes.Datetime()
     if ibis_dtype.is_array():
-        return dtypes.List()
+        return dtypes.List(
+            map_ibis_dtype_to_narwhals_dtype(ibis_dtype.value_type, dtypes)
+        )
     if ibis_dtype.is_struct():
         return dtypes.Struct()
     return dtypes.Unknown()  # pragma: no cover
diff --git a/narwhals/_pandas_like/series.py b/narwhals/_pandas_like/series.py
index f4bcaec0..74e1c492 100644
--- a/narwhals/_pandas_like/series.py
+++ b/narwhals/_pandas_like/series.py
@@ -619,11 +619,9 @@ class PandasLikeSeries:
 
     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, treat_length_one_as_scalar=False
+        )
         other = validate_column_comparand(ser.index, other)
         res = ser.where(mask, other)
         return self._from_native_series(res)
diff --git a/narwhals/_pandas_like/utils.py b/narwhals/_pandas_like/utils.py
index 92fbd519..df87e649 100644
--- a/narwhals/_pandas_like/utils.py
+++ b/narwhals/_pandas_like/utils.py
@@ -7,6 +7,9 @@ from typing import Iterable
 from typing import Literal
 from typing import TypeVar
 
+from narwhals._arrow.utils import (
+    native_to_narwhals_dtype as arrow_native_to_narwhals_dtype,
+)
 from narwhals.utils import Implementation
 from narwhals.utils import isinstance_or_issubclass
 
@@ -29,7 +32,9 @@ PANDAS_LIKE_IMPLEMENTATION = {
 }
 
 
-def validate_column_comparand(index: Any, other: Any) -> Any:
+def validate_column_comparand(
+    index: Any, other: Any, *, treat_length_one_as_scalar: bool = True
+) -> Any:
     """Validate RHS of binary operation.
 
     If the comparison isn't supported, return `NotImplemented` so that the
@@ -50,7 +55,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 treat_length_one_as_scalar:
             # broadcast
             return other.item()
         if other._native_series.index is not index:
@@ -276,7 +281,9 @@ def native_to_narwhals_dtype(column: Any, dtypes: DTypes) -> DType:
     if dtype == "date32[day][pyarrow]":
         return dtypes.Date()
     if dtype.startswith(("large_list", "list")):
-        return dtypes.List()
+        return dtypes.List(
+            arrow_native_to_narwhals_dtype(column.dtype.pyarrow_dtype.value_type, dtypes)
+        )
     if dtype.startswith("fixed_size_list"):
         return dtypes.Array()
     if dtype.startswith("struct"):
diff --git a/narwhals/_polars/utils.py b/narwhals/_polars/utils.py
index b2f06090..95b98b17 100644
--- a/narwhals/_polars/utils.py
+++ b/narwhals/_polars/utils.py
@@ -75,7 +75,7 @@ def native_to_narwhals_dtype(dtype: Any, dtypes: DTypes) -> DType:
     if dtype == pl.Struct:
         return dtypes.Struct()
     if dtype == pl.List:
-        return dtypes.List()
+        return dtypes.List(native_to_narwhals_dtype(dtype.inner, dtypes))
     if dtype == pl.Array:
         return dtypes.Array()
     return dtypes.Unknown()
diff --git a/narwhals/dtypes.py b/narwhals/dtypes.py
index 730f6984..6acfbb76 100644
--- a/narwhals/dtypes.py
+++ b/narwhals/dtypes.py
@@ -173,7 +173,31 @@ class Enum(DType): ...
 class Struct(DType): ...
 
 
-class List(DType): ...
+class List(DType):
+    def __init__(self, inner: DType | type[DType]) -> None:
+        self.inner = inner
+
+    def __eq__(self, other: DType | type[DType]) -> bool:  # type: ignore[override]
+        # This equality check allows comparison of type classes and type instances.
+        # If a parent type is not specific about its inner type, we infer it as equal:
+        # > list[i64] == list[i64] -> True
+        # > list[i64] == list[f32] -> False
+        # > list[i64] == list      -> True
+
+        # allow comparing object instances to class
+        if type(other) is type and issubclass(other, self.__class__):
+            return True
+        elif isinstance(other, self.__class__):
+            return self.inner == other.inner
+        else:
+            return False
+
+    def __hash__(self) -> int:
+        return hash((self.__class__, self.inner))
+
+    def __repr__(self) -> str:
+        class_name = self.__class__.__name__
+        return f"{class_name}({self.inner!r})"
 
 
 class Array(DType): ...
diff --git a/narwhals/stable/v1/__init__.py b/narwhals/stable/v1/__init__.py
index b542b90f..b0cefc3e 100644
--- a/narwhals/stable/v1/__init__.py
+++ b/narwhals/stable/v1/__init__.py
@@ -55,6 +55,7 @@ from narwhals.utils import is_ordered_categorical as nw_is_ordered_categorical
 from narwhals.utils import maybe_align_index as nw_maybe_align_index
 from narwhals.utils import maybe_convert_dtypes as nw_maybe_convert_dtypes
 from narwhals.utils import maybe_get_index as nw_maybe_get_index
+from narwhals.utils import maybe_reset_index as nw_maybe_reset_index
 from narwhals.utils import maybe_set_index as nw_maybe_set_index
 
 if TYPE_CHECKING:
@@ -1802,6 +1803,35 @@ def maybe_set_index(df: T, column_names: str | list[str]) -> T:
     return nw_maybe_set_index(df, column_names)
 
 
+def maybe_reset_index(obj: T) -> T:
+    """
+    Reset the index to the default integer index of a DataFrame or a Series, if it's pandas-like.
+
+    Notes:
+        This is only really intended for backwards-compatibility purposes,
+        for example if your library already resets the index for users.
+        If you're designing a new library, we highly encourage you to not
+        rely on the Index.
+        For non-pandas-like inputs, this is a no-op.
+
+    Examples:
+        >>> import pandas as pd
+        >>> import polars as pl
+        >>> import narwhals.stable.v1 as nw
+        >>> df_pd = pd.DataFrame({"a": [1, 2], "b": [4, 5]}, index=([6, 7]))
+        >>> df = nw.from_native(df_pd)
+        >>> nw.to_native(nw.maybe_reset_index(df))
+           a  b
+        0  1  4
+        1  2  5
+        >>> series_pd = pd.Series([1, 2])
+        >>> series = nw.from_native(series_pd, series_only=True)
+        >>> nw.maybe_get_index(series)
+        RangeIndex(start=0, stop=2, step=1)
+    """
+    return nw_maybe_reset_index(obj)
+
+
 def get_native_namespace(obj: Any) -> Any:
     """
     Get native namespace from object.
@@ -2032,6 +2062,7 @@ __all__ = [
     "maybe_align_index",
     "maybe_convert_dtypes",
     "maybe_get_index",
+    "maybe_reset_index",
     "maybe_set_index",
     "get_native_namespace",
     "get_level",
diff --git a/narwhals/utils.py b/narwhals/utils.py
index 62ae7730..37cce17d 100644
--- a/narwhals/utils.py
+++ b/narwhals/utils.py
@@ -306,6 +306,47 @@ def maybe_set_index(df: T, column_names: str | list[str]) -> T:
     return df_any  # type: ignore[no-any-return]
 
 
+def maybe_reset_index(obj: T) -> T:
+    """
+    Reset the index to the default integer index of a DataFrame or a Series, if it's pandas-like.
+
+    Notes:
+        This is only really intended for backwards-compatibility purposes,
+        for example if your library already resets the index for users.
+        If you're designing a new library, we highly encourage you to not
+        rely on the Index.
+        For non-pandas-like inputs, this is a no-op.
+
+    Examples:
+        >>> import pandas as pd
+        >>> import polars as pl
+        >>> import narwhals as nw
+        >>> df_pd = pd.DataFrame({"a": [1, 2], "b": [4, 5]}, index=([6, 7]))
+        >>> df = nw.from_native(df_pd)
+        >>> nw.to_native(nw.maybe_reset_index(df))
+           a  b
+        0  1  4
+        1  2  5
+        >>> series_pd = pd.Series([1, 2])
+        >>> series = nw.from_native(series_pd, series_only=True)
+        >>> nw.maybe_get_index(series)
+        RangeIndex(start=0, stop=2, step=1)
+    """
+    obj_any = cast(Any, obj)
+    native_obj = to_native(obj_any)
+    if is_pandas_like_dataframe(native_obj):
+        return obj_any._from_compliant_dataframe(  # type: ignore[no-any-return]
+            obj_any._compliant_frame._from_native_frame(native_obj.reset_index(drop=True))
+        )
+    if is_pandas_like_series(native_obj):
+        return obj_any._from_compliant_series(  # type: ignore[no-any-return]
+            obj_any._compliant_series._from_native_series(
+                native_obj.reset_index(drop=True)
+            )
+        )
+    return obj_any  # type: ignore[no-any-return]
+
+
 def maybe_convert_dtypes(obj: T, *args: bool, **kwargs: bool | str) -> T:
     """
     Convert columns or series to the best possible dtypes using dtypes supporting ``pd.NA``, if df is pandas-like.


Test patch:
diff --git a/tests/dtypes_test.py b/tests/dtypes_test.py
index 58061597..c64dfed5 100644
--- a/tests/dtypes_test.py
+++ b/tests/dtypes_test.py
@@ -53,7 +53,21 @@ def test_duration_invalid(time_unit: str) -> None:
         nw.Duration(time_unit=time_unit)  # type: ignore[arg-type]
 
 
-def test_second_tu() -> None:
+def test_list_valid() -> None:
+    dtype = nw.List(nw.Int64)
+    assert dtype == nw.List(nw.Int64)
+    assert dtype == nw.List
+    assert dtype != nw.List(nw.Float32)
+    assert dtype != nw.Duration
+    assert repr(dtype) == "List(<class 'narwhals.dtypes.Int64'>)"
+    dtype = nw.List(nw.List(nw.Int64))
+    assert dtype == nw.List(nw.List(nw.Int64))
+    assert dtype == nw.List
+    assert dtype != nw.List(nw.List(nw.Float32))
+    assert dtype in {nw.List(nw.List(nw.Int64))}
+
+
+def test_second_time_unit() -> None:
     s = pd.Series(np.array([np.datetime64("2020-01-01", "s")]))
     result = nw.from_native(s, series_only=True)
     if parse_version(pd.__version__) < (2,):  # pragma: no cover
diff --git a/tests/series_only/__iter___test.py b/tests/series_only/__iter___test.py
index d190cd80..a0a5c118 100644
--- a/tests/series_only/__iter___test.py
+++ b/tests/series_only/__iter___test.py
@@ -3,13 +3,17 @@ from __future__ import annotations
 from collections.abc import Iterable
 from typing import Any
 
+import pytest
+
 import narwhals.stable.v1 as nw
 from tests.utils import compare_dicts
 
 data = [1, 2, 3]
 
 
-def test_to_list(constructor_eager: Any) -> None:
+def test_iter(constructor_eager: Any, request: pytest.FixtureRequest) -> None:
+    if "cudf" in str(constructor_eager):
+        request.applymarker(pytest.mark.xfail)
     s = nw.from_native(constructor_eager({"a": data}), eager_only=True)["a"]
 
     assert isinstance(s, Iterable)
diff --git a/tests/utils.py b/tests/utils.py
index e8b453f9..15ce2514 100644
--- a/tests/utils.py
+++ b/tests/utils.py
@@ -11,6 +11,7 @@ from typing import Sequence
 import pandas as pd
 
 from narwhals.typing import IntoFrame
+from narwhals.utils import Implementation
 
 if sys.version_info >= (3, 10):
     from typing import TypeAlias  # pragma: no cover
@@ -34,7 +35,13 @@ def compare_dicts(result: Any, expected: dict[str, Any]) -> None:
         for key in result.columns:
             assert key in expected
     for key in expected:
-        for lhs, rhs in zip_strict(result[key], expected[key]):
+        result_key = result[key]
+        if (
+            hasattr(result_key, "_compliant_series")
+            and result_key._compliant_series._implementation is Implementation.CUDF
+        ):  # pragma: no cover
+            result_key = result_key.to_pandas()
+        for lhs, rhs in zip_strict(result_key, expected[key]):
             if hasattr(lhs, "as_py"):
                 lhs = lhs.as_py()  # noqa: PLW2901
             if hasattr(rhs, "as_py"):  # pragma: no cover
diff --git a/tests/utils_test.py b/tests/utils_test.py
index f51c28ea..cea458bc 100644
--- a/tests/utils_test.py
+++ b/tests/utils_test.py
@@ -84,6 +84,30 @@ def test_maybe_get_index_polars() -> None:
     assert result is None
 
 
+def test_maybe_reset_index_pandas() -> None:
+    pandas_df = nw.from_native(
+        pd.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]}, index=[7, 8, 9])
+    )
+    result = nw.maybe_reset_index(pandas_df)
+    expected = pd.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]}, index=[0, 1, 2])
+    assert_frame_equal(nw.to_native(result), expected)
+    pandas_series = nw.from_native(
+        pd.Series([1, 2, 3], index=[7, 8, 9]), series_only=True
+    )
+    result_s = nw.maybe_reset_index(pandas_series)
+    expected_s = pd.Series([1, 2, 3], index=[0, 1, 2])
+    assert_series_equal(nw.to_native(result_s), expected_s)
+
+
+def test_maybe_reset_index_polars() -> None:
+    df = nw.from_native(pl.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]}))
+    result = nw.maybe_reset_index(df)
+    assert result is df
+    series = nw.from_native(pl.Series([1, 2, 3]), series_only=True)
+    result_s = nw.maybe_reset_index(series)
+    assert result_s is series
+
+
 @pytest.mark.skipif(
     parse_version(pd.__version__) < parse_version("1.0.0"),
     reason="too old for convert_dtypes",