Updated `return_output` logic

restrictedpython_code_eval.py

@@ -525,7 +525,6 @@ def _unsafe_execute(check_program, result, timeout,
                     allow_underscore_variable_names: bool = False, return_output: bool = False, output_variable: str = "output"):
 
     with create_tempdir():
-
         # These system calls are needed when cleaning up tempdir.
         import os
         import shutil
@@ -537,9 +536,6 @@ def _unsafe_execute(check_program, result, timeout,
         # Disable functionalities that can make destructive changes to the test.
         reliability_guard()
 
-        if return_output and additional_locals is None:
-            additional_locals = {}
-
         # Run program.
         try:
             builtins = {}
@@ -606,12 +602,15 @@ def _unsafe_execute(check_program, result, timeout,
             if '_apply_' not in exec_globals:
                 exec_globals['_apply_'] = _apply  # type: ignore
 
+            if additional_locals is not None:
+                additional_globals.update(additional_locals)
+
             with swallow_io():
                 policy_class = AllowAugmentedAssignAndUnderscoreVariableNamesRestrictingTransformer if allow_underscore_variable_names else AllowAugmentedAssignRestrictingTransformer
 
                 with time_limit(timeout):
                     byte_code = compile_restricted(check_program, filename="<model output>", mode="exec", policy=policy_class)
-                    exec(byte_code, exec_globals, additional_locals)
+                    exec(byte_code, exec_globals, None)
             
             result['result'] = "passed"
 
@@ -623,8 +622,8 @@ def _unsafe_execute(check_program, result, timeout,
             result['result'] = f"failed ({type(e)}): {str(e)}"
             result['exception'] = e
 
-        if return_output and output_variable in additional_locals:  # type: ignore
-            result[output_variable] = additional_locals[output_variable]  # type: ignore
+        if return_output and output_variable in exec_globals:  # type: ignore
+            result[output_variable] = exec_globals[output_variable]  # type: ignore
 
         # Needed for cleaning up.
         shutil.rmtree = rmtree