Python脚本实现扫描指定目录下.py文件并提取import语句

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
扫描目录下所有 Python 文件，提取所有 import 语句中的模块名，
去重后保存到以时间戳命名的文本文件中。
"""

import os
import sys
import ast
import argparse
from datetime import datetime


def extract_imports_from_file(filepath):
    """
    解析单个 Python 文件，返回该文件中导入的所有模块名（集合）。
    支持：
        - import module
        - import module as alias
        - from module import name
        - from module import name as alias
        - 相对导入（from . import x, from ..sub import y 等）
    """
    modules = set()
    try:
        with open(filepath, 'r', encoding='utf-8') as f:
            content = f.read()
        tree = ast.parse(content, filename=filepath)

        for node in ast.walk(tree):
            if isinstance(node, ast.Import):
                # import module, import module as alias
                for alias in node.names:
                    modules.add(alias.name)  # 完整模块名
            elif isinstance(node, ast.ImportFrom):
                # from module import ...
                if node.module:
                    # 绝对导入或相对导入但指定了模块（如 from .sub import x）
                    if node.level > 0:
                        # 相对导入，模块名前加对应数量的点
                        module = '.' * node.level + node.module
                    else:
                        module = node.module
                else:
                    # 纯相对导入（如 from . import x, from .. import y）
                    module = '.' * node.level
                modules.add(module)
    except Exception as e:
        print(f"警告：解析文件 {filepath} 时出错：{e}", file=sys.stderr)

    return modules


def main():
    parser = argparse.ArgumentParser(
        description='扫描目录下的 Python 文件，提取所有 import 模块名并去重输出。'
    )
    parser.add_argument(
        'directory',
        nargs='?',
        default='.',
        help='要扫描的目录（默认：当前目录）'
    )
    args = parser.parse_args()

    target_dir = args.directory
    if not os.path.isdir(target_dir):
        print(f"错误：'{target_dir}' 不是一个有效的目录。", file=sys.stderr)
        sys.exit(1)

    all_modules = set()

    # 递归遍历目录
    for root, dirs, files in os.walk(target_dir):
        for file in files:
            if file.endswith('.py'):
                filepath = os.path.join(root, file)
                modules = extract_imports_from_file(filepath)
                all_modules.update(modules)

    if not all_modules:
        print("未找到任何 import 语句。")
        return

    # 排序后输出
    sorted_modules = sorted(all_modules)

    # 生成时间戳文件名
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    output_filename = f"imports_{timestamp}.txt"

    with open(output_filename, 'w', encoding='utf-8') as f:
        for module in sorted_modules:
            f.write(module + '\n')

    print(f"共找到 {len(sorted_modules)} 个唯一模块。")
    print(f"结果已写入：{output_filename}")


if __name__ == '__main__':
    main()

python scan_imports.py [要扫描的目录]