Fix a number of mypy ignore cases (#1844)

py-polars/polars/internals/lazy_functions.py

@@ -608,7 +608,7 @@ def map(
 
 
 def apply(
-    exprs: Union[tp.List[str], tp.List["pli.Expr"]],
+    exprs: tp.List[Union[str, "pli.Expr"]],
     f: Callable[[tp.List["pli.Series"]], "pli.Series"],
     return_dtype: Optional[Type[DataType]] = None,
 ) -> "pli.Expr":
@@ -867,13 +867,13 @@ def argsort_by(
 
 
 def _datetime(
-    year: "pli.Expr",
-    month: "pli.Expr",
-    day: "pli.Expr",
-    hour: Optional["pli.Expr"] = None,
-    minute: Optional["pli.Expr"] = None,
-    second: Optional["pli.Expr"] = None,
-    millisecond: Optional["pli.Expr"] = None,
+    year: Union["pli.Expr", str],
+    month: Union["pli.Expr", str],
+    day: Union["pli.Expr", str],
+    hour: Optional[Union["pli.Expr", str]] = None,
+    minute: Optional[Union["pli.Expr", str]] = None,
+    second: Optional[Union["pli.Expr", str]] = None,
+    millisecond: Optional[Union["pli.Expr", str]] = None,
 ) -> "pli.Expr":
     """
     Create polars Datetime from distinct time components.
@@ -919,7 +919,11 @@ def _datetime(
     )
 
 
-def _date(year: "pli.Expr", month: "pli.Expr", day: "pli.Expr") -> "pli.Expr":
+def _date(
+    year: Union["pli.Expr", str],
+    month: Union["pli.Expr", str],
+    day: Union["pli.Expr", str],
+) -> "pli.Expr":
     """
     Create polars Date from distinct time components.
 
@@ -1004,7 +1008,7 @@ def format(fstring: str, *args: Union["pli.Expr", str]) -> "pli.Expr":
     return concat_str(exprs, sep="")
 
 
-def concat_list(exprs: tp.List["pli.Expr"]) -> "pli.Expr":
+def concat_list(exprs: tp.List[Union[str, "pli.Expr"]]) -> "pli.Expr":
     """
     Concat the arrays in a Series dtype List in linear time.
 

py-polars/polars/internals/series.py

@@ -1213,7 +1213,7 @@ def take_every(self, n: int) -> "Series":
         """
         return wrap_s(self._s.take_every(n))
 
-    def sort(self, in_place: bool = False, reverse: bool = False) -> Optional["Series"]:
+    def sort(self, in_place: bool = False, reverse: bool = False) -> "Series":
         """
         Sort this Series.
 
@@ -1249,7 +1249,7 @@ def sort(self, in_place: bool = False, reverse: bool = False) -> Optional["Serie
         """
         if in_place:
             self._s.sort_in_place(reverse)
-            return None
+            return self
         else:
             return wrap_s(self._s.sort(reverse))
 
@@ -1735,7 +1735,7 @@ def to_list(self, use_pyarrow: bool = False) -> tp.List[Optional[Any]]:
     def __iter__(self) -> "SeriesIter":
         return SeriesIter(self.len(), self)
 
-    def rechunk(self, in_place: bool = False) -> Optional["Series"]:
+    def rechunk(self, in_place: bool = False) -> "Series":
         """
         Create a single chunk of memory for this Series.
 
@@ -1746,7 +1746,7 @@ def rechunk(self, in_place: bool = False) -> Optional["Series"]:
         """
         opt_s = self._s.rechunk(in_place)
         if in_place:
-            return None
+            return self
         else:
             return wrap_s(opt_s)
 

py-polars/tests/files/test_functions.py

@@ -13,9 +13,9 @@ def test_date_datetime() -> None:
 
     out = df.select(
         [
-            pl.all(),  # type: ignore
-            pl.datetime("year", "month", "day", "hour").dt.hour().alias("h2"),  # type: ignore
-            pl.date("year", "month", "day").dt.day().alias("date"),  # type: ignore
+            pl.all(),
+            pl.datetime("year", "month", "day", "hour").dt.hour().alias("h2"),
+            pl.date("year", "month", "day").dt.day().alias("date"),
         ]
     )
 

py-polars/tests/test_apply.py

@@ -14,16 +14,16 @@ def test_apply_none() -> None:
 
     out = (
         df.groupby("g", maintain_order=True).agg(
-            pl.apply(  # type: ignore
-                exprs=["a", pl.col("b") ** 4, pl.col("a") / 4],  # type: ignore
+            pl.apply(
+                exprs=["a", pl.col("b") ** 4, pl.col("a") / 4],
                 f=lambda x: x[0] * x[1] + x[2].sum(),
             ).alias("multiple")
         )
     )["multiple"]
     assert out[0].to_list() == [4.75, 326.75, 82.75]
     assert out[1].to_list() == [238.75, 3418849.75, 372.75]
 
-    out = df.select(pl.map(exprs=["a", "b"], f=lambda s: s[0] * s[1]))  # type: ignore
+    out = df.select(pl.map(exprs=["a", "b"], f=lambda s: s[0] * s[1]))
     assert out["a"].to_list() == (df["a"] * df["b"]).to_list()
 
     # check if we can return None

py-polars/tests/test_datelike.py

@@ -19,7 +19,7 @@ def test_fill_null() -> None:
     dt3 = date(2001, 1, 3)
     s = pl.Series("a", [dt1, dt2, dt3, None])
     dt_2 = date(2001, 1, 4)
-    for fill_val in (dt_2, pl.lit(dt_2)):  # type: ignore
+    for fill_val in (dt_2, pl.lit(dt_2)):
         out = s.fill_null(fill_val)  # type: ignore
 
         assert out.null_count() == 0

py-polars/tests/test_exprs.py

@@ -3,8 +3,8 @@
 
 def test_horizontal_agg(fruits_cars: pl.DataFrame) -> None:
     df = fruits_cars
-    out = df.select(pl.max([pl.col("A"), pl.col("B")]))  # type: ignore
+    out = df.select(pl.max([pl.col("A"), pl.col("B")]))
     assert out[:, 0].to_list() == [5, 4, 3, 4, 5]
 
-    out = df.select(pl.min([pl.col("A"), pl.col("B")]))  # type: ignore
+    out = df.select(pl.min([pl.col("A"), pl.col("B")]))
     assert out[:, 0].to_list() == [1, 2, 3, 2, 1]

py-polars/tests/test_series.py

@@ -49,7 +49,7 @@ def test_init_inputs() -> None:
 def test_concat() -> None:
     s = pl.Series("a", [2, 1, 3])
 
-    assert pl.concat([s, s]).len() == 6  # type: ignore
+    assert pl.concat([s, s]).len() == 6
     # check if s remains unchanged
     assert s.len() == 3
 
@@ -193,8 +193,7 @@ def test_to_python() -> None:
 
 def test_sort() -> None:
     a = pl.Series("a", [2, 1, 3])
-    a_sorted: pl.Series = a.sort()  # type: ignore
-    assert a_sorted.to_list() == [1, 2, 3]
+    assert a.sort().to_list() == [1, 2, 3]
     assert a.sort(reverse=True) == [3, 2, 1]
 
 
@@ -203,7 +202,7 @@ def test_rechunk() -> None:
     b = pl.Series("b", [4, 5, 6])
     a.append(b)
     assert a.n_chunks() == 2
-    assert a.rechunk(in_place=False).n_chunks() == 1  # type: ignore
+    assert a.rechunk(in_place=False).n_chunks() == 1
     a.rechunk(in_place=True)
     assert a.n_chunks() == 1
 
@@ -306,13 +305,13 @@ def test_shift() -> None:
 
 
 def test_rolling() -> None:
-    a = pl.Series("a", [1, 2, 3, 2, 1])  # type: ignore
+    a = pl.Series("a", [1, 2, 3, 2, 1])
     assert a.rolling_min(2).to_list() == [None, 1, 2, 2, 1]
     assert a.rolling_max(2).to_list() == [None, 2, 3, 3, 2]
     assert a.rolling_sum(2).to_list() == [None, 3, 5, 5, 3]
     assert a.rolling_mean(2).to_list() == [None, 1.5, 2.5, 2.5, 1.5]
-    assert np.isclose(a.rolling_std(2).to_list()[1], 0.7071067811865476)  # type: ignore
-    assert np.isclose(a.rolling_var(2).to_list()[1], 0.5)  # type: ignore
+    assert a.rolling_std(2).to_list()[1] == pytest.approx(0.7071067811865476)
+    assert a.rolling_var(2).to_list()[1] == pytest.approx(0.5)
     assert a.rolling_median(4).to_list() == [None, None, None, 2, 2]
     assert a.rolling_quantile(3, 0.5).to_list() == [None, None, 2, 2, 2]
     assert a.rolling_skew(4).null_count() == 3
@@ -423,15 +422,15 @@ def test_str_slice() -> None:
 
 def test_arange_expr() -> None:
     df = pl.DataFrame({"a": ["foobar", "barfoo"]})
-    out = df[[pl.arange(0, pl.col("a").count() * 10)]]  # type: ignore
+    out = df[[pl.arange(0, pl.col("a").count() * 10)]]
     assert out.shape == (20, 1)
     assert out.select_at_idx(0)[-1] == 19
 
     # eager arange
-    out = pl.arange(0, 10, 2, eager=True)  # type: ignore
+    out = pl.arange(0, 10, 2, eager=True)
     assert out == [0, 2, 4, 8, 8]
 
-    out = pl.arange(pl.Series([0, 19]), pl.Series([3, 39]), step=2, eager=True)  # type: ignore
+    out = pl.arange(pl.Series([0, 19]), pl.Series([3, 39]), step=2, eager=True)
     assert out.dtype == pl.List
     assert out[0].to_list() == [0, 2]
 
@@ -517,8 +516,8 @@ def test_diff_dispatch() -> None:
 def test_skew_dispatch() -> None:
     s = pl.Series("a", [1, 2, 3, 2, 2, 3, 0])
 
-    assert np.isclose(s.skew(True), -0.5953924651018018)  # type: ignore
-    assert np.isclose(s.skew(False), -0.7717168360221258)  # type: ignore
+    assert s.skew(True) == pytest.approx(-0.5953924651018018)
+    assert s.skew(False) == pytest.approx(-0.7717168360221258)
 
     df = pl.DataFrame([s])
     assert np.isclose(df.select(pl.col("a").skew(False))["a"][0], -0.7717168360221258)
@@ -528,7 +527,7 @@ def test_kurtosis_dispatch() -> None:
     s = pl.Series("a", [1, 2, 3, 2, 2, 3, 0])
     expected = -0.6406250000000004
 
-    assert np.isclose(s.kurtosis(), expected)  # type: ignore
+    assert s.kurtosis() == pytest.approx(expected)
     df = pl.DataFrame([s])
     assert np.isclose(df.select(pl.col("a").kurtosis())["a"][0], expected)
 
@@ -573,7 +572,7 @@ def test_list_concat_dispatch() -> None:
     assert out.series_equal(expected)
 
     df = pl.DataFrame([s0, s1])
-    assert df.select(pl.concat_list(["a", "b"]).alias("concat"))["concat"].series_equal(  # type: ignore
+    assert df.select(pl.concat_list(["a", "b"]).alias("concat"))["concat"].series_equal(
         expected
     )
     assert df.select(pl.col("a").arr.concat("b").alias("concat"))[
@@ -763,7 +762,7 @@ def test_value_counts() -> None:
     s = pl.Series("a", [1, 2, 2, 3])
     result = s.value_counts()
     expected = pl.DataFrame({"a": [1, 2, 3], "counts": [1, 2, 1]})
-    result_sorted: pl.DataFrame = result.sort("a")  # type: ignore
+    result_sorted: pl.DataFrame = result.sort("a")
     assert result_sorted.frame_equal(expected)
 
 

py-polars/tests/test_sort.py

@@ -23,11 +23,11 @@ def test_sort_dates_multiples() -> None:
     expected = [4, 5, 2, 3, 1]
 
     # datetime
-    out: pl.DataFrame = df.sort(["date", "values"])  # type: ignore
+    out: pl.DataFrame = df.sort(["date", "values"])
     assert out["values"].to_list() == expected
 
     # Date
-    out = df.with_column(pl.col("date").cast(pl.Date)).sort(["date", "values"])  # type: ignore
+    out = df.with_column(pl.col("date").cast(pl.Date)).sort(["date", "values"])
     assert out["values"].to_list() == expected