add selector support to "melt" frame method

py-polars/polars/dataframe/frame.py

@@ -70,7 +70,7 @@
     _xl_unique_table_name,
     _XLFormatCache,
 )
-from polars.selectors import expand_selectors, is_selector
+from polars.selectors import expand_selectors
 from polars.slice import PolarsSlice
 from polars.utils import no_default
 from polars.utils._construction import (
@@ -6609,21 +6609,10 @@ def pivot(
         └──────┴──────────┴──────────┘
 
         """  # noqa: W505
-        values = (
-            [values]
-            if isinstance(values, str)
-            else (expand_selectors(self, values) if is_selector(values) else values)
-        )
-        index = (
-            [index]
-            if isinstance(index, str)
-            else (expand_selectors(self, index) if is_selector(index) else index)
-        )
-        columns = (
-            [columns]
-            if isinstance(columns, str)
-            else (expand_selectors(self, columns) if is_selector(columns) else columns)
-        )
+        values = expand_selectors(self, values)
+        index = expand_selectors(self, index)
+        columns = expand_selectors(self, columns)
+
         if aggregate_function is no_default:
             issue_deprecation_warning(
                 "In a future version of polars, the default `aggregate_function` "

py-polars/polars/lazyframe/frame.py

@@ -4520,8 +4520,8 @@ def drop_nulls(
 
     def melt(
         self,
-        id_vars: str | list[str] | None = None,
-        value_vars: str | list[str] | None = None,
+        id_vars: str | Sequence[str] | SelectorType | None = None,
+        value_vars: str | Sequence[str] | SelectorType | None = None,
         variable_name: str | None = None,
         value_name: str | None = None,
         *,
@@ -4578,14 +4578,9 @@ def melt(
         └─────┴──────────┴───────┘
 
         """
-        if isinstance(value_vars, str):
-            value_vars = [value_vars]
-        if isinstance(id_vars, str):
-            id_vars = [id_vars]
-        if value_vars is None:
-            value_vars = []
-        if id_vars is None:
-            id_vars = []
+        value_vars = [] if value_vars is None else expand_selectors(self, value_vars)
+        id_vars = [] if id_vars is None else expand_selectors(self, id_vars)
+
         return self._from_pyldf(
             self._ldf.melt(id_vars, value_vars, value_name, variable_name, streamable)
         )

py-polars/tests/unit/dataframe/test_df.py

@@ -728,46 +728,6 @@ def test_read_missing_file(read_function: Callable[[Any], pl.DataFrame]) -> None
         read_function("fake_file")
 
 
-def test_melt() -> None:
-    df = pl.DataFrame({"A": ["a", "b", "c"], "B": [1, 3, 5], "C": [2, 4, 6]})
-    melted = df.melt(id_vars="A", value_vars=["B", "C"])
-    assert all(melted["value"] == [1, 3, 5, 2, 4, 6])
-
-    melted = df.melt(id_vars="A", value_vars="B")
-    assert all(melted["value"] == [1, 3, 5])
-    n = 3
-    for melted in [df.melt(), df.lazy().melt().collect()]:
-        assert melted["variable"].to_list() == ["A"] * n + ["B"] * n + ["C"] * n
-        assert melted["value"].to_list() == [
-            "a",
-            "b",
-            "c",
-            "1",
-            "3",
-            "5",
-            "2",
-            "4",
-            "6",
-        ]
-
-    for melted in [
-        df.melt(value_name="foo", variable_name="bar"),
-        df.lazy().melt(value_name="foo", variable_name="bar").collect(),
-    ]:
-        assert melted["bar"].to_list() == ["A"] * n + ["B"] * n + ["C"] * n
-        assert melted["foo"].to_list() == [
-            "a",
-            "b",
-            "c",
-            "1",
-            "3",
-            "5",
-            "2",
-            "4",
-            "6",
-        ]
-
-
 def test_shift() -> None:
     df = pl.DataFrame({"A": ["a", "b", "c"], "B": [1, 3, 5]})
     a = df.shift(1)

py-polars/tests/unit/operations/test_melt.py

@@ -1,7 +1,50 @@
 import polars as pl
+import polars.selectors as cs
 from polars.testing import assert_frame_equal
 
 
+def test_melt() -> None:
+    df = pl.DataFrame({"A": ["a", "b", "c"], "B": [1, 3, 5], "C": [2, 4, 6]})
+    for _idv, _vv in (("A", ("B", "C")), (cs.string(), cs.integer())):
+        melted = df.melt(id_vars="A", value_vars=["B", "C"])
+        assert all(melted["value"] == [1, 3, 5, 2, 4, 6])
+
+    melted = df.melt(id_vars="A", value_vars="B")
+    assert all(melted["value"] == [1, 3, 5])
+    n = 3
+
+    for melted in [df.melt(), df.lazy().melt().collect()]:
+        assert melted["variable"].to_list() == ["A"] * n + ["B"] * n + ["C"] * n
+        assert melted["value"].to_list() == [
+            "a",
+            "b",
+            "c",
+            "1",
+            "3",
+            "5",
+            "2",
+            "4",
+            "6",
+        ]
+
+    for melted in [
+        df.melt(value_name="foo", variable_name="bar"),
+        df.lazy().melt(value_name="foo", variable_name="bar").collect(),
+    ]:
+        assert melted["bar"].to_list() == ["A"] * n + ["B"] * n + ["C"] * n
+        assert melted["foo"].to_list() == [
+            "a",
+            "b",
+            "c",
+            "1",
+            "3",
+            "5",
+            "2",
+            "4",
+            "6",
+        ]
+
+
 def test_melt_projection_pd_7747() -> None:
     df = pl.LazyFrame(
         {