Add merge_results: склейка json-документов процессора в один результат
PDF-процессор кладёт по одному json на каждый распознанный документ.
merge_documents() собирает их в {file_id, legal_entity_id, files: [...]},
где files — полное содержимое каждого файла без обработки.
Модуль не зависит от config/S3/NATS: используется и как CLI-скрипт,
и (в дальнейшем) из PollWorker.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01JvCSGmwgL8juhZdXr32qXM
This commit is contained in:
parent
1794654c39
commit
e44ee61548
94
src/core/merge_results.py
Normal file
94
src/core/merge_results.py
Normal file
@ -0,0 +1,94 @@
|
|||||||
|
"""Склейка результатов PDF-процессора в один JSON.
|
||||||
|
|
||||||
|
Процессор кладёт по одному json на каждый распознанный документ
|
||||||
|
({'file_uid': ..., 'error': ..., 'data_ocr': ...}). Здесь они собираются
|
||||||
|
в единый результат вида:
|
||||||
|
|
||||||
|
{'file_id': ..., 'legal_entity_id': ..., 'files': [<содержимое каждого>]}
|
||||||
|
|
||||||
|
Содержимое документов не обрабатывается и не валидируется — кладётся как есть.
|
||||||
|
|
||||||
|
Модуль намеренно не зависит от config/S3/NATS, чтобы его можно было
|
||||||
|
использовать и как отдельный скрипт, и из PollWorker.
|
||||||
|
"""
|
||||||
|
|
||||||
|
import argparse
|
||||||
|
import json
|
||||||
|
import sys
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
|
||||||
|
def load_document(path) -> object:
|
||||||
|
"""Читает один json-документ процессора."""
|
||||||
|
with open(path, 'r', encoding='utf-8') as f:
|
||||||
|
return json.load(f)
|
||||||
|
|
||||||
|
|
||||||
|
def merge_documents(paths, file_id, legal_entity_id) -> dict:
|
||||||
|
"""Собирает документы в один результат.
|
||||||
|
|
||||||
|
Порядок элементов в 'files' — порядок переданных путей.
|
||||||
|
"""
|
||||||
|
paths = list(paths)
|
||||||
|
if not paths:
|
||||||
|
raise ValueError('нечего мержить: не передано ни одного документа')
|
||||||
|
|
||||||
|
return {
|
||||||
|
'file_id': file_id,
|
||||||
|
'legal_entity_id': legal_entity_id,
|
||||||
|
'files': [load_document(path) for path in paths],
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
def merge_directory(job_dir, file_id, legal_entity_id) -> dict:
|
||||||
|
"""То же, но берёт все *.json из директории, отсортированные по имени."""
|
||||||
|
paths = sorted(Path(job_dir).glob('*.json'))
|
||||||
|
if not paths:
|
||||||
|
raise ValueError(f'в директории нет json-файлов: {job_dir}')
|
||||||
|
return merge_documents(paths, file_id, legal_entity_id)
|
||||||
|
|
||||||
|
|
||||||
|
def main(argv=None) -> int:
|
||||||
|
parser = argparse.ArgumentParser(
|
||||||
|
description='Склеить json-документы процессора в один результат.',
|
||||||
|
)
|
||||||
|
parser.add_argument(
|
||||||
|
'documents',
|
||||||
|
nargs='*',
|
||||||
|
help='пути к json-документам (порядок сохраняется в files[])',
|
||||||
|
)
|
||||||
|
parser.add_argument(
|
||||||
|
'-d', '--dir',
|
||||||
|
help='взять все *.json из директории вместо перечисления путей',
|
||||||
|
)
|
||||||
|
parser.add_argument('--file-id', required=True, help='file_id исходного файла')
|
||||||
|
parser.add_argument('--legal-entity-id', required=True, help='legal_entity_id')
|
||||||
|
parser.add_argument(
|
||||||
|
'-o', '--output',
|
||||||
|
help='куда записать результат (по умолчанию — stdout)',
|
||||||
|
)
|
||||||
|
args = parser.parse_args(argv)
|
||||||
|
|
||||||
|
if bool(args.documents) == bool(args.dir):
|
||||||
|
parser.error('укажи либо пути к документам, либо --dir, но не оба сразу')
|
||||||
|
|
||||||
|
try:
|
||||||
|
if args.dir:
|
||||||
|
result = merge_directory(args.dir, args.file_id, args.legal_entity_id)
|
||||||
|
else:
|
||||||
|
result = merge_documents(args.documents, args.file_id, args.legal_entity_id)
|
||||||
|
except (OSError, ValueError, json.JSONDecodeError) as e:
|
||||||
|
print(f'ошибка: {e}', file=sys.stderr)
|
||||||
|
return 1
|
||||||
|
|
||||||
|
payload = json.dumps(result, ensure_ascii=False, indent=2)
|
||||||
|
if args.output:
|
||||||
|
Path(args.output).write_text(payload, encoding='utf-8')
|
||||||
|
print(f'записано {len(result["files"])} документов в {args.output}', file=sys.stderr)
|
||||||
|
else:
|
||||||
|
print(payload)
|
||||||
|
return 0
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == '__main__':
|
||||||
|
raise SystemExit(main())
|
||||||
83
tests/test_merge_results.py
Normal file
83
tests/test_merge_results.py
Normal file
@ -0,0 +1,83 @@
|
|||||||
|
import json
|
||||||
|
import subprocess
|
||||||
|
import sys
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
|
||||||
|
from src.core.merge_results import merge_documents
|
||||||
|
|
||||||
|
DOC_A = {
|
||||||
|
'file_uid': 'uid-a',
|
||||||
|
'error': {'code': '', 'message': ''},
|
||||||
|
'data_ocr': {'label': {'form_name': 'Бухгалтерский баланс'}, 'table_structure': []},
|
||||||
|
}
|
||||||
|
DOC_B = {
|
||||||
|
'file_uid': 'uid-b',
|
||||||
|
'error': {'code': '', 'message': ''},
|
||||||
|
'data_ocr': {'label': {'form_name': 'Отчёт о финансовых результатах'}, 'table_structure': []},
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
def write_doc(directory: Path, name: str, payload: dict) -> Path:
|
||||||
|
path = directory / name
|
||||||
|
path.write_text(json.dumps(payload, ensure_ascii=False), encoding='utf-8')
|
||||||
|
return path
|
||||||
|
|
||||||
|
|
||||||
|
def test_merges_two_documents_into_one_result(tmp_path):
|
||||||
|
paths = [write_doc(tmp_path, 'doc_1.json', DOC_A), write_doc(tmp_path, 'doc_2.json', DOC_B)]
|
||||||
|
|
||||||
|
result = merge_documents(paths, file_id='file-1', legal_entity_id='le-1')
|
||||||
|
|
||||||
|
assert result == {
|
||||||
|
'file_id': 'file-1',
|
||||||
|
'legal_entity_id': 'le-1',
|
||||||
|
'files': [DOC_A, DOC_B],
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
def test_keeps_document_content_untouched(tmp_path):
|
||||||
|
# содержимое кладётся как есть: незнакомые поля и вложенность не теряются
|
||||||
|
weird = {'file_uid': 'uid-c', 'unexpected': [1, {'deep': None}], 'кириллица': 'значение'}
|
||||||
|
paths = [write_doc(tmp_path, 'doc_1.json', weird)]
|
||||||
|
|
||||||
|
result = merge_documents(paths, file_id='file-1', legal_entity_id='le-1')
|
||||||
|
|
||||||
|
assert result['files'] == [weird]
|
||||||
|
|
||||||
|
|
||||||
|
def test_preserves_given_order(tmp_path):
|
||||||
|
# порядок в files[] — порядок переданных путей, а не порядок записи на диск
|
||||||
|
second = write_doc(tmp_path, 'z_doc.json', DOC_B)
|
||||||
|
first = write_doc(tmp_path, 'a_doc.json', DOC_A)
|
||||||
|
|
||||||
|
result = merge_documents([second, first], file_id='f', legal_entity_id='l')
|
||||||
|
|
||||||
|
assert [doc['file_uid'] for doc in result['files']] == ['uid-b', 'uid-a']
|
||||||
|
|
||||||
|
|
||||||
|
def test_rejects_empty_input():
|
||||||
|
with pytest.raises(ValueError):
|
||||||
|
merge_documents([], file_id='f', legal_entity_id='l')
|
||||||
|
|
||||||
|
|
||||||
|
def test_cli_writes_merged_file(tmp_path):
|
||||||
|
a = write_doc(tmp_path, 'doc_1.json', DOC_A)
|
||||||
|
b = write_doc(tmp_path, 'doc_2.json', DOC_B)
|
||||||
|
out = tmp_path / 'merged.json'
|
||||||
|
|
||||||
|
completed = subprocess.run(
|
||||||
|
[sys.executable, '-m', 'src.core.merge_results', str(a), str(b),
|
||||||
|
'--file-id', 'file-1', '--legal-entity-id', 'le-1', '-o', str(out)],
|
||||||
|
cwd=Path(__file__).resolve().parent.parent,
|
||||||
|
capture_output=True,
|
||||||
|
text=True,
|
||||||
|
)
|
||||||
|
|
||||||
|
assert completed.returncode == 0, completed.stderr
|
||||||
|
assert json.loads(out.read_text(encoding='utf-8')) == {
|
||||||
|
'file_id': 'file-1',
|
||||||
|
'legal_entity_id': 'le-1',
|
||||||
|
'files': [DOC_A, DOC_B],
|
||||||
|
}
|
||||||
Loading…
x
Reference in New Issue
Block a user