Python-自动生成 requirements.txt

generate_requirements.py：

"""
generate_requirements.py - 自动生成 requirements.txt
    功能：
    1. 扫描 code 目录下所有 .py 文件
    2. 提取所有 import 语句
    3. 过滤出第三方库（排除标准库和本地模块）
    4. 生成 requirements.txt 文件
"""

import os
import ast
import sys
import argparse
from pathlib import Path

# Python 标准库列表（部分常用的）
STDLIB_MODULES = {
    'os', 'sys', 're', 'time', 'datetime', 'json', 'csv', 'math', 'random',
    'collections', 'itertools', 'functools', 'operator', 'copy', 'pickle',
    'subprocess', 'threading', 'multiprocessing', 'queue', 'socket', 'urllib',
    'http', 'email', 'html', 'xml', 'logging', 'argparse', 'io', 'pathlib',
    'shutil', 'tempfile', 'glob', 'fnmatch', 'gzip', 'zipfile', 'tarfile',
    'hashlib', 'hmac', 'secrets', 'base64', 'struct', 'codecs', 'warnings',
    'contextlib', 'abc', 'typing', 'enum', 'dataclasses', 'ast', 'dis',
    'inspect', 'gc', 'weakref', 'platform', 'ctypes', 'unittest', 'doctest'
}

# 第三方库的包名映射（import 名 -> pip 包名）
PACKAGE_MAPPING = {
    'cv2': 'opencv-python-headless',
    'sklearn': 'scikit-learn',
    'PIL': 'Pillow',
    'yaml': 'PyYAML',
}


def extract_imports_from_file(file_path):
    """
    从单个 Python 文件中提取所有 import 语句
    返回: set of module names
    """
    imports = set()
    
    try:
        with open(file_path, 'r', encoding='utf-8') as f:
            tree = ast.parse(f.read())
        
        for node in ast.walk(tree):
            # import xxx
            if isinstance(node, ast.Import):
                for alias in node.names:
                    imports.add(alias.name.split('.')[0])
            
            # from xxx import yyy
            elif isinstance(node, ast.ImportFrom):
                if node.module:
                    imports.add(node.module.split('.')[0])
    
    except Exception as e:
        print(f"⚠️  解析文件失败: {file_path}, 错误: {e}")
    
    return imports


def scan_directory(directory):
    """
    扫描目录下所有 .py 文件，提取所有导入的模块
    返回: set of module names
    """
    all_imports = set()
    
    py_files = list(Path(directory).rglob('*.py'))
    print(f"📂 扫描目录: {directory}")
    print(f"📄 找到 {len(py_files)} 个 Python 文件\n")
    
    for py_file in py_files:
        imports = extract_imports_from_file(py_file)
        all_imports.update(imports)
    
    return all_imports


def filter_third_party_packages(imports, local_modules):
    """
    过滤出第三方库，排除标准库和本地模块
    """
    third_party = set()
    
    for module in imports:
        # 排除标准库
        if module in STDLIB_MODULES:
            continue
        
        # 排除本地模块（当前项目中的 .py 文件）
        if module in local_modules:
            continue
        
        # 添加到第三方库列表
        third_party.add(module)
    
    return third_party


def get_package_names(modules):
    """
    将 import 的模块名转换为 pip 包名
    """
    packages = set()
    
    for module in modules:
        # 使用映射表
        package_name = PACKAGE_MAPPING.get(module, module)
        packages.add(package_name)
    
    return sorted(packages)


def get_local_modules(directory):
    """
    获取项目中所有本地模块名（不含 .py 后缀）
    """
    local_modules = set()
    
    for py_file in Path(directory).rglob('*.py'):
        module_name = py_file.stem
        if module_name != '__init__':
            local_modules.add(module_name)
    
    return local_modules


def generate_requirements(directory, output_file, with_comments=True):
    """
    生成 requirements.txt 文件
    """
    # 1. 获取本地模块列表
    local_modules = get_local_modules(directory)
    print(f"🔍 本地模块: {', '.join(sorted(local_modules))}\n")
    
    # 2. 扫描所有导入
    all_imports = scan_directory(directory)
    print(f"📦 总共发现 {len(all_imports)} 个导入模块\n")
    
    # 3. 过滤第三方库
    third_party = filter_third_party_packages(all_imports, local_modules)
    print(f"✅ 识别出 {len(third_party)} 个第三方库:")
    for lib in sorted(third_party):
        print(f"   - {lib}")
    print()
    
    # 4. 转换为 pip 包名
    packages = get_package_names(third_party)
    
    # 5. 写入 requirements.txt
    with open(output_file, 'w', encoding='utf-8') as f:
        if with_comments:
            f.write("# requirements.txt - 环境依赖\n")
            f.write("# 自动生成，列出项目所需的Python包\n\n")
        
        for package in packages:
            f.write(f"{package}\n")
    
    print(f"💾 已生成 requirements.txt 到: {output_file}")
    print(f"📋 包含 {len(packages)} 个依赖包")


def main():
    parser = argparse.ArgumentParser(description='自动生成 requirements.txt')
    
    parser.add_argument('--dir', type=str, default='.',
                        help='要扫描的代码目录（默认: 当前目录）')
    
    parser.add_argument('--output', type=str, default='requirements.txt',
                        help='输出文件路径（默认: requirements.txt）')
    
    parser.add_argument('--no-comments', action='store_true',
                        help='不添加注释')
    
    args = parser.parse_args()
    
    # 执行生成
    generate_requirements(
        directory=args.dir,
        output_file=args.output,
        with_comments=not args.no_comments
    )


if __name__ == "__main__":
    main()