60 lines
2.0 KiB
Python
60 lines
2.0 KiB
Python
import json
|
|
import os
|
|
from typing import Dict, Any
|
|
|
|
|
|
class OutputWriter:
|
|
"""将 AI 生成的数据写入文件系统。"""
|
|
|
|
def __init__(self, output_dir: str):
|
|
self.output_dir = output_dir
|
|
|
|
def write(self, program_id: str, ai_result: Dict[str, Any],
|
|
input_type: str) -> Dict[str, str]:
|
|
"""写入所有组的输出文件。
|
|
|
|
Returns:
|
|
{group_folder: written_file_path} 映射
|
|
"""
|
|
written = {}
|
|
groups = ai_result.get('groups', ai_result)
|
|
|
|
for group_key in sorted(groups.keys()):
|
|
group_data = groups[group_key]
|
|
group_dir = os.path.join(self.output_dir, program_id, group_key)
|
|
os.makedirs(group_dir, exist_ok=True)
|
|
|
|
if input_type in ('file', 'mixed'):
|
|
json_path = self._write_json(group_dir, program_id, group_key, group_data)
|
|
written[f"{group_key}/json"] = json_path
|
|
|
|
if input_type in ('db', 'mixed'):
|
|
sql_path = self._write_sql(group_dir, program_id, group_key, group_data)
|
|
written[f"{group_key}/sql"] = sql_path
|
|
|
|
return written
|
|
|
|
def _write_json(self, group_dir: str, program_id: str,
|
|
group_key: str, data: Any) -> str:
|
|
"""写入 JSON 文件。"""
|
|
filename = f"{program_id}_{group_key}.json"
|
|
filepath = os.path.join(group_dir, filename)
|
|
|
|
with open(filepath, 'w', encoding='utf-8') as f:
|
|
json.dump(data, f, ensure_ascii=False, indent=2)
|
|
|
|
return filepath
|
|
|
|
def _write_sql(self, group_dir: str, program_id: str,
|
|
group_key: str, data: Any) -> str:
|
|
"""写入 SQL 文件。"""
|
|
filename = f"{program_id}_{group_key}.sql"
|
|
filepath = os.path.join(group_dir, filename)
|
|
|
|
sql_content = data if isinstance(data, str) else data.get('sql', json.dumps(data, ensure_ascii=False))
|
|
|
|
with open(filepath, 'w', encoding='utf-8') as f:
|
|
f.write(sql_content)
|
|
|
|
return filepath
|