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

"""
掃描目錄下所有 Python 文件，提取所有 import 語(yǔ)句中的模塊名，
去重后保存到以時(shí)間戳命名的文本文件中。
"""

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


def extract_imports_from_file(filepath):
    """
    解析單個(gè) Python 文件，返回該文件中導(dǎo)入的所有模塊名（集合）。
    支持：
        - import module
        - import module as alias
        - from module import name
        - from module import name as alias
        - 相對(duì)導(dǎo)入（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:
                    # 絕對(duì)導(dǎo)入或相對(duì)導(dǎo)入但指定了模塊（如 from .sub import x）
                    if node.level > 0:
                        # 相對(duì)導(dǎo)入，模塊名前加對(duì)應(yīng)數(shù)量的點(diǎn)
                        module = '.' * node.level + node.module
                    else:
                        module = node.module
                else:
                    # 純相對(duì)導(dǎo)入（如 from . import x, from .. import y）
                    module = '.' * node.level
                modules.add(module)
    except Exception as e:
        print(f"警告：解析文件 {filepath} 時(shí)出錯(cuò)：{e}", file=sys.stderr)

    return modules


def main():
    parser = argparse.ArgumentParser(
        description='掃描目錄下的 Python 文件，提取所有 import 模塊名并去重輸出。'
    )
    parser.add_argument(
        'directory',
        nargs='?',
        default='.',
        help='要掃描的目錄（默認(rèn)：當(dāng)前目錄）'
    )
    args = parser.parse_args()

    target_dir = args.directory
    if not os.path.isdir(target_dir):
        print(f"錯(cuò)誤：'{target_dir}' 不是一個(gè)有效的目錄。", 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 語(yǔ)句。")
        return

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

    # 生成時(shí)間戳文件名
    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)} 個(gè)唯一模塊。")
    print(f"結(jié)果已寫(xiě)入：{output_filename}")


if __name__ == '__main__':
    main()

python scan_imports.py [要掃描的目錄](méi)