fix.py

#!/usr/bin/env python3
"""
修复 quant_strategy.json 中的列表结构问题
将 [{...}] 格式转换为 {...} 格式

用法：
    python fix_quant_strategy.py quant_strategy.json
    python fix_quant_strategy.py quant_strategy.json -o fixed_strategy.json
"""

import json
import argparse
from pathlib import Path


def fix_strategy_structure(input_file, output_file=None):
    """
    修复 strategy 文件中的列表结构
    
    Args:
        input_file: 输入的 quant_strategy.json 文件路径
        output_file: 输出文件路径（如果为 None，则覆盖原文件）
    """
    input_path = Path(input_file)
    
    if not input_path.exists():
        print(f"❌ 错误: 文件不存在 {input_path}")
        return False
    
    print(f"📖 读取文件: {input_path}")
    
    # 读取原始 JSON
    try:
        with open(input_path, 'r', encoding='utf-8') as f:
            data = json.load(f)
    except json.JSONDecodeError as e:
        print(f"❌ JSON 解析错误: {e}")
        return False
    
    if 'measurement' not in data:
        print("❌ 错误: JSON 中没有 'measurement' 字段")
        return False
    
    measurement = data['measurement']
    new_measurement = {}
    fixed_count = 0
    kept_count = 0
    
    # 遍历所有层并修复结构
    for key, value in measurement.items():
        if isinstance(value, list):
            if len(value) > 0:
                # 从列表中取出字典
                new_measurement[key] = value[0]
                fixed_count += 1
                print(f"  ✅ 修复: {key} (list -> dict)")
            else:
                print(f"  ⚠️  警告: {key} 是空列表，跳过")
                new_measurement[key] = value
        elif isinstance(value, dict):
            # 已经是字典，保持不变
            new_measurement[key] = value
            kept_count += 1
        else:
            print(f"  ⚠️  警告: {key} 的值类型未知 ({type(value)})")
            new_measurement[key] = value
    
    # 更新数据
    data['measurement'] = new_measurement
    
    # 确定输出文件路径
    if output_file is None:
        output_path = input_path
        print(f"\n💾 覆盖原文件: {output_path}")
    else:
        output_path = Path(output_file)
        print(f"\n💾 保存到: {output_path}")
    
    # 写入修复后的 JSON
    try:
        with open(output_path, 'w', encoding='utf-8') as f:
            json.dump(data, f, indent=2, ensure_ascii=False)
    except Exception as e:
        print(f"❌ 写入文件失败: {e}")
        return False
    
    # 打印统计信息
    print(f"\n📊 统计:")
    print(f"   - 修复的层数: {fixed_count}")
    print(f"   - 保持不变: {kept_count}")
    print(f"   - 总层数: {len(new_measurement)}")
    print(f"\n✅ 完成!")
    
    return True


def main():
    parser = argparse.ArgumentParser(
        description='修复 quant_strategy.json 中的列表结构问题',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
示例:
  # 直接覆盖原文件
  python fix_quant_strategy.py quant_strategy.json
  
  # 保存到新文件
  python fix_quant_strategy.py quant_strategy.json -o fixed_strategy.json
        """
    )
    
    parser.add_argument(
        'input_file',
        help='输入的 quant_strategy.json 文件路径'
    )
    
    parser.add_argument(
        '-o', '--output',
        dest='output_file',
        default=None,
        help='输出文件路径（默认覆盖原文件）'
    )
    
    args = parser.parse_args()
    
    success = fix_strategy_structure(args.input_file, args.output_file)
    
    return 0 if success else 1


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