Passing a callable to DataFrame.mask's 'cond' argument results in arg-type error (mypy)

[georgesittas]

Describe the bug

The cond argument of the DataFrame.mask method currently doesn't expect a callable, even though it should, according to the method's documentation page:

If cond is callable, it is computed on the Series/DataFrame and should return boolean Series/DataFrame or array. The callable must not change input Series/DataFrame (though pandas doesn’t check it).

To Reproduce

1. Minimally runnable pandas example that is not properly checked by the stubs:

import pandas as pd
df = pd.DataFrame({"a": [1]}).mask(cond=lambda x: x.notna(), other=2)
print(df)
#    a
# 0  2

2. Type checker: mypy

3. Show the error message received from that type checker while checking your example.

(.venv) ➜  python -m mypy test.py
test.py:2: error: Argument "cond" to "mask" of "DataFrame" has incompatible type "Callable[[Any], Any]"; expected "Series[Any] | DataFrame | ndarray[Any, Any]"  [arg-type]
Found 1 error in 1 file (checked 1 source file)

Please complete the following information:

• OS: MacOS
• OS Version: Sonoma 14.2.1
• python version: Python 3.10.14
• version of type checker: 1.10.0
• version of installed pandas-stubs: 2.2.1.240316

Additional context

I've also tried using typing.cast to cast the lambda into Callable[[DataFrame], DataFrame], but it doesn't work, as expected. Happy to provide more details if needed.

[Dr-Irv]

Thanks for the report. It appears we have Callable in Series.mask(), but not DataFrame.mask().

PR with tests welcome.

[georgesittas]

Gotcha! Appreciate the quick response, will try to make a PR for this 👍