Save snapshots in testing even if doesn't finish

Author: saulshanabrook

python/egglog/exp/array_api_jit.py

@@ -14,16 +14,22 @@
 X = TypeVar("X", bound=Callable)
 
 
-def jit(fn: X) -> X:
+def jit(
+    fn: X,
+    *,
+    handle_expr: Callable[[NDArray], None] | None = None,
+    handle_optimized_expr: Callable[[NDArray], None] | None = None,
+) -> X:
     """
     Jit compiles a function
     """
     egraph, res, res_optimized, program = function_to_program(fn, save_egglog_string=False)
+    if handle_expr:
+        handle_expr(res)
+    if handle_optimized_expr:
+        handle_optimized_expr(res_optimized)
     fn_program = EvalProgram(program, {"np": np})
-    fn = cast("X", try_evaling(egraph, array_api_program_gen_schedule, fn_program, fn_program.as_py_object))
-    fn.initial_expr = res  # type: ignore[attr-defined]
-    fn.expr = res_optimized  # type: ignore[attr-defined]
-    return fn
+    return cast("X", try_evaling(egraph, array_api_program_gen_schedule, fn_program, fn_program.as_py_object))
 
 
 def function_to_program(fn: Callable, save_egglog_string: bool) -> tuple[EGraph, NDArray, NDArray, Program]:

python/tests/test_array_api.py

@@ -1,6 +1,7 @@
 # mypy: disable-error-code="empty-body"
 import inspect
 from collections.abc import Callable
+from functools import partial
 from itertools import product
 from pathlib import Path
 from types import FunctionType
@@ -352,9 +353,12 @@ def lda(X: NDArray, y: NDArray):
     ],
 )
 def test_jit(program, snapshot_py, benchmark):
-    jitted = benchmark(jit, program)
-    assert str(jitted.initial_expr) == snapshot_py(name="initial_expr")
-    assert str(jitted.expr) == snapshot_py(name="expr")
+    def save_expr(name, expr):
+        assert str(expr) == snapshot_py(name=name)
+
+    jitted = benchmark(
+        jit, program, handle_expr=partial(save_expr, "initial_expr"), handle_optimized_expr=partial(save_expr, "expr")
+    )
     assert inspect.getsource(jitted) == snapshot_py(name="code")