如何使用LibCST在.py文件指定位置插入导入语句？

核心思路 #

借助LibCST的CSTTransformer遍历语法树，精准定位到原导入块结束后、第一个非导入语句之前的位置，插入新导入语句，全程保留原代码的注释、空白符等格式细节。

完整代码实现 #

import libcst as cst

class ImportInjector(cst.CSTTransformer):
    def __init__(self, new_imports):
        self.new_imports = new_imports
        self.import_block_processed = False

    def visit_Module(self, node):
        self.import_block_processed = False
        return super().visit_Module(node)

    def leave_Module(self, original_node, updated_node):
        # 处理全文件都是导入语句的情况
        if not self.import_block_processed:
            new_import_lines = [cst.parse_statement(imp) for imp in self.new_imports]
            updated_body = list(updated_node.body) + [cst.EmptyLine()] + new_import_lines
            return updated_node.with_changes(body=updated_body)
        return updated_node

    def leave_Import(self, original_node, updated_node):
        self.import_block_processed = False
        return updated_node

    def leave_ImportFrom(self, original_node, updated_node):
        self.import_block_processed = False
        return updated_node

    def visit_Statement(self, node):
        # 遇到第一个非导入语句时插入新导入
        if not self.import_block_processed:
            self.import_block_processed = True
            new_import_lines = [cst.parse_statement(imp) for imp in self.new_imports]
            return cst.FlattenSentinel([
                cst.EmptyLine(),
                *new_import_lines,
                cst.EmptyLine(),
                node
            ])
        return node

def inject_imports(source_code, new_imports):
    tree = cst.parse_module(source_code)
    transformer = ImportInjector(new_imports)
    modified_tree = tree.visit(transformer)
    return modified_tree.code

# 示例调用
if __name__ == "__main__":
    source = """import car
import horse
import cheese

x = 1"""
    new_imports = ["import new", "import packages"]
    print(inject_imports(source, new_imports))

代码说明 #

• ImportInjector类通过节点遍历实现位置定位：
  • 遍历Module节点时重置状态，保证多模块处理的正确性。
  • 识别Import/ImportFrom节点，标记当前处于导入块区间。
  • 遇到第一个非导入的Statement节点时，在其前方插入空行和新导入语句。
  • 特殊处理全文件都是导入的场景，在文件末尾追加新导入。
• inject_imports函数封装了解析语法树、应用转换、生成目标代码的完整流程。

注意事项 #

• 先安装依赖：pip install libcst
• 代码会完整保留原代码的注释、空行、缩进等格式
• 支持批量插入任意数量的导入语句，只需传入对应的字符串列表即可

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