Repository: pandas-dev/pandas-stubs (language: python)

Issue:
No overload variant of "apply" of "Styler" matches argument types "Callable[[Series[int], str], list[str]]", "str", "int"
**Describe the bug**
Following the example in [pandas website](https://pandas.pydata.org/pandas-docs/stable/reference/api/pandas.io.formats.style.Styler.apply.html), I defined a function to highlight the maximum of each row. However, `mypy` does not recognize the call signature.

**To Reproduce**
1. Minimal working example
```
from typing import List

import numpy as np
import pandas as pd


def highlight_max(x: 'pd.Series[int]', color: str) -> List[str]:
    """Highlight max."""
    return [f"color: {color}" if val == x.max() else "" for val in x]


df = pd.DataFrame(np.random.randint(0, high=100, size=(10, 3)))
styler = df.style.apply(highlight_max, color='red', axis=1)
styler.to_html("output.html")
```
2. I'm using `mypy`
3. Show the error message received from that type checker while checking your example.
In terminal
```
$ mypy --explicit-package-bases --strict .
proj.py:14: error: No overload variant of "apply" of "Styler" matches argument types "Callable[[Series[int], str], list[str]]", "str", "int"  [call-overload]
proj.py:14: note: Possible overload variants:
proj.py:14: note:     def apply(self, func: Callable[[Series[Any]], Union[list[Any], Series[Any]]], axis: Union[Literal['index', 0], Literal['columns', 1]] = ..., subset: Optional[Union[_IndexSlice, slice, tuple[slice, ...], list[Any], Index[Any]]] = ..., **kwargs: Any) -> Styler
proj.py:14: note:     def apply(self, func: Callable[[DataFrame], Union[ndarray[Any, dtype[Any]], DataFrame]], axis: None, subset: Optional[Union[_IndexSlice, slice, tuple[slice, ...], list[Any], Index[Any]]] = ..., **kwargs: Any) -> Styler
Found 1 error in 1 file (checked 1 source file)
```


**Please complete the following information:**
 - OS: Mac OS
 - OS Version 14.4.1 
 - python version: Python 3.9.6
 - version of type checker: mypy 1.8.0 (compiled: yes)
 - version of installed `pandas-stubs`: from `pip list`: `pandas-stubs                      2.2.1.240316`



Discussion / hints:
Thanks for the report.  This is non-trivial to fix.  We have overloads that test whether you are passing a `Callable` that takes a `Series` or `DataFrame`, as that pairs with the `axis` argument of `Styler.apply()`.  But there really aren't signatures for callables that have keyword arguments, and the recommended way to handle that is to use `Protocol`.  But that has issues - see https://github.com/microsoft/pyright/issues/7898 because it forces the callable to have a specific name for the first argument.

I'll have to see what the typing experts have to say about this.


Fix patch:
diff --git a/pandas-stubs/io/formats/style.pyi b/pandas-stubs/io/formats/style.pyi
index dc2b5382..4724726f 100644
--- a/pandas-stubs/io/formats/style.pyi
+++ b/pandas-stubs/io/formats/style.pyi
@@ -5,6 +5,7 @@ from collections.abc import (
 from typing import (
     Any,
     Literal,
+    Protocol,
     overload,
 )
 
@@ -40,6 +41,16 @@ from pandas.io.formats.style_render import (
     Subset,
 )
 
+class _SeriesFunc(Protocol):
+    def __call__(
+        self, series: Series, /, *args: Any, **kwargs: Any
+    ) -> list | Series: ...
+
+class _DataFrameFunc(Protocol):
+    def __call__(
+        self, series: DataFrame, /, *args: Any, **kwargs: Any
+    ) -> npt.NDArray | DataFrame: ...
+
 class Styler(StylerRenderer):
     def __init__(
         self,
@@ -198,7 +209,7 @@ class Styler(StylerRenderer):
     @overload
     def apply(
         self,
-        func: Callable[[Series], list | Series],
+        func: _SeriesFunc | Callable[[Series], list | Series],
         axis: Axis = ...,
         subset: Subset | None = ...,
         **kwargs: Any,
@@ -206,7 +217,7 @@ class Styler(StylerRenderer):
     @overload
     def apply(
         self,
-        func: Callable[[DataFrame], npt.NDArray | DataFrame],
+        func: _DataFrameFunc | Callable[[DataFrame], npt.NDArray | DataFrame],
         axis: None,
         subset: Subset | None = ...,
         **kwargs: Any,


Test patch:
diff --git a/tests/test_styler.py b/tests/test_styler.py
index 7dd4003e..42f7830d 100644
--- a/tests/test_styler.py
+++ b/tests/test_styler.py
@@ -52,6 +52,14 @@ def h(df: DataFrame) -> DataFrame:
 
     check(assert_type(DF.style.apply(h, axis=None), Styler), Styler)
 
+    # GH 919
+    def highlight_max(x: Series[int], /, color: str) -> list[str]:
+        return [f"color: {color}" if val == x.max() else "" for val in x]
+
+    check(
+        assert_type(DF.style.apply(highlight_max, color="red", axis=1), Styler), Styler
+    )
+
 
 def test_apply_index() -> None:
     def f(s: Series) -> npt.NDArray[np.str_]: