data-extract / core /output.py
Soumik-404's picture
ok
abcd0c2
Raw
History Blame Contribute Delete
4.64 kB
"""
Output writer for the MarkItDown API.
Writes conversion results to disk in Markdown, plain text, or JSON format.
Intended for use by the CLI batch pipeline; the API server handles output
directly via HTTP responses and does not use this module.
"""
from __future__ import annotations
import json
from pathlib import Path
from typing import Literal
from .converter import ConversionResult
from logger import get_logger
logger = get_logger(__name__)
OutputFormat = Literal["markdown", "json", "txt"]
class OutputWriter:
"""Write ConversionResult objects to files in a specified format.
Parameters
----------
output_dir:
Destination directory. Created automatically if it does not exist.
format:
Output format: ``"markdown"`` (default), ``"json"``, or ``"txt"``.
"""
def __init__(self, output_dir: str | Path, format: OutputFormat = "markdown") -> None:
self._output_dir = Path(output_dir)
self._format = format
self._output_dir.mkdir(parents=True, exist_ok=True)
def write(self, result: ConversionResult) -> Path:
"""Serialise *result* and write it to the output directory.
Collisions are resolved by appending an incrementing counter to the stem.
Returns
-------
Path
The path of the written file.
"""
stem = Path(result.source).stem if not result.source.startswith("http") else "web_content"
suffix = ".md" if self._format == "markdown" else f".{self._format}"
output_path = self._resolve_collision(self._output_dir / f"{stem}{suffix}")
content = self._render(result)
output_path.write_text(content, encoding="utf-8")
logger.debug("write | path=%s | chars=%d", output_path, len(content))
return output_path
def write_batch_report(self, report, output_path: str | Path) -> None:
"""Write a BatchReport summary as JSON to *output_path*."""
path = Path(output_path)
data = {
"summary": {
"total": report.total,
"succeeded": report.succeeded,
"failed": report.failed,
"success_rate_pct": round(report.success_rate, 2),
"total_chars": report.total_chars,
"total_words": report.total_words,
"total_duration_ms": round(report.total_duration_ms, 2),
},
"results": [
{
"source": r.source,
"char_count": r.char_count,
"word_count": r.word_count,
"line_count": r.line_count,
"duration_ms": round(r.duration_ms, 2),
"content_hash": r.content_hash,
"mime_type": r.mime_type,
}
for r in report.results
],
"errors": [
{
"source": e.source,
"error_type": e.error_type,
"message": e.message,
"duration_ms": round(e.duration_ms, 2),
}
for e in report.errors
],
}
path.write_text(json.dumps(data, indent=2), encoding="utf-8")
logger.debug("write_batch_report | path=%s | total=%d", path, report.total)
def _render(self, result: ConversionResult) -> str:
"""Serialise *result* to the configured output format."""
if self._format == "json":
return json.dumps(
{
"source": result.source,
"markdown": result.markdown,
"meta": {
"char_count": result.char_count,
"word_count": result.word_count,
"line_count": result.line_count,
"duration_ms": round(result.duration_ms, 2),
"mime_type": result.mime_type,
"content_hash": result.content_hash,
},
},
indent=2,
)
return result.markdown
@staticmethod
def _resolve_collision(path: Path) -> Path:
"""Return *path* if it does not exist, otherwise append a counter to the stem."""
if not path.exists():
return path
counter = 1
while True:
candidate = path.with_stem(f"{path.stem}_{counter}")
if not candidate.exists():
return candidate
counter += 1