ENH: Improve hashing and eval (#347)

* ENH: Improve hash functions

* ENH: Improve eval

Author: bashtage

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,

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: ...

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]: ...

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,
+    )