Add ability to override hash

Author: saulshanabrook

python/egglog/egraph.py

@@ -110,7 +110,6 @@
     "__weakref__",
     "__orig_bases__",
     "__annotations__",
-    "__hash__",
     "__qualname__",
     "__firstlineno__",
     "__static_attributes__",

python/egglog/exp/array_api.py

@@ -154,6 +154,16 @@ def __le__(self, other: IntLike) -> Boolean: ...
     def __eq__(self, other: IntLike) -> Boolean:  # type: ignore[override]
         ...
 
+    # add a hash so that this test can pass
+    # https://github.com/scikit-learn/scikit-learn/blob/6fd23fca53845b32b249f2b36051c081b65e2fab/sklearn/utils/validation.py#L486-L487
+    @method(preserve=True)
+    def __hash__(self) -> int:
+        egraph = _get_current_egraph()
+        egraph.register(self)
+        egraph.run(array_api_schedule)
+        simplified = egraph.extract(self)
+        return hash(cast("RuntimeExpr", simplified).__egg_typed_expr__)
+
     # TODO: Fix this?
     # Make != always return a Bool, so that numpy.unique works on a tuple of ints
     # In _unique1d

python/egglog/runtime.py

@@ -567,6 +567,8 @@ def __setstate__(self, d: tuple[Declarations, TypedExprDecl]) -> None:
         self.__egg_typed_expr_thunk__ = Thunk.value(d[1])
 
     def __hash__(self) -> int:
+        if (method := _get_expr_method(self, "__hash__")) is not None:
+            return cast("int", cast("Any", method()))
         return hash(self.__egg_typed_expr__)
 
     # Implement this directly to special case behavior where it transforms to an egraph equality, if it is not a

python/tests/test_high_level.py

@@ -947,6 +947,17 @@ def test_type_param_sub():
     assert str(V[Unit]) == str(Vec[Unit] | int)  # type: ignore[misc]
 
 
+def test_override_hash(self):
+    class A(Expr):
+        def __init__(self) -> None: ...
+
+        @method(preserve=True)
+        def __hash__(self) -> int:
+            return 42
+
+    assert hash(A()) == 42
+
+
 EXAMPLE_FILES = list((pathlib.Path(__file__).parent / "../egglog/examples").glob("*.py"))