如何自动折叠Python函数中的中间变量以优化内存占用评估？

要解决函数因中间变量赋值导致内存差异的问题，核心是实现表达式内联——将仅使用一次的中间变量定义直接替换到其被调用的位置，同时移除原赋值语句。以下是几种可行的实现方式：

1. 基于AST抽象语法树的通用实现 #

Python的AST模块可以解析代码结构，通过修改语法树再生成代码，这是最灵活通用的方法，适合处理大量函数：

实现代码 #

import ast
import astunparse

def inline_single_use_variables(func_code):
    # 解析代码为抽象语法树
    tree = ast.parse(func_code)
    
    # 收集可内联的变量：单变量赋值且仅被引用一次（赋值本身+一次调用）
    assign_map = {}
    for node in ast.walk(tree):
        if isinstance(node, ast.Assign) and len(node.targets) == 1:
            target_node = node.targets[0]
            if isinstance(target_node, ast.Name):
                var_name = target_node.id
                # 统计该变量的总引用次数
                ref_count = sum(1 for n in ast.walk(tree) if isinstance(n, ast.Name) and n.id == var_name)
                if ref_count == 2:
                    assign_map[var_name] = node.value
    
    # 定义AST转换器，替换变量引用为对应表达式
    class Inliner(ast.NodeTransformer):
        def visit_Name(self, node):
            return assign_map.get(node.id, node)
    
    # 应用转换并移除原赋值语句
    transformed_tree = Inliner().visit(tree)
    cleaned_body = [
        node for node in transformed_tree.body
        if not (isinstance(node, ast.Assign) and len(node.targets) == 1 
                and isinstance(node.targets[0], ast.Name) 
                and node.targets[0].id in assign_map)
    ]
    transformed_tree.body = cleaned_body
    
    # 修复语法树位置信息，生成代码
    ast.fix_missing_locations(transformed_tree)
    return astunparse.unparse(transformed_tree).strip()

# 测试示例函数
se1_code = """
def se1(vals1, vals2):
    diff = [(v1 - v2) for v1, v2 in zip(vals1, vals2)]
    diff_sq = [d**2 for d in diff]
    return(sum(diff_sq))
"""

print(inline_single_use_variables(se1_code))

输出效果 #

生成的代码会自动内联中间变量，结果类似：

def se1(vals1, vals2):
    return(sum([(d**2) for d in [(v1 - v2) for v1, v2 in zip(vals1, vals2)]]))

你可以进一步优化嵌套推导式的合并，但核心内存优化的目标已经达成。

注意事项 #

• 仅处理单变量赋值且仅被使用一次的场景，避免多引用变量替换后重复计算的问题
• 不支持复杂赋值（如多变量赋值、对象属性赋值等），可根据需求扩展AST处理逻辑

2. 使用第三方库简化AST操作 #

如果不想手动处理AST的细节，可以用redbaron这类专门处理Python代码的库，它能更友好地修改代码结构并保留格式：

from redbaron import RedBaron

def inline_with_redbaron(func_code):
    red_tree = RedBaron(func_code)
    assign_map = {}
    
    # 收集可内联的变量
    for assign_node in red_tree.find_all("assignment"):
        if len(assign_node.target) == 1 and isinstance(assign_node.target[0], str):
            var_name = assign_node.target[0]
            ref_nodes = red_tree.find_all("name", value=var_name)
            if len(ref_nodes) == 2:
                assign_map[var_name] = assign_node.value
    
    # 替换变量引用
    for var_name, expr in assign_map.items():
        for ref_node in red_tree.find_all("name", value=var_name):
            if not isinstance(ref_node.parent, RedBaron.assignment_node_class):
                ref_node.replace(expr.copy())
    
    # 删除原赋值语句
    for assign_node in red_tree.find_all("assignment"):
        if len(assign_node.target) == 1 and assign_node.target[0] in assign_map:
            assign_node.delete()
    
    return red_tree.dumps()

3. 正则表达式快速替换（仅适用于简单场景） #

如果你的函数都是类似示例中的链式赋值结构，可以用正则做快速替换，但通用性极差，仅适合小批量固定结构的函数：

import re

def inline_simple_chain(func_code):
    pattern = r'(\w+) = (.+)\n\s+(\w+) = (.+?)\1(.+)\n\s+return\((.+?)\3(.+)\)'
    replacement = r'return(\6(\4\2\5)\7)'
    return re.sub(pattern, replacement, func_code, flags=re.DOTALL)

验证优化效果 #

用你提供的tracemalloc代码测试重构后的函数，其内存峰值会与se2一致，说明中间变量的冗余内存占用已被消除。

内容的提问来源于stack exchange，提问作者Mathhead