Refactor operations count state setting (#631)
This commit is contained in:
parent
3316dd7b11
commit
41a388dac6
|
@ -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)
|
||||||
|
|
|
@ -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
|
||||||
|
|
||||||
|
|
||||||
|
|
Loading…
Reference in New Issue