lint/typing

pola-rs · Aug 4, 2023 · aa52b3a · aa52b3a
1 parent 91b1256
commit aa52b3a
Show file tree

Hide file tree

Showing 4 changed files with 43 additions and 41 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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 = (