Complete signature of value_counts (#995)

For Index, Series, DataFrame

Author: refack

pandas-stubs/core/base.pyi

@@ -7,10 +7,14 @@ from typing import (
     Generic,
     Literal,
     final,
+    overload,
 )
 
 import numpy as np
-from pandas import Index
+from pandas import (
+    Index,
+    Series,
+)
 from pandas.core.arraylike import OpsMixin
 from pandas.core.arrays import ExtensionArray
 from pandas.core.arrays.categorical import Categorical
@@ -76,14 +80,24 @@ class IndexOpsMixin(OpsMixin, Generic[S1]):
     def __iter__(self) -> Iterator[S1]: ...
     @property
     def hasnans(self) -> bool: ...
+    @overload
+    def value_counts(
+        self,
+        normalize: Literal[False] = ...,
+        sort: bool = ...,
+        ascending: bool = ...,
+        bins=...,
+        dropna: bool = ...,
+    ) -> Series[int]: ...
+    @overload
     def value_counts(
         self,
-        normalize: bool = ...,
+        normalize: Literal[True],
         sort: bool = ...,
         ascending: bool = ...,
         bins=...,
         dropna: bool = ...,
-    ): ...
+    ) -> Series[float]: ...
     def nunique(self, dropna: bool = ...) -> int: ...
     @property
     def is_unique(self) -> bool: ...

pandas-stubs/core/frame.pyi

@@ -993,14 +993,24 @@ class DataFrame(NDFrame, OpsMixin):
         ignore_index: _bool = ...,
         key: Callable | None = ...,
     ) -> DataFrame | None: ...
+    @overload
     def value_counts(
         self,
         subset: Sequence[Hashable] | None = ...,
-        normalize: _bool = ...,
+        normalize: Literal[False] = ...,
         sort: _bool = ...,
         ascending: _bool = ...,
         dropna: _bool = ...,
     ) -> Series[int]: ...
+    @overload
+    def value_counts(
+        self,
+        normalize: Literal[True],
+        subset: Sequence[Hashable] | None = ...,
+        sort: _bool = ...,
+        ascending: _bool = ...,
+        dropna: _bool = ...,
+    ) -> Series[float]: ...
     def nlargest(
         self,
         n: int,

pandas-stubs/core/series.pyi

@@ -1477,14 +1477,24 @@ class Series(IndexOpsMixin[S1], NDFrame):
     ) -> Series[S1]: ...
     def first_valid_index(self) -> Scalar: ...
     def last_valid_index(self) -> Scalar: ...
+    @overload
     def value_counts(
         self,
-        normalize: _bool = ...,
+        normalize: Literal[False] = ...,
         sort: _bool = ...,
         ascending: _bool = ...,
         bins: int | None = ...,
         dropna: _bool = ...,
     ) -> Series[int]: ...
+    @overload
+    def value_counts(
+        self,
+        normalize: Literal[True],
+        sort: _bool = ...,
+        ascending: _bool = ...,
+        bins: int | None = ...,
+        dropna: _bool = ...,
+    ) -> Series[float]: ...
     def transpose(self, *args, **kwargs) -> Series[S1]: ...
     @property
     def T(self) -> Self: ...

tests/test_frame.py

@@ -610,7 +610,12 @@ def test_types_idxmax() -> None:
 # This was added in 1.1.0 https://pandas.pydata.org/docs/whatsnew/v1.1.0.html
 def test_types_value_counts() -> None:
     df = pd.DataFrame(data={"col1": [1, 2], "col2": [1, 4]})
-    df.value_counts()
+    check(assert_type(df.value_counts(), "pd.Series[int]"), pd.Series, np.integer)
+    check(
+        assert_type(df.value_counts(normalize=True), "pd.Series[float]"),
+        pd.Series,
+        float,
+    )
 
 
 def test_types_unique() -> None:

tests/test_indexes.py

@@ -1135,3 +1135,13 @@ def test_get_loc() -> None:
         np.ndarray,
         np.bool_,
     )
+
+
+def test_value_counts() -> None:
+    nmi = pd.Index(list("abcb"))
+    check(assert_type(nmi.value_counts(), "pd.Series[int]"), pd.Series, np.integer)
+    check(
+        assert_type(nmi.value_counts(normalize=True), "pd.Series[float]"),
+        pd.Series,
+        float,
+    )

tests/test_series.py

@@ -557,6 +557,11 @@ def test_types_idxmax() -> None:
 def test_types_value_counts() -> None:
     s = pd.Series(["a", "b"])
     check(assert_type(s.value_counts(), "pd.Series[int]"), pd.Series, np.integer)
+    check(
+        assert_type(s.value_counts(normalize=True), "pd.Series[float]"),
+        pd.Series,
+        float,
+    )
 
 
 def test_types_unique() -> None: