Python import 语句的智能重构：基于 AST 实现精细化管理

import ast

def collect_attribute_usage(code):
    """
    解析代码并收集模块属性的使用情况。
    例如，对于 'math.sin'，将记录 'math' 使用了 'sin'。
    """
    tree = ast.parse(code)
    attr_usage = {}  # 存储 {module_name: {attribute1, attribute2, ...}}

    for node in ast.walk(tree):
        if isinstance(node, ast.Attribute):
            # 检查属性访问是否是 'module.attribute' 形式
            # 确保 node.value 是一个 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)
"""

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

class IndividualizeImportNames(ast.NodeTransformer):
    """
    一个 AST 转换器，用于：
    1. 将 'import module' 替换为 'from module import attr1, attr2'。
    2. 将 'module.attr' 替换为 'attr'。
    """
    def __init__(self, attr_usage):
        self.attr_usage = attr_usage

    def visit_Import(self, node):
        """
        处理 'import module1, module2' 形式的导入语句。
        """
        new_imports = []

        # 遍历当前 import 语句中的所有别名（模块名）
        for alias in node.names:
            module_name = alias.name

            if module_name in self.attr_usage:
                # 如果该模块的属性被使用了，则创建 'from module import attr1, attr2'
                imported_attrs = [ast.alias(name=attr) for attr in self.attr_usage[module_name]]
                new_imports.append(ast.ImportFrom(module=module_name, names=imported_attrs, level=0))
            else:
                # 如果该模块的属性没有被使用（但模块本身被导入了），则保留 'import module'
                # 这也处理了原始 'import module1, module2' 中未使用的模块
                new_imports.append(ast.Import(names=[alias]))

        # 返回一个包含新导入语句的列表。如果列表为空，则表示该原始导入语句被完全移除。
        # 如果原始 import 语句中包含多个模块，这里会为每个模块生成一个导入节点。
        return new_imports

    def visit_Attribute(self, node):
        """
        处理 'module.attribute' 形式的属性访问。
        """
        self.generic_visit(node)  # 确保子节点也被访问和转换

        # 如果 node.value 是一个 Name 节点，且其 ID（模块名）在 attr_usage 中
        if isinstance(node.value, ast.Name) and node.value.id in self.attr_usage:
            # 将 'module.attribute' 替换为 'attribute'
            return ast.Name(id=node.attr, ctx=ast.Load())
        return node

def individualize_import_names(code):
    """
    将 Python 源代码中的 'import module' 语句重构为 
    'from module import name1, name2, ...' 形式，
    并相应地修改模块属性的调用。
    """
    # 1. 解析代码并收集属性使用情况
    tree = ast.parse(code)
    attr_usage = {}
    for node in ast.walk(tree):
        if isinstance(node, ast.Attribute):
            if isinstance(node.value, ast.Name): # 确保是 module.attribute 形式
                attr_usage.setdefault(node.value.id, set()).add(node.attr)

    # 2. 应用 AST 转换
    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)
"""

transformed_code = individualize_import_names(original_code)
print(transformed_code)

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