如何用ast将Python的if-else块封装为函数并保留原格式与注释？

核心问题根源 #

ast模块生成的是抽象语法树（AST），仅存储代码的逻辑结构，不包含注释、缩进、换行等格式信息，因此用它做转换必然丢失原始格式。要实现需求，必须使用能保留完整语法细节的**具象语法树（CST）**工具。

方案一：基于Parso的工业级实现 #

Parso是专门为Python设计的语法分析库，生成的CST会完整保留代码的注释、空白字符和缩进格式，是这类场景的首选工具。

1. 安装依赖

pip install parso

2. 转换代码实现

import parso

def encapsulate_if_block(code, func_name="if_encapsulated"):
    # 解析代码为保留格式的具象语法树
    module = parso.parse(code)
    
    # 遍历并定位第一个if语句（如需批量处理可移除break）
    for stmt_idx, stmt in enumerate(module.children):
        if isinstance(stmt, parso.python.tree.IfStmt):
            # 克隆原if-else节点，保留所有格式和注释
            cloned_if = stmt.clone()
            
            # 构建函数定义节点
            func_def = parso.parse(f"def {func_name}():\n").children[0]
            # 将克隆的if-else加入函数体
            func_def.body.append(cloned_if)
            
            # 替换原if语句为函数定义
            module.children[stmt_idx] = func_def
            # 在函数定义后插入函数调用
            call_stmt = parso.parse(f"{func_name}()\n").children[0]
            module.children.insert(stmt_idx + 1, call_stmt)
            break
    
    # 生成保留原始格式的代码
    return module.get_code()

3. 使用示例
   原代码：

# 初始化变量
x = 10
# 判断x的范围
if x > 5:
    print("x大于5")
    # 内部注释：执行后续逻辑
    x += 1
else:
    print("x小于等于5")
# 后续业务代码
y = x * 2

转换后输出：

# 初始化变量
x = 10
# 判断x的范围
def if_encapsulated():
    if x > 5:
        print("x大于5")
        # 内部注释：执行后续逻辑
        x += 1
    else:
        print("x小于等于5")
if_encapsulated()
# 后续业务代码
y = x * 2

方案二：简单场景的文本正则处理 #

如果不想引入第三方库，且代码格式规范（如统一用4空格缩进），可使用正则匹配实现快速转换，仅适合无嵌套的简单if-else结构：

import re

def encapsulate_if_text(code, func_name="if_encapsulated"):
    # 匹配规范缩进的if-else块（可根据实际缩进调整正则）
    pattern = re.compile(
        r'(^if .+?:\n)((\s{4}.+\n)+)(^else:\n)?((\s{4}.+\n)+)?',
        re.MULTILINE
    )
    
    def replace_handler(match):
        if_head = match.group(1)
        if_body = match.group(2)
        else_head = match.group(3) or ""
        else_body = match.group(4) or ""
        
        # 拼接函数代码，保留原格式
        func_code = f"def {func_name}():\n    {if_head.strip()}\n{if_body}{else_head}{else_body}"
        return f"{func_code}\n{func_name}()"
    
    return pattern.sub(replace_handler, code)

关键注意事项 #

• 绝对避免使用ast/astor/unparse：这类工具仅处理语法逻辑，完全不保留格式和注释，无法满足需求。
• Parso支持复杂场景：嵌套if-else、不规则缩进、跨行注释等情况都能正确处理，是生产环境的可靠选择。

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