Python: exec() a code block and eval() the last line

Question

I have a string literal containing one or more lines of (trusted) Python code, and I would like to exec() the block, while capturing the results of the last line. More concretely, I would like a function exec_then_eval that returns the following:

code = """
x = 4
y = 5
x + y
"""

assert exec_then_eval(code) == 9

A couple things I've tried:

1. By splitting off the last line, you can exec the first block, then eval the last line; e.g.

   def exec_then_eval(code):
       first_block = '\n'.join(code.splitlines()[:-1])
       last_line = code.splitlines()[-1]
       globals = {}
       locals = {}
       exec(first_block, globals, locals)
       return eval(last_line, globals, locals)
   

   This works, but will fail if the last statement has multiple lines.

2. If the code itself is modified so the result is stored as a local variable, this variable can then be recovered; e.g.

   code = """
   x = 4
   y = 5
   z = x + y
   """
   
   globals = {}
   locals = {}
   exec(code, globals, locals)
   assert locals['z'] == 9
   

   again, this works, but only if you're able to first parse the code block in a general way and modify it appropriately.

Is there an easy way to write a general exec_and_eval function?

Answer 1

Based on @kalzekdor's suggestion of using the ast module, I came up with this solution which is similar in spirit to @vaultah's solution posted above:

import ast

def exec_then_eval(code):
    block = ast.parse(code, mode='exec')

    # assumes last node is an expression
    last = ast.Expression(block.body.pop().value)

    _globals, _locals = {}, {}
    exec(compile(block, '<string>', mode='exec'), _globals, _locals)
    return eval(compile(last, '<string>', mode='eval'), _globals, _locals)

It also passes @vaultah's nice set of test-cases:

exec_then_eval('''x = 4
y = 5
x + y''')) # 9

exec_then_eval('''x = 4
y = 5;x + y''')) # 9

exec_then_eval('''x = 4
y = 5;(
x + y *

2)''') # 14