refactor(python): Fewer typing.no_type_check (#16497)

Co-authored-by: Stijn de Gooijer <stijndegooijer@gmail.com>

py-polars/docs/source/conf.py

@@ -264,14 +264,22 @@ def _minify_classpaths(s: str) -> str:
     )
 
 
-def process_signature(app, what, name, obj, opts, sig, ret):  # noqa: D103
+def process_signature(  # noqa: D103
+    app: object,
+    what: object,
+    name: object,
+    obj: object,
+    opts: object,
+    sig: str,
+    ret: str,
+) -> tuple[str, str]:
     return (
         _minify_classpaths(sig) if sig else sig,
         _minify_classpaths(ret) if ret else ret,
     )
 
 
-def setup(app):  # noqa: D103
+def setup(app: Any) -> None:  # noqa: D103
     # TODO: a handful of methods do not seem to trigger the event for
     #  some reason (possibly @overloads?) - investigate further...
     app.connect("autodoc-process-signature", process_signature)

py-polars/tests/unit/interop/numpy/test_ufunc_expr.py

@@ -136,8 +136,8 @@ def test_grouped_ufunc() -> None:
 def make_gufunc_mean() -> Callable[[pl.Series], pl.Series]:
     numba = pytest.importorskip("numba")
 
-    @numba.guvectorize([(numba.float64[:], numba.float64[:])], "(n)->(n)")
-    def gufunc_mean(arr, result):  # type: ignore[no-untyped-def]
+    @numba.guvectorize([(numba.float64[:], numba.float64[:])], "(n)->(n)")  # type: ignore[misc]
+    def gufunc_mean(arr: Any, result: Any) -> None:
         mean = arr.mean()
         for i in range(len(arr)):
             result[i] = mean + i

py-polars/tests/unit/interop/numpy/test_ufunc_series.py

@@ -1,4 +1,4 @@
-from typing import Callable, cast
+from typing import Any, Callable, cast
 
 import numpy as np
 import pytest
@@ -125,8 +125,8 @@ def test_numpy_string_array() -> None:
 def make_add_one() -> Callable[[pl.Series], pl.Series]:
     numba = pytest.importorskip("numba")
 
-    @numba.guvectorize([(numba.float64[:], numba.float64[:])], "(n)->(n)")
-    def add_one(arr, result):  # type: ignore[no-untyped-def]
+    @numba.guvectorize([(numba.float64[:], numba.float64[:])], "(n)->(n)")  # type: ignore[misc]
+    def add_one(arr: Any, result: Any) -> None:
         for i in range(len(arr)):
             result[i] = arr[i] + 1.0
 
@@ -164,8 +164,8 @@ def make_divide_by_sum() -> Callable[[pl.Series, pl.Series], pl.Series]:
     numba = pytest.importorskip("numba")
     float64 = numba.float64
 
-    @numba.guvectorize([(float64[:], float64[:], float64[:])], "(n),(m)->(m)")
-    def divide_by_sum(arr, arr2, result):  # type: ignore[no-untyped-def]
+    @numba.guvectorize([(float64[:], float64[:], float64[:])], "(n),(m)->(m)")  # type: ignore[misc]
+    def divide_by_sum(arr: Any, arr2: Any, result: Any) -> None:
         total = arr.sum()
         for i in range(len(arr2)):
             result[i] = arr2[i] / total

py-polars/tests/unit/lazyframe/cuda/test_node_visitor.py

@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-import typing
 from functools import lru_cache, partial
 from typing import TYPE_CHECKING, Any, Callable
 
@@ -12,41 +11,42 @@
     import pandas as pd
 
 
-@typing.no_type_check
 def test_run_on_pandas() -> None:
     # Simple join example, missing multiple columns, slices, etc.
-    @typing.no_type_check
-    def join(inputs: list[Callable], obj: Any, _node_traverer: Any) -> Callable:
+    def join(
+        inputs: list[Callable[[], pd.DataFrame]], obj: Any, _node_traverer: Any
+    ) -> Callable[[], pd.DataFrame]:
         assert len(obj.left_on) == 1
         assert len(obj.right_on) == 1
         left_on = obj.left_on[0].output_name
         right_on = obj.right_on[0].output_name
 
         assert len(inputs) == 2
 
-        def run(inputs: list[Callable]):
+        def run(inputs: list[Callable[[], pd.DataFrame]]) -> pd.DataFrame:
             # materialize inputs
-            inputs = [call() for call in inputs]
-            return inputs[0].merge(inputs[1], left_on=left_on, right_on=right_on)
+            dataframes = [call() for call in inputs]
+            return dataframes[0].merge(
+                dataframes[1], left_on=left_on, right_on=right_on
+            )
 
         return partial(run, inputs)
 
     # Simple scan example, missing predicates, columns pruning, slices, etc.
