Reset TraceEvaluation[] params on an exception (#1204)

This should address the problem seen via Rubi testing where TraceEvalution was called and we hit a recursion limit.

See Mathics3/Mathics3-Rubi#2 (comment)

---------

Co-authored-by: Aravindh Krishnamoorthy <[email protected]>

mathics/builtin/trace.py

@@ -388,13 +388,16 @@ def eval(self, expr, evaluation: Evaluation, options: dict):
         evaluation.definitions.timing_trace_evaluation = (
             options["System`ShowTimeBySteps"] is SymbolTrue
         )
-        result = expr.evaluate(evaluation)
-        evaluation.definitions.trace_evaluation = curr_trace_evaluation
-        evaluation.definitions.timing_trace_evaluation = curr_time_by_steps
-
-        mathics.eval.tracing.trace_evaluate_on_call = old_evaluation_hook
-
-        return result
+        try:
+            result = expr.evaluate(evaluation)
+            return result
+        except Exception:
+            raise
+        finally:
+            evaluation.definitions.trace_evaluation = curr_trace_evaluation
+            evaluation.definitions.timing_trace_evaluation = curr_time_by_steps
+
+            mathics.eval.tracing.trace_evaluate_on_call = old_evaluation_hook
 
 
 class TraceEvaluationVariable(Builtin):

test/builtin/test_exp_structure.py

@@ -3,9 +3,7 @@
 Unit tests for mathics.builtin.exp_structure
 """
 
-import sys
-import time
-from test.helper import check_evaluation, evaluate
+from test.helper import check_evaluation
 
 import pytest
 

test/builtin/test_trace.py

@@ -0,0 +1,60 @@
+# -*- coding: utf-8 -*-
+"""
+Unit tests for mathics.builtin.trace
+"""
+
+from test.helper import evaluate
+
+import pytest
+
+import mathics.eval.tracing
+from mathics.core.interrupt import AbortInterrupt
+
+trace_evaluation_calls = 0
+
+
+def test_TraceEvaluation():
+    """
+    Basic test of TraceEvaluate[]
+    """
+    old_recursion_limit = evaluate("$RecursionLimit")
+    old_evaluation_hook = mathics.eval.tracing.print_evaluate
+
+    def counting_print_evaluate(expr, evaluation, status: str, orig_expr=None) -> bool:
+        """
+        A replacement for mathics.eval.tracing.print_evaluate() that counts the
+        number of evaluation calls.
+        """
+        global trace_evaluation_calls
+        trace_evaluation_calls += 1
+        return False
+
+    try:
+        # Set a small recursion limit,
+        # Replace TraceEvaluation's print function something that counts evaluation
+        # calls, and then force a RecursionLimit Error.
+        evaluate("$RecursionLimit = 20")
+        assert mathics.eval.tracing.print_evaluate == old_evaluation_hook
+        evaluate("f[x_] := x + f[x-1]; f[0] = 0")
+        global trace_evaluation_calls
+        trace_evaluation_calls = 0
+        mathics.eval.tracing.print_evaluate = counting_print_evaluate
+        evaluate("f[30] // TraceEvaluation")
+
+    except AbortInterrupt:
+        # We should get an AbortInterrupt from exceeding RecursionLimit in evaluating f[30]
+        assert trace_evaluation_calls != 0, "TraceEvaluate[] should have counted steps"
+
+        # Clear evaluation-call count and then check that TraceEvaluation restored
+        # ts print hook. We do this by running another evaluate and checking
+        # that nothing happened.
+        trace_evaluation_calls = 0
+        evaluate("1+2")
+        assert trace_evaluation_calls == 0
+    else:
+        pytest.xfail("We should have raised an AbortInterrupt in evaluation")
+    finally:
+        # Just in case, restore everything back to what it was before running this test.
+        mathics.eval.tracing.print_evaluate = old_evaluation_hook
+        old_recursion_limit = evaluate(f"$RecursionLimit = {old_recursion_limit.value}")
+    assert mathics.eval.tracing.print_evaluate == old_evaluation_hook