strepsiades
61 lines
2.1 KiB
Python
61 lines
2.1 KiB
Python
"""Units Domain - JSONL Handler"""
|
||
import json
|
||
import re
|
||
from typing import List
|
||
|
||
from ...datasource import BaseJsonlHandler, JsonlDataSource
|
||
from .sync_node import UnitsSyncNode, SyncNode
|
||
from schemas.project_extensions.generator_unit import GeneratorUnitProject
|
||
|
||
|
||
class UnitsJsonlHandler(BaseJsonlHandler):
|
||
"""机组 JSONL Handler"""
|
||
def __init__(self, datasource: JsonlDataSource):
|
||
super().__init__(datasource, UnitsSyncNode, schema=GeneratorUnitProject)
|
||
|
||
async def load(self) -> List['SyncNode']:
|
||
"""加载机组数据（支持 generator_unit_N.jsonl 与 units_N.jsonl）"""
|
||
nodes: List[SyncNode] = []
|
||
seen_ids: set[str] = set()
|
||
|
||
if not self.datasource.dir_path.exists():
|
||
return nodes
|
||
|
||
# 查找所有匹配的文件：generator_unit_N.jsonl 或 units_N.jsonl
|
||
patterns = (
|
||
re.compile(r"^generator_unit_\d+\.jsonl$"),
|
||
re.compile(r"^units_\d+\.jsonl$"),
|
||
)
|
||
|
||
for file_path in self.datasource.dir_path.iterdir():
|
||
if not any(pattern.match(file_path.name) for pattern in patterns):
|
||
continue
|
||
|
||
# 读取文件
|
||
with open(file_path, "r", encoding="utf-8") as f:
|
||
for line_num, line in enumerate(f, 1):
|
||
line = line.strip()
|
||
if not line:
|
||
continue
|
||
|
||
try:
|
||
item = json.loads(line)
|
||
item_id = self.extract_id(item)
|
||
if item_id and item_id in seen_ids:
|
||
continue
|
||
|
||
# 创建节点
|
||
node = self.create_node(item)
|
||
nodes.append(node)
|
||
|
||
# 同时加载到缓存
|
||
if item_id:
|
||
seen_ids.add(item_id)
|
||
bucket = self.datasource._data.setdefault(self.node_type, {})
|
||
bucket[item_id] = item
|
||
|
||
except json.JSONDecodeError as e:
|
||
print(f"Warning: Invalid JSON in {file_path.name}:{line_num}: {e}")
|
||
|
||
return nodes
|