Refactor operations count state setting (#631)

This commit is contained in:
Albert Villanova del Moral 2025-02-13 11:22:02 +01:00 committed by GitHub
parent 3316dd7b11
commit 41a388dac6
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 2 additions and 4 deletions

View File

@ -1184,7 +1184,7 @@ def evaluate_ast(
The list of modules that can be imported by the code. By default, only a few safe modules are allowed. The list of modules that can be imported by the code. By default, only a few safe modules are allowed.
If it contains "*", it will authorize any import. Use this at your own risk! If it contains "*", it will authorize any import. Use this at your own risk!
""" """
if state["_operations_count"] >= MAX_OPERATIONS: if state.setdefault("_operations_count", 0) >= MAX_OPERATIONS:
raise InterpreterError( raise InterpreterError(
f"Reached the max number of operations of {MAX_OPERATIONS}. Maybe there is an infinite loop somewhere in the code, or you're just asking too many calculations." f"Reached the max number of operations of {MAX_OPERATIONS}. Maybe there is an infinite loop somewhere in the code, or you're just asking too many calculations."
) )
@ -1363,7 +1363,6 @@ def evaluate_python_code(
custom_tools = custom_tools if custom_tools is not None else {} custom_tools = custom_tools if custom_tools is not None else {}
result = None result = None
state["_print_outputs"] = PrintContainer() state["_print_outputs"] = PrintContainer()
state["_operations_count"] = 0
def final_answer(value): def final_answer(value):
raise FinalAnswerException(value) raise FinalAnswerException(value)

View File

@ -1149,7 +1149,6 @@ def test_evaluate_python_code_with_evaluate_delete(code, expected_error_message)
], ],
) )
def test_evaluate_delete(code, state, expectation): def test_evaluate_delete(code, state, expectation):
state["_operations_count"] = 0
delete_node = ast.parse(code).body[0] delete_node = ast.parse(code).body[0]
if isinstance(expectation, Exception): if isinstance(expectation, Exception):
with pytest.raises(type(expectation)) as exception_info: with pytest.raises(type(expectation)) as exception_info:
@ -1157,7 +1156,7 @@ def test_evaluate_delete(code, state, expectation):
assert str(expectation) in str(exception_info.value) assert str(expectation) in str(exception_info.value)
else: else:
evaluate_delete(delete_node, state, {}, {}, []) evaluate_delete(delete_node, state, {}, {}, [])
del state["_operations_count"] _ = state.pop("_operations_count", None)
assert state == expectation assert state == expectation