diff --git a/pandas-stubs/core/base.pyi b/pandas-stubs/core/base.pyi index 5b30bd13..744228c7 100644 --- a/pandas-stubs/core/base.pyi +++ b/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: ... diff --git a/pandas-stubs/core/frame.pyi b/pandas-stubs/core/frame.pyi index 51150c91..21974c42 100644 --- a/pandas-stubs/core/frame.pyi +++ b/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, diff --git a/pandas-stubs/core/series.pyi b/pandas-stubs/core/series.pyi index e12fa2cc..24bac1a1 100644 --- a/pandas-stubs/core/series.pyi +++ b/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: ... diff --git a/tests/test_frame.py b/tests/test_frame.py index e9d327aa..3e09e7b3 100644 --- a/tests/test_frame.py +++ b/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: diff --git a/tests/test_indexes.py b/tests/test_indexes.py index 6777736d..d151ba48 100644 --- a/tests/test_indexes.py +++ b/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, + ) diff --git a/tests/test_series.py b/tests/test_series.py index 2b18960e..5a211096 100644 --- a/tests/test_series.py +++ b/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: