如何将Python代码的import语句转为from...import并替换模块调用

要实现你需要的代码转换，不能用单一正则一步完成，得分三步处理：先收集代码中用到的模块方法，再替换import语句，最后替换方法调用。以下是具体实现：

步骤1：收集模块与对应调用的方法 #

用正则匹配所有[模块名].[方法名]()的调用，提取每个模块下的方法名并去重：

import re

def transform_code(code):
    # 定义需要处理的模块
    target_modules = {'time', 'math'}
    module_methods = {mod: set() for mod in target_modules}
    
    # 匹配模块方法调用：比如 time.sleep(...)
    pattern = re.compile(rf'({"|".join(target_modules)})\.(\w+)\(')
    for match in pattern.finditer(code):
        mod = match.group(1)
        method = match.group(2)
        module_methods[mod].add(method)
    
    # 步骤2：替换import语句
    for mod, methods in module_methods.items():
        if not methods:
            continue
        # 匹配 import 模块名 的语句
        import_pattern = re.compile(rf'^import\s+{re.escape(mod)}\s*$', re.MULTILINE)
        # 替换为 from 模块名 import 方法1, 方法2...
        code = import_pattern.sub(f'from {mod} import {", ".join(methods)}', code)
    
    # 步骤3：替换模块方法调用为直接方法名
    for mod, methods in module_methods.items():
        if not methods:
            continue
        # 匹配 模块名.方法名(
        call_pattern = re.compile(rf'{re.escape(mod)}\.({"|".join(methods)})\(')
        code = call_pattern.sub(r'\1(', code)
    
    return code

示例测试 #

输入代码：

import time
import math

def demo():
    time.sleep(3)
    print(math.sin(90))
    time.sleep(1)
    print(math.cos(45))

转换后输出：

from time import sleep
from math import sin, cos

def demo():
    sleep(3)
    print(sin(90))
    sleep(1)
    print(cos(45))

原正则问题分析 #

你之前的正则试图同时匹配import语句和紧跟的方法调用，但实际代码中import和调用往往不在同一行，甚至间隔大量代码，所以这种单一正则无法覆盖所有场景，必须先收集所有调用的方法，再分别处理import和调用语句。

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