AcidMap/scripts/demos/conda_env_export_tool.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Conda环境导出工具

该脚本结合conda显式安装的包（包含版本信息）和pip显式安装的包，创建一个完整的environment.yml文件。

特性:
- 📦 导出conda显式安装的包，包含精确的版本信息
- 🐍 导出pip安装的包，保持完整的版本约束
- ✨ 自动合并两种包管理器的依赖项
- 🎯 只包含用户显式安装的包，避免冗余依赖

使用方法:
python scripts/demos/conda_env_export_tool.py --env-name 环境名 --output environment.yml

python scripts/demos/conda_env_export_tool.py -e base -o environment.yml
作者: AcidMap项目组
"""

import argparse
import subprocess
import yaml
import tempfile
import os
import json
import re
from pathlib import Path


def run_command(command):
    """执行命令并返回输出"""
    try:
        result = subprocess.run(command, shell=True, capture_output=True, text=True, encoding='utf-8')
        if result.returncode != 0:
            print(f"命令执行失败: {command}")
            print(f"错误信息: {result.stderr}")
            return None
        return result.stdout
    except Exception as e:
        print(f"执行命令时出错: {e}")
        return None


def get_conda_explicit_packages(env_name):
    """获取conda显式安装的包（包含版本信息）"""
    print(f"正在获取环境 '{env_name}' 的conda显式安装包...")
    
    with tempfile.NamedTemporaryFile(mode='w+', suffix='.yml', delete=False) as temp_file:
        temp_path = temp_file.name
    
    try:
        # 首先获取显式安装的包名（无版本）
        command = f"conda env export --name {env_name} --from-history --file {temp_path}"
        output = run_command(command)
        
        if output is None:
            return None
            
        # 读取显式安装的包名
        with open(temp_path, 'r', encoding='utf-8') as f:
            conda_env = yaml.safe_load(f)
        
        # 获取显式安装的包名列表
        explicit_packages = []
        if 'dependencies' in conda_env:
            for dep in conda_env['dependencies']:
                if isinstance(dep, str):
                    # 使用更健壮的方法提取包名
                    # 匹配包名（字母、数字、下划线、连字符、点）直到遇到版本操作符或特殊字符
                    match = re.match(r'^([a-zA-Z0-9_.-]+)', dep.strip())
                    if match:
                        package_name = match.group(1)
                        explicit_packages.append(package_name)
                    else:
                        print(f"警告: 无法解析包名 '{dep}'，跳过")
                        continue
        
        # 获取所有已安装包的详细信息
        command = f"conda list --name {env_name} --json"
        list_output = run_command(command)
        
        if list_output is None:
            print("警告: 无法获取包的详细信息，将返回无版本的包列表")
            return conda_env
        
        all_packages = json.loads(list_output)
        
        # 创建包名到版本的映射
        package_versions = {}
        for pkg in all_packages:
            if pkg.get('channel') != 'pypi':  # 只处理conda包
                package_versions[pkg['name']] = f"{pkg['name']}={pkg['version']}"
        
        # 更新dependencies列表，添加版本信息
        versioned_dependencies = []
        for pkg_name in explicit_packages:
            if pkg_name in package_versions:
                versioned_dependencies.append(package_versions[pkg_name])
            else:
                print(f"警告: 未找到包 '{pkg_name}' 的版本信息，保持原样")
                versioned_dependencies.append(pkg_name)
        
        conda_env['dependencies'] = versioned_dependencies
        print(f"✅ 成功获取 {len(versioned_dependencies)} 个显式安装的conda包（包含版本信息）")
        
        return conda_env
        
    finally:
        # 清理临时文件
        if os.path.exists(temp_path):
            os.unlink(temp_path)


def get_pip_packages(env_name):
    """获取pip安装的包"""
    print(f"正在获取环境 '{env_name}' 的pip安装包...")
    
    with tempfile.NamedTemporaryFile(mode='w+', suffix='.yml', delete=False) as temp_file:
        temp_path = temp_file.name
    
    try:
        # 导出完整环境（包含pip包）
        command = f"conda env export --name {env_name} --file {temp_path}"
        output = run_command(command)
        
        if output is None:
            return None
            
        # 读取生成的文件
        with open(temp_path, 'r', encoding='utf-8') as f:
            full_env = yaml.safe_load(f)
        
        # 提取pip包
        pip_packages = []
        if 'dependencies' in full_env:
            for dep in full_env['dependencies']:
                if isinstance(dep, dict) and 'pip' in dep:
                    pip_packages = dep['pip']
                    break
        
        return pip_packages
        
    finally:
        # 清理临时文件
        if os.path.exists(temp_path):
            os.unlink(temp_path)


def create_combined_environment_file(env_name, output_file):
    """创建结合conda和pip包的环境文件"""
    
    # 获取conda显式包
    conda_env = get_conda_explicit_packages(env_name)
    if conda_env is None:
        print("获取conda包失败")
        return False
    
    # 获取pip包
    pip_packages = get_pip_packages(env_name)
    if pip_packages is None:
        print("获取pip包失败")
        return False
    
    # 合并
    if pip_packages:
        # 确保dependencies是列表
        if 'dependencies' not in conda_env:
            conda_env['dependencies'] = []
        
        # 添加pip包
        conda_env['dependencies'].append({'pip': pip_packages})
    
    # 检查文件是否已存在
    file_exists = os.path.exists(output_file)
    if file_exists:
        print(f"⚠️  文件 '{output_file}' 已存在，将被替换")
    
    # 写入输出文件
    try:
        with open(output_file, 'w', encoding='utf-8') as f:
            yaml.dump(conda_env, f, default_flow_style=False, allow_unicode=True, sort_keys=False)
        
        if file_exists:
            print(f"✅ 成功替换环境文件: {output_file}")
        else:
            print(f"✅ 成功创建环境文件: {output_file}")
        print(f"📦 Conda显式包数量: {len([dep for dep in conda_env['dependencies'] if isinstance(dep, str)])}")
        print(f"🐍 Pip包数量: {len(pip_packages) if pip_packages else 0}")
        
        return True
        
    except Exception as e:
        print(f"写入文件时出错: {e}")
        return False


def main():
    """主函数"""
    parser = argparse.ArgumentParser(
        description="导出包含conda显式包（含版本）和pip包的环境文件",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
示例:
  python scripts/demos/conda_env_export_tool.py --env-name GeoSys --output my_environment.yml
  python scripts/demos/conda_env_export_tool.py -e myenv -o environment.yml
  
注意: 该工具会自动获取显式安装包的版本信息，确保环境的精确重现。
        """
    )
    
    parser.add_argument(
        '--env-name', '-e',
        required=True,
        help="conda环境名称"
    )
    
    parser.add_argument(
        '--output', '-o',
        default='environment_combined.yml',
        help="输出文件名 (默认: environment_combined.yml)"
    )
    
    args = parser.parse_args()
    
    print(f"🚀 开始导出环境 '{args.env_name}' 到文件 '{args.output}'")
    print("=" * 50)
    
    success = create_combined_environment_file(args.env_name, args.output)
    
    if success:
        print("=" * 50)
        print("✨ 导出完成！")
        print(f"\n📝 使用此文件创建新环境:")
        print(f"   conda env create -f {args.output}")
    else:
        print("❌ 导出失败")
        return 1
    
    return 0


if __name__ == "__main__":
    exit(main())