diff --git a/src/smolagents/local_python_executor.py b/src/smolagents/local_python_executor.py index 8e150cd5c..5df3802fa 100644 --- a/src/smolagents/local_python_executor.py +++ b/src/smolagents/local_python_executor.py @@ -1627,6 +1627,20 @@ def evaluate_python_code( state["_print_outputs"] = PrintContainer() state["_operations_count"] = {"counter": 0} + # Bind the `print` tool to this run's print outputs so that every invocation + # is captured — including when `print` is passed around as a callback (e.g. + # `map(print, values)`), which bypasses the direct-call interception in + # `evaluate_call`. Only swap it in when the caller already allows `print`: + # with a custom static_tools lacking it, `print` stays forbidden, as before. + print_outputs = state["_print_outputs"] + + def capturing_print(*args, **kwargs): + print_outputs.append(" ".join(map(str, args)) + "\n") + return None + + if "print" in static_tools: + static_tools["print"] = capturing_print + if "final_answer" in static_tools: previous_final_answer = static_tools["final_answer"] diff --git a/tests/test_local_python_executor.py b/tests/test_local_python_executor.py index aede6cd2d..645ddf0c5 100644 --- a/tests/test_local_python_executor.py +++ b/tests/test_local_python_executor.py @@ -830,6 +830,26 @@ def function(): evaluate_python_code(code, {"print": print, "range": range}, state=state) assert state["_print_outputs"].value == "1\n2\n2\n2\n2\n2\n2\n2\n2\n2\n2\n" + def test_print_as_callback_is_captured(self): + # print passed as a callback (e.g. to map) must capture like a direct call + code = "list(map(print, ['red', 'blue']))" + state = {} + result, _ = evaluate_python_code(code, BASE_PYTHON_TOOLS, state=state) + assert result == [None, None] + assert state["_print_outputs"].value == "red\nblue\n" + + # aliased print must capture as well + code = "f = print\nf('aliased')" + state = {} + evaluate_python_code(code, BASE_PYTHON_TOOLS, state=state) + assert state["_print_outputs"].value == "aliased\n" + + def test_print_stays_forbidden_without_print_in_static_tools(self): + # callers passing a custom static_tools without print keep the + # allowlist behavior: print is forbidden, as before the fix + with pytest.raises(InterpreterError): + evaluate_python_code("print('hi')", {}, state={}) + def test_tuple_target_in_iterator(self): code = "for a, b in [('Ralf Weikert', 'Austria'), ('Samuel Seungwon Lee', 'South Korea')]:res = a.split()[0]" result, _ = evaluate_python_code(code, BASE_PYTHON_TOOLS, state={})