使用抽象语法树`ast`统计哪些Python包与模块被导入了

长话短说，我的Gist。

给定一个没有requirements.txt的Python项目，如果想知道需要安装哪些包才能满足这个项目的依赖需求，一个容易想到的方法就是对每一个.py文件，用模式匹配（如正则表达式）找import xxx，并记录xxx为需要的包。然而import语句有很多形式，如：import xxx、import xxx as aaa、import xxx as aaa, yyy as bbb、from xxx.yyy import fff as ccc、from .zzz import ggg。因此，更好的方法是利用抽象语法树ast模块来找出所有import语句。

Python的import语句对应ast的两种节点：ast.Import和ast.ImportFrom。要从ast.Import获取导入包的列表，可用：

[a.name for a in node.names]  # 其中node是ast.Import类型的

要从ast.ImportFrom获取导入的包，可用：

node.module  # 其中node是ast.ImportFrom类型的

值得注意的是如果当前import语句是from . import xxx，node.module将会是None，此时node.level > 0，意味着相对导入。因此，要想获得所有导入的包（除了相对导入外，因为相对导入的包绝不会是需要安装的依赖），可以这样：

import ast
# 假设source包含待解析源码
root = ast.parse(source)
result = []
for node in ast.walk(root):
    if isinstance(node, ast.Import):
        for a in node.names:
            result.append(a.name.split('.', maxsplit=1)[0])
    elif isinstance(node, ast.ImportFrom):
        if node.level == 0:
            result.append(node.module.split('.', maxsplit=1)[0])

然而绝对导入的包也有可能是工作目录中已存在的模块或包啊，此时我们就可以根据导入路径判断它是不是指工作目录下的包：

def exists_local(path, rootpkg):
    filepath = os.path.join(rootpkg, path.replace('.', os.path.sep))
    # see if path is a local package
    if os.path.isdir(filepath) and os.path.isfile(
            os.path.join(filepath, '__init__.py')):
        return True
    # see if path is a local module
    if os.path.isfile(filepath + '.py'):
        return True

    return False

其中path是导入路径，rootpkg是根包所在目录（定义见这里）。

把这个核心功能稍作包装，便可写出下面的完整可执行代码：

from __future__ import print_function

import argparse
import os
import ast
import sys
import pkgutil
import itertools
import logging
import json


def make_parser():
    parser = argparse.ArgumentParser(
        description=(