-    @typing.no_type_check
-    def df_scan(_inputs: None, obj: Any, _: Any) -> pd.DataFrame:
+    def df_scan(_inputs: None, obj: Any, _: Any) -> Callable[[], pd.DataFrame]:
         assert obj.selection is None
         return lambda: wrap_df(obj.df).to_pandas()
 
     @lru_cache(1)
-    @typing.no_type_check
-    def get_node_converters():
+    def get_node_converters() -> (
+        dict[type, Callable[[Any, Any, Any], Callable[[], pd.DataFrame]]]
+    ):
         return {
             _ir_nodes.Join: join,
             _ir_nodes.DataFrameScan: df_scan,
         }
 
-    @typing.no_type_check
-    def get_input(node_traverser):
+    def get_input(node_traverser: Any) -> Callable[[], pd.DataFrame]:
         current_node = node_traverser.get_node()
 
         inputs_callable = []
@@ -60,13 +60,11 @@ def get_input(node_traverser):
             inputs_callable, ir_node, node_traverser
         )
 
-    @typing.no_type_check
-    def run_on_pandas(node_traverser) -> None:
+    def run_on_pandas(node_traverser: Any) -> None:
         current_node = node_traverser.get_node()
 
         callback = get_input(node_traverser)
 
-        @typing.no_type_check
         def run_callback(
             columns: list[str] | None, _: Any, n_rows: int | None
         ) -> pl.DataFrame:
@@ -83,7 +81,9 @@ def run_callback(
     q1 = pl.LazyFrame({"foo": [1, 2, 3]})
     q2 = pl.LazyFrame({"foo": [1], "bar": [2]})
     q = q1.join(q2, on="foo")
-    assert q.collect(post_opt_callback=run_on_pandas).to_dict(as_series=False) == {
+    assert q.collect(
+        post_opt_callback=run_on_pandas  # type: ignore[call-overload]
+    ).to_dict(as_series=False) == {
         "foo": [2],
         "bar": [4],
     }

py-polars/tests/unit/operations/map/test_inefficient_map_warning.py

@@ -506,7 +506,7 @@ def test_omit_implicit_bool() -> None:
 
 
 def test_partial_functions_13523() -> None:
-    def plus(value, amount: int):  # type: ignore[no-untyped-def]
+    def plus(value: int, amount: int) -> int:
         return value + amount
 
     data = {"a": [1, 2], "b": [3, 4]}

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

@@ -966,8 +966,7 @@ def test_cross_join_slice_pushdown() -> None:
 
 
 @pytest.mark.parametrize("how", ["left", "inner"])
-@typing.no_type_check
-def test_join_coalesce(how: str) -> None:
+def test_join_coalesce(how: JoinStrategy) -> None:
     a = pl.LazyFrame({"a": [1, 2], "b": [1, 2]})
     b = pl.LazyFrame(
         {
@@ -995,8 +994,7 @@ def test_join_coalesce(how: str) -> None:
 
 
 @pytest.mark.parametrize("how", ["left", "inner", "full", "outer"])
-@typing.no_type_check
-def test_join_empties(how: str) -> None:
+def test_join_empties(how: JoinStrategy) -> None:
     df1 = pl.DataFrame({"col1": [], "col2": [], "col3": []})
     df2 = pl.DataFrame({"col2": [], "col4": [], "col5": []})
 

py-polars/tests/unit/test_cse.py

@@ -1,5 +1,6 @@
+from __future__ import annotations
+
 import re
-import typing
 from datetime import date, datetime, timedelta
 from tempfile import NamedTemporaryFile
 from typing import Any
@@ -598,7 +599,6 @@ def test_cse_11958() -> None:
     }
 
 
-@typing.no_type_check
 def test_cse_14047() -> None:
     ldf = pl.LazyFrame(
         {
@@ -615,7 +615,7 @@ def test_cse_14047() -> None:
 
     def count_diff(
         price: pl.Expr, upper_bound: float = 0.1, lower_bound: float = 0.001
-    ):
+    ) -> pl.Expr:
         span_end_to_curr = (
             price.count()
             .cast(int)
@@ -630,7 +630,7 @@ def count_diff(
             f"count_diff_{upper_bound}_{lower_bound}"
         )
 
-    def s_per_count(count_diff, span) -> pl.Expr:
+    def s_per_count(count_diff: pl.Expr, span: tuple[float, float]) -> pl.Expr:
         return (span[1] * 1000 - span[0] * 1000) / count_diff
 
     spans = [(0.001, 0.1), (1, 10)]