Skills › Scientific Agent Skills › liteparse › batch_parse_dir.py
batch_parse_dir.py
189 lines · skills/liteparse/scripts/batch_parse_dir.py at 5f2bc03
1#!/usr/bin/env python3
2"""
3Batch-parse documents with LiteParse; missing OCR language data may be downloaded.
5Usage:
6 python batch_parse_dir.py INPUT_DIR OUTPUT_DIR [--format json|text|markdown] [--no-ocr] [--recursive] [--extension .pdf]
7"""
9from __future__ import annotations
11import argparse
12import json
13import sys
14from pathlib import Path
15from typing import Iterable, List, Optional
17from liteparse import LiteParse, ParseResult
19DEFAULT_EXTENSIONS = {
20 ".pdf",
21 ".doc",
22 ".docx",
23 ".docm",
24 ".odt",
25 ".rtf",
26 ".ppt",
27 ".pptx",
28 ".pptm",
29 ".odp",
30 ".xls",
31 ".xlsx",
32 ".xlsm",
33 ".ods",
34 ".csv",
35 ".tsv",
36 ".png",
37 ".jpg",
38 ".jpeg",
39 ".gif",
40 ".bmp",
41 ".tiff",
42 ".tif",
43 ".webp",
44 ".svg",
45}
48def _text_item_dict(item) -> dict:
49 return {
50 "text": item.text,
51 "x": item.x,
52 "y": item.y,
53 "width": item.width,
54 "height": item.height,
55 "font_name": item.font_name,
56 "font_size": item.font_size,
57 "confidence": item.confidence,
58 }
61def _result_to_dict(result: ParseResult) -> dict:
62 return {
63 "text": result.text,
64 "total_pages": result.total_pages,
65 "page_errors": [
66 {"page_number": e.page_number, "message": e.message}
67 for e in result.page_errors
68 ],
69 "pages": [
70 {
71 "page_num": p.page_num,
72 "width": p.width,
73 "height": p.height,
74 "text": p.text,
75 "markdown": p.markdown,
76 "page_label": p.page_label,
77 "text_items": [_text_item_dict(i) for i in p.text_items],
78 }
79 for p in result.pages
80 ],
81 }
84def iter_files(
85 input_dir: Path,
86 *,
87 recursive: bool,
88 extension: Optional[str],
89) -> Iterable[Path]:
90 ext_filter = {extension.lower()} if extension else DEFAULT_EXTENSIONS
91 pattern = "**/*" if recursive else "*"
92 for path in sorted(input_dir.glob(pattern)):
93 if path.is_file() and path.suffix.lower() in ext_filter:
94 yield path
97def parse_one(
98 parser: LiteParse,
99 file_path: Path,
100 output_dir: Path,
101 fmt: str,
102) -> tuple[bool, str, str]:
103 try:
104 extensions = {"json": "json", "text": "txt", "markdown": "md"}
105 if fmt not in extensions:
106 raise ValueError(f"Unsupported output format: {fmt}")
107 # Retain the source suffix so report.pdf and report.docx cannot collide.
108 out_name = f"{file_path.name}.{extensions[fmt]}"
109 out_path = output_dir / out_name
110 if out_path.exists():
111 raise FileExistsError(f"Output already exists: {out_path}")
112 result = parser.parse(file_path)
113 if result.page_errors:
114 failures = "; ".join(
115 f"page {error.page_number}: {error.message}" for error in result.page_errors
116 )
117 raise ValueError(f"Incomplete extraction: {failures}")
118 if len(result.pages) != result.total_pages:
119 raise ValueError(
120 f"Incomplete extraction: returned {len(result.pages)} of "
121 f"{result.total_pages} source pages; check the parser page cap"
122 )
123 if fmt == "json":
124 content = json.dumps(_result_to_dict(result), indent=2, ensure_ascii=False)
125 else:
126 content = result.text
127 output_dir.mkdir(parents=True, exist_ok=True)
128 # Exclusive creation also protects against collisions between concurrent runs.
129 with out_path.open("x", encoding="utf-8") as output:
130 output.write(content)
131 return True, str(file_path), f"OK -> {out_name}"
132 except Exception as exc:
133 return False, str(file_path), str(exc)
136def main(argv: Optional[List[str]] = None) -> int:
137 p = argparse.ArgumentParser(description="Batch-parse documents with LiteParse")
138 p.add_argument("input_dir", type=Path, help="Directory of input files")
139 p.add_argument("output_dir", type=Path, help="Directory for parsed output")
140 p.add_argument(
141 "--format",
142 choices=("json", "text", "markdown"),
143 default="text",
144 help="Output format (default: text)",
145 )
146 p.add_argument("--no-ocr", action="store_true", help="Disable OCR")
147 p.add_argument("--recursive", action="store_true", help="Search subdirectories")
148 p.add_argument(
149 "--extension",
150 help="Only process this extension (e.g. .pdf); default: common document/image types",
151 )
152 p.add_argument("-q", "--quiet", action="store_true", help="Less console output")
153 args = p.parse_args(argv)
155 if not args.input_dir.is_dir():
156 print(f"Input directory not found: {args.input_dir}", file=sys.stderr)
157 return 1
159 args.output_dir.mkdir(parents=True, exist_ok=True)
161 parser = LiteParse(
162 ocr_enabled=not args.no_ocr,
163 output_format=args.format,
164 quiet=args.quiet,
165 )
167 files = list(iter_files(args.input_dir, recursive=args.recursive, extension=args.extension))
168 if not files:
169 print("No matching files found.", file=sys.stderr)
170 return 1
172 ok, fail = 0, 0
173 for fp in files:
174 destination = args.output_dir / fp.relative_to(args.input_dir).parent
175 success, path, msg = parse_one(parser, fp, destination, args.format)
176 if success:
177 ok += 1
178 if not args.quiet:
179 print(f"[OK] {path}: {msg}")
180 else:
181 fail += 1
182 print(f"[FAIL] {path}: {msg}", file=sys.stderr)
184 print(f"Done: {ok} succeeded, {fail} failed, {len(files)} total")
185 return 0 if fail == 0 else 1
188if __name__ == "__main__":
189 sys.exit(main())