fix(python): Correct type hint for map_columns function parameter

[veeceey]

Summary

Fixes #26371

The type hint for the function parameter in DataFrame.map_columns was Callable[[Series], Series], which incorrectly implied the function only accepts a single Series argument. However, the implementation passes *args and **kwargs to the function (line 6964).

This PR updates the type hint to Callable[Concatenate[Series, P], Series] using ParamSpec to accurately reflect that the function can accept additional positional and keyword arguments after the Series parameter.

Changes

• Updated function parameter type hint from Callable[[Series], Series] to Callable[Concatenate[Series, P], Series]
• Updated *args type hint from Any to P.args
• Updated **kwargs type hint from Any to P.kwargs
• The ParamSpec P was already defined in the TYPE_CHECKING block (line 206)

Test Plan

Created and ran a test script demonstrating the fix:

import polars as pl

def function(s: pl.Series, a: int, b: int = 0) -> pl.Series:
    return a * s + b

df = pl.DataFrame({"a": [1, 2]})
df.map_columns("a", function, 1, b=2)  # Now correctly type-checked

The code executes correctly and type checkers can now properly validate that functions passed to map_columns can accept additional arguments beyond the Series parameter.

[veeceey]

Manual Test Results

I've verified the fix works correctly with various test cases:

import polars as pl


def function_with_args(s: pl.Series, a: int, b: int = 0) -> pl.Series:
    """Function that takes Series plus additional positional and keyword arguments."""
    return a * s + b


def simple_function(s: pl.Series) -> pl.Series:
    """Function that only takes a Series (original functionality)."""
    return s * 2


# Test 1: Function with additional positional and keyword arguments
df = pl.DataFrame({"a": [1, 2]})
result1 = df.map_columns("a", function_with_args, 1, b=2)

# Test 2: Multiple columns with the same function
df2 = pl.DataFrame({"a": [1, 2], "b": [3, 4]})
result2 = df2.map_columns(["a", "b"], function_with_args, 2, b=1)

# Test 3: Simple function without extra args (backward compatibility)
result3 = df.map_columns("a", simple_function)

# Test 4: Lambda with extra args
result4 = df.map_columns("a", lambda s, mult: s * mult, 3)

All tests execute successfully and produce correct results. The new type hints properly support:

• Functions with additional positional arguments
• Functions with keyword arguments
• Simple functions with only a Series parameter (backward compatibility)
• Lambda functions with extra parameters

Type checkers (mypy/pyright) will now correctly validate that custom functions passed to map_columns can accept the additional arguments being passed.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 80.83%. Comparing base (8a81b0e) to head (4619921).

Additional details and impacted files

@@            Coverage Diff             @@
##             main   #26487      +/-   ##
==========================================
- Coverage   81.08%   80.83%   -0.25%     
==========================================
  Files        1783     1783              
  Lines      243288   243288              
  Branches     3074     3074              
==========================================
- Hits       197263   196668     -595     
- Misses      45222    45817     +595     
  Partials      803      803              

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[veeceey]

Friendly ping - any chance someone could take a look at this when they get a chance? Happy to make any changes if needed.