pola-rs / polars

Dataframes powered by a multithreaded, vectorized query engine, written in Rust
https://docs.pola.rs
Other
26.9k stars 1.65k forks source link

fix(python): Improve `map_elements` typing #16257

Closed max-muoto closed 1 week ago

max-muoto commented 2 weeks ago

Map elements accepts a callable Callable[[Series], Series] | Callable[[Any], Any], this for all intensive purposes should just accept Callable[[Any], Any]. Previously, the former was fine in type-checkers such as Pyright, however as Eric Traut explains here, there's not a great way for Pyright to deterministically determine which to use when evaluating lambdas, which can lead to Pyright errors when calling this function. Take this example, which will lead to an error in the latest version of Pyright, which I would imagine should be fully supported:

import polars as pl

df = pl.DataFrame(
    {
        "a": [1, 2, None],
        "b": [1.0, 2.0, 3.0],
        "c": ["a", "b", "c"],
    }
)

def transform_int(x: int) -> int:
    return x + 1

# Argument of type "Series" cannot be assigned to parameter "x" of type "int" in function "transform_int" "Series" is incompatible with "int"
df.with_columns([pl.col("a").map_elements(lambda x: transform_int(x) if x else None)])
codecov[bot] commented 2 weeks ago

Codecov Report

All modified and coverable lines are covered by tests :white_check_mark:

Project coverage is 80.87%. Comparing base (79a2cf1) to head (93e2fc0). Report is 8 commits behind head on main.

Additional details and impacted files ```diff @@ Coverage Diff @@ ## main #16257 +/- ## ========================================== - Coverage 80.92% 80.87% -0.06% ========================================== Files 1393 1394 +1 Lines 179568 179488 -80 Branches 2909 2913 +4 ========================================== - Hits 145321 145166 -155 - Misses 33742 33819 +77 + Partials 505 503 -2 ```

:umbrella: View full report in Codecov by Sentry.
:loudspeaker: Have feedback on the report? Share it here.