126 lines
4.1 KiB
Python
126 lines
4.1 KiB
Python
#!/usr/bin/env python
|
|
"""Regenerate the byte-equivalence golden fixtures for the native docx
|
|
sidecar pipeline.
|
|
|
|
The fixtures live at
|
|
``tests/parser/docx/golden/native_docx/<scenario>/``
|
|
and capture the exact on-disk artifacts ``LightRAG.parse_native`` produces
|
|
for each scenario in ``tests/parser/docx/_native_docx_fixtures.py``.
|
|
|
|
Usage::
|
|
|
|
python scripts/regen_native_docx_golden.py
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import shutil
|
|
import sys
|
|
import tempfile
|
|
from pathlib import Path
|
|
from unittest import mock
|
|
|
|
PROJECT_ROOT = Path(__file__).resolve().parent.parent
|
|
sys.path.insert(0, str(PROJECT_ROOT))
|
|
sys.path.insert(0, str(PROJECT_ROOT / "tests" / "parser" / "docx"))
|
|
|
|
|
|
async def _regen() -> None:
|
|
from lightrag.constants import (
|
|
FULL_DOCS_FORMAT_PENDING_PARSE,
|
|
PARSED_DIR_NAME,
|
|
)
|
|
from lightrag.parser.base import ParseContext
|
|
from lightrag.parser.debug import (
|
|
FrozenDateTime,
|
|
build_debug_rag,
|
|
)
|
|
from lightrag.parser.registry import get_parser
|
|
import lightrag.pipeline as pipeline_module
|
|
|
|
from _native_docx_fixtures import SCENARIOS # type: ignore[import]
|
|
|
|
fixtures_root = (
|
|
PROJECT_ROOT / "tests" / "parser" / "docx" / "golden" / "native_docx"
|
|
)
|
|
fixtures_root.mkdir(parents=True, exist_ok=True)
|
|
|
|
async def _noop_archive(_p: str) -> None:
|
|
return None
|
|
|
|
for scenario in SCENARIOS:
|
|
scenario_dir = fixtures_root / scenario.name
|
|
if scenario_dir.exists():
|
|
shutil.rmtree(scenario_dir)
|
|
scenario_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
def _stub_extract(
|
|
file_path,
|
|
*,
|
|
drawing_context=None,
|
|
parse_warnings=None,
|
|
parse_metadata=None,
|
|
**_kwargs,
|
|
):
|
|
if drawing_context is not None and scenario.assets:
|
|
drawing_context.export_dir_path.mkdir(parents=True, exist_ok=True)
|
|
for name, data in scenario.assets.items():
|
|
(drawing_context.export_dir_path / name).write_bytes(data)
|
|
if parse_metadata is not None:
|
|
parse_metadata.update(scenario.parse_metadata)
|
|
return [dict(b) for b in scenario.blocks]
|
|
|
|
with tempfile.TemporaryDirectory(prefix="regen_") as tmp_root:
|
|
input_dir = Path(tmp_root) / "inputs"
|
|
input_dir.mkdir()
|
|
source_path = input_dir / scenario.file_path
|
|
source_path.parent.mkdir(parents=True, exist_ok=True)
|
|
source_path.write_bytes(b"fake-docx")
|
|
|
|
rag = build_debug_rag()
|
|
|
|
with (
|
|
mock.patch.dict("os.environ", {"INPUT_DIR": str(input_dir)}),
|
|
mock.patch(
|
|
"lightrag.parser.docx.parse_document.extract_docx_blocks",
|
|
_stub_extract,
|
|
),
|
|
mock.patch.object(
|
|
pipeline_module,
|
|
"archive_docx_source_after_full_docs_sync",
|
|
_noop_archive,
|
|
),
|
|
mock.patch("lightrag.sidecar.writer.datetime", FrozenDateTime),
|
|
):
|
|
await get_parser("native").parse(
|
|
ParseContext(
|
|
rag,
|
|
scenario.doc_id,
|
|
str(source_path),
|
|
{
|
|
"parse_format": FULL_DOCS_FORMAT_PENDING_PARSE,
|
|
"content": "",
|
|
},
|
|
)
|
|
)
|
|
|
|
produced_dir = input_dir / PARSED_DIR_NAME / f"{scenario.file_path}.parsed"
|
|
for item in produced_dir.rglob("*"):
|
|
rel = item.relative_to(produced_dir)
|
|
target = scenario_dir / rel
|
|
target.parent.mkdir(parents=True, exist_ok=True)
|
|
if item.is_dir():
|
|
target.mkdir(exist_ok=True)
|
|
else:
|
|
shutil.copyfile(item, target)
|
|
print(f"[regen] {scenario.name}: wrote {scenario_dir}")
|
|
|
|
|
|
def main() -> None:
|
|
asyncio.run(_regen())
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|