From aa52b3a7d4351a03fe29708fcd30dc00adb04d6f Mon Sep 17 00:00:00 2001
From: Alexander Beedie <alexander-beedie@users.noreply.github.com>
Date: Fri, 4 Aug 2023 10:52:31 +0000
Subject: [PATCH] lint/typing

---
 py-polars/polars/lazyframe/frame.py           |  1 +
 py-polars/polars/type_aliases.py              |  2 +-
 py-polars/tests/unit/dataframe/test_df.py     | 79 ++++++++++---------
 py-polars/tests/unit/datatypes/test_struct.py |  2 +-
 4 files changed, 43 insertions(+), 41 deletions(-)

diff --git a/py-polars/polars/lazyframe/frame.py b/py-polars/polars/lazyframe/frame.py
index 0f9767ad5f485..2577d2578ef32 100644
--- a/py-polars/polars/lazyframe/frame.py
+++ b/py-polars/polars/lazyframe/frame.py
@@ -4443,6 +4443,7 @@ def unique(
         return self._from_pyldf(self._ldf.unique(maintain_order, subset, keep))
 
     def drop_nulls(
+        self,
         subset: (
             str | SelectorType | Collection[str] | Collection[SelectorType] | None
         ) = None,
diff --git a/py-polars/polars/type_aliases.py b/py-polars/polars/type_aliases.py
index 459656ac8374f..3281f6753c945 100644
--- a/py-polars/polars/type_aliases.py
+++ b/py-polars/polars/type_aliases.py
@@ -69,7 +69,7 @@
 ComparisonOperator: TypeAlias = Literal["eq", "neq", "gt", "lt", "gt_eq", "lt_eq"]
 
 # selector type
-SelectorType: TypeAlias = "_selector_proxy_"
+SelectorType = TypeVar("SelectorType", "Expr", "_selector_proxy_")
 
 # User-facing string literal types
 # The following all have an equivalent Rust enum with the same name
diff --git a/py-polars/tests/unit/dataframe/test_df.py b/py-polars/tests/unit/dataframe/test_df.py
index 1169bf76a4884..ca99fee6adb8e 100644
--- a/py-polars/tests/unit/dataframe/test_df.py
+++ b/py-polars/tests/unit/dataframe/test_df.py
@@ -2972,12 +2972,16 @@ def test_selection_regex_and_multicol() -> None:
 
 
 def test_unique_on_sorted() -> None:
+    df = pl.DataFrame(data={"a": [1, 1, 3], "b": [1, 2, 3]})
     for subset in ("a", cs.starts_with("x", "a")):
-        assert (
-            pl.DataFrame({"a": [1, 1, 3], "b": [1, 2, 3]})
-            .with_columns([pl.col("a").set_sorted()])
-            .unique(subset=subset, keep="last")  # type: ignore[arg-type]
-        ).to_dict(False) == {"a": [1, 3], "b": [2, 3]}
+        assert df.with_columns([pl.col("a").set_sorted()]).unique(  # type: ignore[type-var]
+            subset=subset, keep="last"
+        ).to_dict(
+            False
+        ) == {
+            "a": [1, 3],
+            "b": [2, 3],
+        }
 
 
 def test_len_compute(df: pl.DataFrame) -> None:
@@ -3605,46 +3609,43 @@ def test_unstack() -> None:
             "col3": pl.int_range(-9, 0, eager=True),
         }
     )
-    assert_frame_equal(
-        df.unstack(step=3, how="vertical").to_dict(False),
-        {
-            "col1_0": ["A", "B", "C"],
-            "col1_1": ["D", "E", "F"],
-            "col1_2": ["G", "H", "I"],
-            "col2_0": [0, 1, 2],
-            "col2_1": [3, 4, 5],
-            "col2_2": [6, 7, 8],
-            "col3_0": [-9, -8, -7],
-            "col3_1": [-6, -5, -4],
-            "col3_2": [-3, -2, -1],
-        },
-    )
-    assert_frame_equal(
-        df.unstack(step=3, how="horizontal").to_dict(False),
-        {
-            "col1_0": ["A", "D", "G"],
-            "col1_1": ["B", "E", "H"],
-            "col1_2": ["C", "F", "I"],
+    assert df.unstack(step=3, how="vertical").to_dict(False) == {
+        "col1_0": ["A", "B", "C"],
+        "col1_1": ["D", "E", "F"],
+        "col1_2": ["G", "H", "I"],
+        "col2_0": [0, 1, 2],
+        "col2_1": [3, 4, 5],
+        "col2_2": [6, 7, 8],
+        "col3_0": [-9, -8, -7],
+        "col3_1": [-6, -5, -4],
+        "col3_2": [-3, -2, -1],
+    }
+
+    assert df.unstack(step=3, how="horizontal").to_dict(False) == {
+        "col1_0": ["A", "D", "G"],
+        "col1_1": ["B", "E", "H"],
+        "col1_2": ["C", "F", "I"],
+        "col2_0": [0, 3, 6],
+        "col2_1": [1, 4, 7],
+        "col2_2": [2, 5, 8],
+        "col3_0": [-9, -6, -3],
+        "col3_1": [-8, -5, -2],
+        "col3_2": [-7, -4, -1],
+    }
+
+    for column_subset in (("col2", "col3"), cs.integer()):
+        assert df.unstack(  # type: ignore[type-var]
+            step=3,
+            how="horizontal",
+            columns=column_subset,
+        ).to_dict(False) == {
             "col2_0": [0, 3, 6],
             "col2_1": [1, 4, 7],
             "col2_2": [2, 5, 8],
             "col3_0": [-9, -6, -3],
             "col3_1": [-8, -5, -2],
             "col3_2": [-7, -4, -1],
-        },
-    )
-    for column_subset in (("col2", "col3"), cs.integer()):
-        assert_frame_equal(
-            df.unstack(step=3, how="horizontal", columns=column_subset).to_dict(False),  # type: ignore[arg-type]
-            {
-                "col2_0": [0, 3, 6],
-                "col2_1": [1, 4, 7],
-                "col2_2": [2, 5, 8],
-                "col3_0": [-9, -6, -3],
-                "col3_1": [-8, -5, -2],
-                "col3_2": [-7, -4, -1],
-            },
-        )
+        }
 
 
 def test_window_deadlock() -> None:
diff --git a/py-polars/tests/unit/datatypes/test_struct.py b/py-polars/tests/unit/datatypes/test_struct.py
index 50a4f21315aa6..705714795f75f 100644
--- a/py-polars/tests/unit/datatypes/test_struct.py
+++ b/py-polars/tests/unit/datatypes/test_struct.py
@@ -115,7 +115,7 @@ def test_struct_unnesting() -> None:
         }
     )
     for cols in ("foo", cs.ends_with("oo")):
-        out = df.unnest(cols)  # type: ignore[arg-type]
+        out = df.unnest(cols)  # type: ignore[type-var]
         assert_frame_equal(out, expected)
 
     out = (