Experimenting with how I return results

restrictedpython_code_eval.py

@@ -453,8 +453,8 @@ def _check_correctness(check_program, timeout, task_id, completion_id,
 
     out_dict = dict(
         task_id=task_id,
-        passed=result[-1] == "passed",
-        result=result[0],
+        passed=result[0] == "passed",
+        result=result,
         completion_id=completion_id,
     )
 
@@ -612,9 +612,6 @@ def _unsafe_execute(check_program, result, timeout,
                 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)
-
-            if return_output:
-                result.append(additional_locals[output_variable])
             
             result.append("passed")
 
@@ -626,6 +623,9 @@ def _unsafe_execute(check_program, result, timeout,
             result.append(f"failed ({type(e)}): {str(e)}")
             result.append(e)
 
+        if return_output and output_variable in additional_locals:  # type: ignore
+            result.append(additional_locals[output_variable])  # type: ignore
+
         # Needed for cleaning up.
         shutil.rmtree = rmtree
         os.rmdir = rmdir