diff --git a/pandas-stubs/__init__.pyi b/pandas-stubs/__init__.pyi index bbac8bd73..b1dba7b4b 100644 --- a/pandas-stubs/__init__.pyi +++ b/pandas-stubs/__init__.pyi @@ -6,6 +6,7 @@ from . import ( plotting as plotting, testing as testing, tseries as tseries, + util as util, ) from ._config import ( describe_option as describe_option, diff --git a/pandas-stubs/core/computation/eval.pyi b/pandas-stubs/core/computation/eval.pyi index 388513783..ad4e2ed11 100644 --- a/pandas-stubs/core/computation/eval.pyi +++ b/pandas-stubs/core/computation/eval.pyi @@ -1,12 +1,30 @@ +from typing import ( + Any, + Literal, + Mapping, +) + +from pandas import ( + DataFrame, + Series, +) +from pandas.core.computation.ops import BinOp + +from pandas._typing import ( + Scalar, + npt, +) + def eval( - expr, - parser=..., - engine: str | None = ..., - truediv=..., - local_dict=..., - global_dict=..., - resolvers=..., - level=..., - target=..., - inplace=..., -): ... + expr: str | BinOp, + parser: Literal["pandas", "python"] = ..., + engine: Literal["python", "numexpr"] | None = ..., + # Keyword only due to omitted deprecated argument + *, + local_dict: dict[str, Any] | None = ..., + global_dict: dict[str, Any] | None = ..., + resolvers: list[Mapping] | None = ..., + level: int = ..., + target: object | None = ..., + inplace: bool = ..., +) -> npt.NDArray | Scalar | DataFrame | Series | None: ... diff --git a/pandas-stubs/core/util/hashing.pyi b/pandas-stubs/core/util/hashing.pyi index 538f5b2a3..cedae1308 100644 --- a/pandas-stubs/core/util/hashing.pyi +++ b/pandas-stubs/core/util/hashing.pyi @@ -1,10 +1,22 @@ +import numpy as np +from pandas import ( + DataFrame, + Index, + Series, +) + +from pandas._typing import ( + ArrayLike, + npt, +) + def hash_pandas_object( - obj, + obj: Index | Series | DataFrame, index: bool = ..., encoding: str = ..., hash_key: str | None = ..., categorize: bool = ..., -): ... +) -> Series: ... def hash_array( - vals, encoding: str = ..., hash_key: str = ..., categorize: bool = ... -): ... + vals: ArrayLike, encoding: str = ..., hash_key: str = ..., categorize: bool = ... +) -> npt.NDArray[np.uint64]: ... diff --git a/tests/test_pandas.py b/tests/test_pandas.py index cd42caa34..f44afc595 100644 --- a/tests/test_pandas.py +++ b/tests/test_pandas.py @@ -14,6 +14,8 @@ import pytest from typing_extensions import assert_type +from pandas._typing import Scalar + from tests import check @@ -275,3 +277,42 @@ def test_arrow_dtype() -> None: ), pd.ArrowDtype, ) + + +def test_hashing(): + a = np.array([1, 2, 3]) + check(assert_type(pd.util.hash_array(a), npt.NDArray[np.uint64]), np.ndarray) + check( + assert_type( + pd.util.hash_array(a, encoding="latin1", hash_key="1", categorize=True), + npt.NDArray[np.uint64], + ), + np.ndarray, + ) + + b = pd.Series(a) + c = pd.DataFrame({"a": a, "b": a}) + d = pd.Index(b) + check(assert_type(pd.util.hash_pandas_object(b), pd.Series), pd.Series) + check(assert_type(pd.util.hash_pandas_object(c), pd.Series), pd.Series) + check(assert_type(pd.util.hash_pandas_object(d), pd.Series), pd.Series) + check( + assert_type( + pd.util.hash_pandas_object( + d, index=True, encoding="latin1", hash_key="apple", categorize=True + ), + pd.Series, + ), + pd.Series, + ) + + +def test_eval(): + df = pd.DataFrame({"animal": ["dog", "pig"], "age": [10, 20]}) + check( + assert_type( + pd.eval("double_age = df.age * 2", target=df), + Union[npt.NDArray, Scalar, pd.DataFrame, pd.Series, None], + ), + pd.DataFrame, + )