Python 轻松读取 py 模块中的所有导入语句的方法

Question

我正在尝试创建一个辅助函数，用来读取一个文件，并在单元测试中模拟所有的导入。因为我没有那些东西在Python的路径上，所以我必须读取文件，而不是直接导入。

下面是示例代码：

---

#module.py
import com.stackoverflow.question
from com.stackoverflow.util import test_func
from com.stackoverflow.util import TestClass

---

#magic helper: what i want
magic = process('<path_to>/module.py')
for module in magic.modules_as_strings():
    #todo  would have to recuirsively add each path
    # so i would first create com, then com.stackoverflow, etc
    setattr(self, module, StubModules(module)
for obj in magic.sink:
    #these would be "from"  from x import Y
    #its basically just creating self.Y = object
    setattr(self, object)

上面的代码是模拟的部分，我其实想找到一个更好的方法来处理文件中的“from/import语句”。

这样说你明白了吗？我知道我可以逐行读取文件，但我希望能有一种更简洁的方法。

如果你有任何问题，请告诉我。

Answer 1

使用AST模块，其实很简单：

import ast
from collections import namedtuple

Import = namedtuple("Import", ["module", "name", "alias"])

def get_imports(path):
    with open(path) as fh:        
       root = ast.parse(fh.read(), path)

    for node in ast.iter_child_nodes(root):
        if isinstance(node, ast.Import):
            module = []
        elif isinstance(node, ast.ImportFrom):  
            module = node.module.split('.')
        else:
            continue

        for n in node.names:
            yield Import(module, n.name.split('.'), n.asname)

对于这样的模块：

from coco import bunny
from coco.bungy import carrot
from meta import teta
from rocket import spaceship as sp
import bingo
import com.stackoverflow
import motorbike as car
import module1, module2

s="a random variable"

def func():
    """And a function"""

输出结果是：

>>> for imp in get_imports("/path/to/file.py"): print imp
Import(module=['coco'], name=['bunny'], alias=None)
Import(module=['coco', 'bungy'], name=['carrot'], alias=None)
Import(module=['meta'], name=['teta'], alias=None)
Import(module=['rocket'], name=['spaceship'], alias='sp')
Import(module=[], name=['bingo'], alias=None)
Import(module=[], name=['com', 'stackoverflow'], alias=None)
Import(module=[], name=['motorbike'], alias='car')
Import(module=[], name=['module1'], alias=None)
Import(module=[], name=['module2'], alias=None)