在源代码中将 `import ...` 转换为 `from ... import ...`

Question

我在 Python 中有一个表示某些代码的字符串，如果代码包含像 time.sleep() 这样的方法调用，我想用相应的 from time import sleep 替换 time 的 import 语句。另外，我想用 sleep() 替换方法调用 time.sleep() 。

下面是代码示例：

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

我想将此代码转换为：

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

我尝试过使用正则表达式，但无法达到预期的结果。有人可以帮助我为这项任务提供正确的正则表达式和替换逻辑吗？

这是我尝试过的代码：

re.sub(rf'import\s+{re.escape("time")}\s*\n*(\s*{re.escape("time")}\.(\w+)\(([^)]*)\))', lambda m: f'from {m.group(2)} import {m.group(3)}, {m.group(2)}' if 'time' in m.group(2) else f'from {m.group(2)} import {m.group(3)}', a)

Answer 1

要转换import module为from module import names, ...语句，您应该首先收集这些names，即代码中的属性访问。将代码解析为 AST，找到ast.Attribute节点，并将名称到属性的映射存储在集合的字典中：

import ast

tree = ast.parse(code)
attr_usage = {}
for node in ast.walk(tree):
    if isinstance(node, ast.Attribute):
        attr_usage.setdefault(node.value.id, set()).add(node.attr)

然后，您可以使用子类将节点ast.NodeTransformer替换为相应集合中名称的节点，并将字典中名称为这些模块的节点替换为相应的节点。如果代码中存在应保留在语句中的未使用模块，您可以使用单独的列表推导式从导入名称列表中过滤属性映射中的名称：ImportImportFromAttributeNameimport

class IndividualizeImportNames(ast.NodeTransformer):
    def __init__(self, attr_usage):
        self.attr_usage = attr_usage

    def visit_Import(self, node):
        leftover_imports = []
        if unused_names := [
            alias for alias in node.names
            if alias.name not in self.attr_usage
        ]:
            leftover_imports.append(ast.Import(names=unused_names))
        return leftover_imports + [
            ast.ImportFrom(
                module=alias.name,
                names=[
                    ast.alias(name=attr) 
                    for attr in self.attr_usage[alias.name]
                ]
            )
            for alias in node.names if alias.name in self.attr_usage
        ]

    def visit_Attribute(self, node):
        self.generic_visit(node)
        if node.value.id in self.attr_usage:
            return ast.Name(id=node.attr, ctx=ast.Load())
        return node

IndividualizeImportNames(attr_usage).visit(tree)

将上面的两个步骤合并成一个函数：

def individualize_import_names(code):
    tree = ast.parse(code)
    attr_usage = {}
    for node in ast.walk(tree):
        if isinstance(node, ast.Attribute):
            attr_usage.setdefault(node.value.id, set()).add(node.attr)
    IndividualizeImportNames(attr_usage).visit(tree)
    return ast.unparse(tree)

以便：

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

输出：

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

演示在这里