将嵌套调用分成几行

Question

我在Python中有一个这样的函数（大写字母可以代表常量，函数，任何东西，但不能代表函数调用）：

def f(x):
    a = foo1(A, B, foo3(E, foo2(A, B)))
    b = foo3(a, E)
    return b

我想将其分解为这样的“原子”操作：

def f(x):
    tmp1 = foo2(A, B)
    tmp2 = foo3(E, tmp1)
    a = foo1(A, B, tmp2)
    b = foo3(a, E)
    return b

换句话说，每行一次函数调用和一次赋值。

有没有办法可以在Python中实现这个源代码转换？接受此类函数的字符串表示形式并返回转换后的版本的程序。我知道我需要使用 AST 表示，但我真的不知道如何继续。

Answer 1

我认为这正是您想要的：

s = """def f(x):
    a = foo1(A, B, foo3(E, foo2(A, B)))
    b = foo3(a, E)
    return b"""

import re

result = []
reg = re.compile(r'^(\s+).*?(\w+\([^\(\)]*\))')
count = 1
for l in s.splitlines():
    r = reg.search(l)
    if not r:
        result.append(l)
        continue
    while r:
        indent = r[1]
        var = f'tmp_{count}'
        result.append(f'{indent}{var} = {r[2]}')
        count += 1
        l = l[:r.start(2)] + var + l[r.end(2):]
        if re.search(r'^\s*\w+ = \w+$', l):
            result.append(l)
            break
        r = reg.search(l)

print('\n'.join(result))

它会显示：

def f(x):
    tmp_1 = foo2(A, B)
    tmp_2 = foo3(E, tmp_1)
    tmp_3 = foo1(A, B, tmp_2)
    a = tmp_3
    tmp_4 = foo3(a, E)
    b = tmp_4
    return b