Move py eval fn

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "egglog-python"
-version = "3.1.0"
+version = "4.0.0"
 edition = "2021"
 
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html

docs/changelog.md

@@ -2,9 +2,12 @@
 
 _This project uses semantic versioning_
 
-## UNRELEASED
+## Unreleased
+
+## 4.0.0 (2023-11-24)
 
 - Fix `as_egglog_string` proprety.
+- Move `EGraph.eval_fn` to `py_eval_fn` since it doesn't need the `EGraph` anymore.
 
 ## 3.1.0 (2023-11-21)
 

docs/reference/python-integration.md

@@ -105,7 +105,7 @@ assert egraph.eval(evalled) == 3
 
 ### Simpler Eval
 
-Instead of using the above low level primitive for evaluating, there is a higher level wrapper function, `egraph.eval_fn`.
+Instead of using the above low level primitive for evaluating, there is a higher level wrapper function, `eval_fn`.
 
 It takes in a Python function and converts it to a function of PyObjects, by using `py_eval` under the hood.
 
@@ -115,7 +115,7 @@ The above code code be re-written like this:
 def my_add(a, b):
     return a + b
 
-evalled = egraph.eval_fn(lambda a: my_add(a, 2))(1)
+evalled = eval_fn(lambda a: my_add(a, 2))(1)
 assert egraph.eval(evalled) == 3
 ```
 

python/egglog/builtins.py

@@ -5,11 +5,14 @@
 
 from __future__ import annotations
 
-from typing import Generic, TypeVar, Union
+from typing import TYPE_CHECKING, Generic, Protocol, TypeVar, Union
 
 from .egraph import BUILTINS, Expr, Unit
 from .runtime import converter
 
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
 __all__ = [
     "BUILTINS",
     "i64",
@@ -28,6 +31,7 @@
     "PyObject",
     "py_eval",
     "py_exec",
+    "py_eval_fn",
 ]
 
 
@@ -552,6 +556,32 @@ def py_eval(code: StringLike, globals: object = PyObject.dict(), locals: object
     ...
 
 
+class PyObjectFunction(Protocol):
+    def __call__(self, *__args: PyObject) -> PyObject:
+        ...
+
+
+def py_eval_fn(fn: Callable) -> PyObjectFunction:
+    """
+    Takes a python callable and maps it to a callable which takes and returns PyObjects.
+
+    It translates it to a call which uses `py_eval` to call the function, passing in the
+    args as locals, and using the globals from function.
+    """
+
+    def inner(*__args: PyObject, __fn: Callable = fn) -> PyObject:
+        new_kvs: list[object] = []
+        eval_str = "__fn("
+        for i, arg in enumerate(__args):
+            new_kvs.append(f"__arg_{i}")
+            new_kvs.append(arg)
+            eval_str += f"__arg_{i}, "
+        eval_str += ")"
+        return py_eval(eval_str, PyObject({"__fn": __fn}).dict_update(*new_kvs), __fn.__globals__)
+
+    return inner
+
+
 @BUILTINS.function(egg_fn="py-exec")
 def py_exec(code: StringLike, globals: object = PyObject.dict(), locals: object = PyObject.dict()) -> PyObject:
     """

python/egglog/egraph.py

@@ -17,7 +17,6 @@
     Generic,
     Literal,
     NoReturn,
-    Protocol,
     TypedDict,
     TypeVar,
     Union,
@@ -101,11 +100,6 @@
 ALWAYS_MUTATES_SELF = {"__setitem__", "__delitem__"}
 
 
-class PyObjectFunction(Protocol):
-    def __call__(self, *__args: PyObject) -> PyObject:
-        ...
-
-
 @dataclass
 class _BaseModule(ABC):
     """
@@ -1080,27 +1074,6 @@ def eval(self, expr: Expr) -> object:
                 return self._egraph.eval_py_object(egg_expr)
         raise NotImplementedError(f"Eval not implemented for {typed_expr.tp.name}")
 
-    def eval_fn(self, fn: Callable) -> PyObjectFunction:
-        """
-        Takes a python callable and maps it to a callable which takes and returns PyObjects.
-
-        It translates it to a call which uses `py_eval` to call the function, passing in the
-        args as locals, and using the globals from function.
-        """
-        from .builtins import PyObject, py_eval
-
-        def inner(*__args: PyObject, __fn: Callable = fn) -> PyObject:
-            new_kvs: list[object] = []
-            eval_str = "__fn("
-            for i, arg in enumerate(__args):
-                new_kvs.append(f"__arg_{i}")
-                new_kvs.append(arg)
-                eval_str += f"__arg_{i}, "
-            eval_str += ")"
-            return py_eval(eval_str, PyObject({"__fn": __fn}).dict_update(*new_kvs), __fn.__globals__)
-
-        return inner
-
     def saturate(
         self, *, max: int = 1000, performance: bool = False, **kwargs: Unpack[GraphvizKwargs]
     ) -> ipywidgets.Widget:

python/tests/test_high_level.py

@@ -536,3 +536,28 @@ def test_no_egglog_string():
     egraph.register((i64(1)))
     with pytest.raises(ValueError):
         egraph.as_egglog_string
+
+
+
+def test_eval_fn():
+    egraph = EGraph()
+
+    assert egraph.eval(py_eval_fn(lambda x: (x,))(PyObject.from_int(1))) == (1,)
+
+
+def _global_make_tuple(x):
+    return (x,)
+
+def test_eval_fn_globals():
+    egraph = EGraph()
+
+    assert egraph.eval(py_eval_fn(lambda x: _global_make_tuple(x))(PyObject.from_int(1))) == (1,)
+
+def test_eval_fn_locals():
+    egraph = EGraph()
+
+
+    def _locals_make_tuple(x):
+        return (x,)
+
+    assert egraph.eval(py_eval_fn(lambda x: _locals_make_tuple(x))(PyObject.from_int(1))) == (1,)