Python AST实战：动态重构导入语句以优化代码引用

import ast

def collect_attribute_usage(code):
    """
    解析代码，收集每个模块使用的属性。
    返回一个字典，键是模块名，值是该模块使用的属性名集合。
    """
    tree = ast.parse(code)
    attr_usage = {}
    for node in ast.walk(tree):
        # 查找所有 ast.Attribute 节点，例如 'time.sleep'
        if isinstance(node, ast.Attribute):
            # 确保 node.value 是一个 ast.Name 节点，表示直接的模块名
            if isinstance(node.value, ast.Name):
                module_name = node.value.id
                attribute_name = node.attr
                attr_usage.setdefault(module_name, set()).add(attribute_name)
    return attr_usage

# 示例代码
original_code = """
import math, numpy, random
import time
from PIL import Image
a = math.sin(90)
time.sleep(3)
"""

# 收集属性使用情况
usage_map = collect_attribute_usage(original_code)
print("模块属性使用情况:", usage_map)
# 预期输出: 模块属性使用情况: {'math': {'sin'}, 'time': {'sleep'}}

class IndividualizeImportNames(ast.NodeTransformer):
    """
    AST转换器，用于将 'import module' 转换为 'from module import name'
    并简化 'module.name()' 调用为 'name()'。
    """
    def __init__(self, attr_usage):
        self.attr_usage = attr_usage

    def visit_Import(self, node):
        """
        处理 'import module, other_module' 形式的导入语句。
        将其分解为 'from module import name' 和保留未使用的 'import other_module'。
        """
        new_imports = []
        # 遍历当前 import 语句中的所有别名 (例如 'math', 'numpy', 'random')
        for alias in node.names:
            module_name = alias.name
            # 如果该模块有被使用的属性
            if module_name in self.attr_usage:
                # 创建一个新的 'from module import name1, name2...' 语句
                new_imports.append(
                    ast.ImportFrom(
                        module=module_name,
                        names=[
                            ast.alias(name=attr, asname=None) # asname=None表示不使用别名
                            for attr in sorted(list(self.attr_usage[module_name])) # 排序以保持输出稳定
                        ],
                        level=0 # 绝对导入
                    )
                )
            else:
                # 如果模块没有被使用的属性，或者它是一个 'import module as alias'
                # （本方案主要针对 'import module' 形式），则保留原始导入。
                # 注意：本解决方案不直接处理 'import module as alias' 的情况。
                new_imports.append(ast.Import(names=[alias]))

        # 返回一个包含新导入语句的列表。
        # 如果原始 import 语句被完全替换，则返回空列表或新的 ImportFrom 列表。
        # 如果有未使用的模块，则会保留其原始 Import 语句。
        # 重要的是，此方法返回一个列表，以支持一个 Import 节点被替换为多个节点。
        return new_imports

    def visit_Attribute(self, node):
        """
        处理 'module.attribute' 形式的属性访问。
        如果 'module' 的 'attribute' 已被 'from module import attribute' 导入，
        则将其替换为简单的 'attribute'。
        """
        self.generic_visit(node) # 首先遍历子节点，确保内部结构被正确处理

        # 检查 node.value 是否是 ast.Name (即直接的模块名)
        if isinstance(node.value, ast.Name):
            module_name = node.value.id
            attribute_name = node.attr

            # 如果该模块的该属性在我们的使用记录中
            if module_name in self.attr_usage and attribute_name in self.attr_usage[module_name]:
                # 将 'module.attribute' 替换为 'attribute' (一个 ast.Name 节点)
                return ast.Name(id=attribute_name, ctx=ast.Load())
        return node # 否则，返回原始节点

def individualize_import_names(code):
    """
    将Python源代码中的 'import module' 转换为 'from module import name'
    并简化 'module.name()' 调用为 'name()'。
    """
    # 1. 解析代码并收集属性使用情况
    tree = ast.parse(code)
    attr_usage = {}
    for node in ast.walk(tree):
        if isinstance(node, ast.Attribute) and isinstance(node.value, ast.Name):
            attr_usage.setdefault(node.value.id, set()).add(node.attr)

    # 2. 转换AST
    # 注意：ast.NodeTransformer 的 visit 方法返回的节点或节点列表会替换原始节点。
    # 对于 ast.Import 节点，我们可能会返回多个节点（ImportFrom 和/或 Import）。
    # 对于其他节点，通常返回单个节点或原始节点。
    IndividualizeImportNames(attr_usage).visit(tree)

    # 3. 将修改后的AST转换回代码字符串
    return ast.unparse(tree)

# 完整示例
original_code = """
import math, numpy, random
import time
from PIL import Image
a = math.sin(90)
time.sleep(3)
"""

print("--- 原始代码 ---")
print(original_code)

transformed_code = individualize_import_names(original_code)

print("\n--- 转换后的代码 ---")
print(transformed_code)

--- 原始代码 ---

import math, numpy, random
import time
from PIL import Image
a = math.sin(90)
time.sleep(3)


--- 转换后的代码 ---
import numpy, random
from math import sin
from time import sleep
from PIL import Image
a = sin(90)
sleep(3)