test: Fully cover selectors in drop

Author: dangotbanned

narwhals/_plan/_parse.py

@@ -173,15 +173,15 @@ def parse_into_combined_selector_ir(
     return _any_of(selectors)._ir
 
 
-def _any_of(selectors: Iterable[Selector], /) -> Selector:
+def _any_of(selectors: Collection[Selector], /) -> Selector:
     import narwhals._plan.selectors as cs
 
-    if isinstance(selectors, Collection):
-        if not selectors:
-            return cs.empty()
-        if len(selectors) == 1:
-            return next(iter(selectors))  # type: ignore[no-any-return]
-    s: Selector = reduce(operator.or_, selectors)
+    if not selectors:
+        s: Selector = cs.empty()
+    elif len(selectors) == 1:
+        s = next(iter(selectors))
+    else:
+        s = reduce(operator.or_, selectors)
     return s
 
 

narwhals/_plan/dataframe.py

@@ -128,6 +128,9 @@ def drop_nulls(
     def rename(self, mapping: Mapping[str, str]) -> Self:
         return self._with_compliant(self._compliant.rename(mapping))
 
+    def collect_schema(self) -> Schema:
+        return self.schema
+
 
 class DataFrame(
     BaseFrame[NativeDataFrameT_co], Generic[NativeDataFrameT_co, NativeSeriesT]

tests/plan/compliant_test.py

@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+import re
+from collections.abc import Iterable
 from typing import TYPE_CHECKING, Any
 
 import pytest
@@ -528,6 +530,42 @@ def test_row_is_py_literal(
     assert result == polars_result
 
 
+@pytest.mark.parametrize(
+    ("columns", "expected"),
+    [
+        ("a", ["b", "c"]),
+        (["a"], ["b", "c"]),
+        (ncs.first(), ["b", "c"]),
+        ([ncs.first()], ["b", "c"]),
+        (["a", "b"], ["c"]),
+        (~ncs.last(), ["c"]),
+        ([ncs.integer() | ncs.enum()], ["c"]),
+        ([ncs.first(), "b"], ["c"]),
+        (ncs.all(), []),
+        ([], ["a", "b", "c"]),
+        (ncs.struct(), ["a", "b", "c"]),
+    ],
+)
+def test_drop(columns: OneOrIterable[ColumnNameOrSelector], expected: list[str]) -> None:
+    data = {"a": [1, 3, 2], "b": [4, 4, 6], "c": [7.0, 8.0, 9.0]}
+    df = dataframe(data)
+    if isinstance(columns, (str, nwp.Selector, list)):
+        assert df.drop(columns).collect_schema().names() == expected
+    if not isinstance(columns, str) and isinstance(columns, Iterable):
+        assert df.drop(*columns).collect_schema().names() == expected
+
+
+def test_drop_strict() -> None:
+    data = {"a": [1, 3, 2], "b": [4, 4, 6]}
+    df = dataframe(data)
+    with pytest.raises(ColumnNotFoundError):
+        df.drop("z")
+    with pytest.raises(ColumnNotFoundError, match=re.escape("not found: ['z']")):
+        df.drop(ncs.last(), "z")
+    assert df.drop("z", strict=False).collect_schema().names() == ["a", "b"]
+    assert df.drop(ncs.last(), "z", strict=False).collect_schema().names() == ["a"]
+
+
 def test_drop_nulls(data_small_dh: Data) -> None:
     df = dataframe(data_small_dh)
     expected: Data = {"d": [], "e": [], "f": [], "g": [], "h": []}