ENH: Improve typing of crosstab

pandas-stubs/core/reshape/pivot.pyi

@@ -1,18 +1,25 @@
 from typing import (
     Callable,
+    Hashable,
     Literal,
     Sequence,
+    TypeVar,
 )
 
+import numpy as np
 from pandas.core.frame import DataFrame
 from pandas.core.groupby.grouper import Grouper
 from pandas.core.series import Series
 
 from pandas._typing import (
+    HashableT,
     IndexLabel,
     Scalar,
+    npt,
 )
 
+_HashableT2 = TypeVar("_HashableT2", bound=Hashable)
+
 def pivot_table(
     data: DataFrame,
     values: str | None = ...,
@@ -33,12 +40,12 @@ def pivot(
     values: IndexLabel = ...,
 ) -> DataFrame: ...
 def crosstab(
-    index: Sequence | Series,
-    columns: Sequence | Series,
-    values: Sequence | None = ...,
-    rownames: Sequence | None = ...,
-    colnames: Sequence | None = ...,
-    aggfunc: Callable | None = ...,
+    index: list | npt.NDArray | Series | list[Sequence | npt.NDArray | Series],
+    columns: list | npt.NDArray | Series | list[Sequence | npt.NDArray | Series],
+    values: list | npt.NDArray | Series | None = ...,
+    rownames: list[HashableT] | None = ...,
+    colnames: list[_HashableT2] | None = ...,
+    aggfunc: str | np.ufunc | Callable[[Series], float] | None = ...,
     margins: bool = ...,
     margins_name: str = ...,
     dropna: bool = ...,

tests/test_pandas.py

@@ -1103,3 +1103,99 @@ def test_merge_asof() -> None:
         ),
         pd.DataFrame,
     )
+
+
+def test_crosstab_args() -> None:
+    a = [1, 2, 3, 4, 1, 2, 3, 4, 1, 2]
+    b: list = [4, 5, 6, 3, 4, 3, 5, 6, 5, 5]
+    c = [1, 3, 2, 3, 1, 2, 3, 1, 3, 2]
+    check(assert_type(pd.crosstab(a, b), pd.DataFrame), pd.DataFrame)
+    check(assert_type(pd.crosstab(a, [b, c]), pd.DataFrame), pd.DataFrame)
+    check(
+        assert_type(pd.crosstab(np.array(a), np.array(b)), pd.DataFrame), pd.DataFrame
+    )
+    check(
+        assert_type(pd.crosstab(np.array(a), [np.array(b), np.array(c)]), pd.DataFrame),
+        pd.DataFrame,
+    )
+    check(
+        assert_type(pd.crosstab(pd.Series(a), pd.Series(b)), pd.DataFrame), pd.DataFrame
+    )
+    check(
+        assert_type(
+            pd.crosstab(pd.Series(a), [pd.Series(b), pd.Series(c)]), pd.DataFrame
+        ),
+        pd.DataFrame,
+    )
+    values = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
+    check(
+        assert_type(pd.crosstab(a, b, values=values, aggfunc=np.sum), pd.DataFrame),
+        pd.DataFrame,
+    )
+    check(
+        assert_type(
+            pd.crosstab(a, b, values=np.array(values), aggfunc=np.sum), pd.DataFrame
+        ),
+        pd.DataFrame,
+    )
+
+    check(
+        assert_type(
+            pd.crosstab(a, b, values=pd.Series(values), aggfunc=np.sum), pd.DataFrame
+        ),
+        pd.DataFrame,
+    )
+    check(
+        assert_type(pd.crosstab(a, b, values=values, aggfunc=np.mean), pd.DataFrame),
+        pd.DataFrame,
+    )
+    check(
+        assert_type(pd.crosstab(a, b, values=values, aggfunc="mean"), pd.DataFrame),
+        pd.DataFrame,
+    )
+
+    def m(x: pd.Series) -> float:
+        return x.sum() / len(x)
+
+    check(
+        assert_type(pd.crosstab(a, b, values=values, aggfunc=m), pd.DataFrame),
+        pd.DataFrame,
+    )
+
+    def m2(x: pd.Series) -> int:
+        return int(x.sum())
+
+    check(
+        assert_type(pd.crosstab(a, b, values=values, aggfunc=m2), pd.DataFrame),
+        pd.DataFrame,
+    )
+    check(
+        assert_type(
+            pd.crosstab(a, b, margins=True, margins_name="something"), pd.DataFrame
+        ),
+        pd.DataFrame,
+    )
+    check(
+        assert_type(pd.crosstab(a, b, margins=True, dropna=True), pd.DataFrame),
+        pd.DataFrame,
+    )
+    check(
+        assert_type(pd.crosstab(a, b, colnames=["a"], rownames=["b"]), pd.DataFrame),
+        pd.DataFrame,
+    )
+    rownames: list[tuple] = [("b", 1)]
+    colnames: list[tuple] = [("a",)]
+    check(
+        assert_type(
+            pd.crosstab(a, b, colnames=colnames, rownames=rownames),
+            pd.DataFrame,
+        ),
+        pd.DataFrame,
+    )
+    check(assert_type(pd.crosstab(a, b, normalize=0), pd.DataFrame), pd.DataFrame)
+    check(assert_type(pd.crosstab(a, b, normalize=1), pd.DataFrame), pd.DataFrame)
+    check(assert_type(pd.crosstab(a, b, normalize="all"), pd.DataFrame), pd.DataFrame)
+    check(assert_type(pd.crosstab(a, b, normalize="index"), pd.DataFrame), pd.DataFrame)
+    check(
+        assert_type(pd.crosstab(a, b, normalize="columns"), pd.DataFrame), pd.DataFrame
+    )