Coverage for src/lilbee/data/analyze.py: 100%
153 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-11-06 08:17 +0000
1"""Corpus analysis: file types over every file, and what native extraction reads from a sample."""
3from __future__ import annotations
5import asyncio
6import statistics
7from collections import Counter
8from collections.abc import Mapping, Sequence
9from dataclasses import dataclass
10from pathlib import Path
11from typing import TYPE_CHECKING
13from lilbee.core.config import active_config
14from lilbee.data.extract.document import analysis_config
15from lilbee.data.extract.xberg import BatchItem, aextract_batch
16from lilbee.data.ingest.discovery import archive_content_types, classify_file
17from lilbee.data.types import (
18 CODE_CONTENT_TYPE,
19 IMAGE_CONTENT_TYPE,
20 PDF_CONTENT_TYPE,
21 OcrBackendName,
22)
23from lilbee.runtime.cancellation import CancelSignal, TaskCancelledError
24from lilbee.runtime.progress import (
25 AnalyzeEvent,
26 DetailedProgressCallback,
27 EventType,
28 noop_callback,
29)
31if TYPE_CHECKING:
32 from xberg import ExtractedDocument
34# How often a running batch checks its cancel signal.
35_CANCEL_POLL_S = 0.1
38@dataclass(frozen=True)
39class FileFailure:
40 """A sampled file analyze could not read, and why."""
42 file: str
43 error: str
46@dataclass(frozen=True)
47class PdfSignals:
48 """PDF pages, scans and tables; ``scanned_share`` also counts image files at the sample rate."""
50 files: int
51 pages: int
52 scanned_pages: int
53 scanned_share: float
54 files_with_tables: int
55 tables: int
56 median_pages: float | None
59@dataclass(frozen=True)
60class LanguageShare:
61 """A detected language (ISO 639-3) and its share of the files with a text layer."""
63 code: str
64 share: float
67@dataclass(frozen=True)
68class CorpusSignals:
69 """What analyze read: type mix and code share over every file, the rest over the sample.
71 ``files_read`` counts the documents extracted without error; code, image and archive
72 files are counted, never read.
73 """
75 files_total: int
76 documents_total: int
77 files_read: int
78 cap: int
79 failed: tuple[FileFailure, ...]
80 file_types: Mapping[str, int]
81 code_share: float
82 pdf: PdfSignals
83 median_chars: float | None
84 languages: tuple[LanguageShare, ...]
85 image_files: int
86 unreadable_folders: tuple[str, ...] = ()
88 @property
89 def files_counted(self) -> int:
90 """The code, image and archive files counted without being read."""
91 return self.files_total - self.documents_total
94@dataclass(frozen=True)
95class _Reading:
96 """What one extracted file contributes to the signals."""
98 content_type: str
99 pdf_pages: int
100 scanned_pages: int
101 tables: int
102 chars: int
103 language: str | None
106class _ExtractedView:
107 """The fields analyze reads from an xberg ``ExtractedDocument``."""
109 def __init__(self, doc: ExtractedDocument) -> None:
110 self._doc = doc
112 def reading(self, content_type: str) -> _Reading:
113 """The signals of this document, which discovery classified as *content_type*."""
114 fmt = self._doc.metadata.format
115 pdf = None if fmt is None else fmt.pdf
116 pages = 0 if pdf is None else pdf.page_count or self._doc.counts.pages
117 scanned = 0 if pdf is None else len(pdf.scanned_pages or [])
118 languages = self._doc.detected_languages
119 return _Reading(
120 content_type=content_type,
121 pdf_pages=pages,
122 scanned_pages=scanned,
123 tables=self._doc.counts.tables,
124 chars=len(self._doc.content),
125 language=languages[0] if languages else None,
126 )
129def sample_keys(keys: Sequence[str], cap: int) -> list[str]:
130 """The sorted keys, thinned to *cap* at an even stride when there are more."""
131 ordered = sorted(keys)
132 if len(ordered) <= cap:
133 return ordered
134 return [ordered[i * len(ordered) // cap] for i in range(cap)]
137def _counted_only() -> frozenset[str]:
138 """Types analyze counts without extracting: code, images, and archives of other files."""
139 return frozenset({CODE_CONTENT_TYPE, IMAGE_CONTENT_TYPE, *archive_content_types()})
142def _classify(files: Mapping[str, Path]) -> dict[str, str]:
143 """Each file's type; a file lilbee cannot ingest is left out."""
144 return {key: kind for key, path in files.items() if (kind := classify_file(path))}
147def _read_all(paths: list[Path]) -> list[bytes | OSError]:
148 results: list[bytes | OSError] = []
149 for path in paths:
150 try:
151 results.append(path.read_bytes())
152 except OSError as exc:
153 results.append(exc)
154 return results
157async def _extract_batch(
158 keys: list[str], files: Mapping[str, Path], types: Mapping[str, str]
159) -> tuple[list[_Reading], list[FileFailure]]:
160 """Read and extract one batch; a file that cannot be read or extracted becomes a failure."""
161 contents = await asyncio.to_thread(_read_all, [files[key] for key in keys])
162 # a file that could not be read is carried as its OSError
163 failures = [
164 FileFailure(key, str(data))
165 for key, data in zip(keys, contents, strict=True)
166 if isinstance(data, OSError)
167 ]
168 readable = [
169 (key, data) for key, data in zip(keys, contents, strict=True) if isinstance(data, bytes)
170 ]
171 items = [BatchItem(data, None, files[key].name, None) for key, data in readable]
172 docs = await aextract_batch(items, analysis_config()) if items else []
173 readings: list[_Reading] = []
174 for (key, _), doc in zip(readable, docs, strict=True):
175 # aextract_batch returns each failed file's error in its slot
176 if isinstance(doc, Exception):
177 failures.append(FileFailure(key, str(doc)))
178 else:
179 readings.append(_ExtractedView(doc).reading(types[key]))
180 return readings, failures
183def _share(part: float, whole: float) -> float:
184 return part / whole if whole else 0.0
187def _median(values: list[int]) -> float | None:
188 return float(statistics.median(values)) if values else None
191def _pdf_signals(readings: list[_Reading], image_weight: float) -> PdfSignals:
192 """PDF signals over the read sample; *image_weight* is the image files at the sample rate."""
193 pdfs = [r for r in readings if r.content_type == PDF_CONTENT_TYPE]
194 pages = sum(r.pdf_pages for r in pdfs)
195 scanned = sum(r.scanned_pages for r in pdfs)
196 return PdfSignals(
197 files=len(pdfs),
198 pages=pages,
199 scanned_pages=scanned,
200 scanned_share=_share(scanned + image_weight, pages + image_weight),
201 files_with_tables=sum(1 for r in pdfs if r.tables > 0),
202 tables=sum(r.tables for r in pdfs),
203 median_pages=_median([r.pdf_pages for r in pdfs]),
204 )
207def _language_shares(readings: list[_Reading]) -> tuple[LanguageShare, ...]:
208 counts = Counter(r.language for r in readings if r.language is not None)
209 total = sum(counts.values())
210 ranked = sorted(counts.items(), key=lambda item: (-item[1], item[0]))
211 return tuple(LanguageShare(code, count / total) for code, count in ranked)
214def _signals(
215 types: Mapping[str, str],
216 documents: list[str],
217 sample: list[str],
218 cap: int,
219 readings: list[_Reading],
220 failures: list[FileFailure],
221) -> CorpusSignals:
222 type_counts = Counter(types.values())
223 image_files = type_counts[IMAGE_CONTENT_TYPE]
224 # images are never read, so they count at the rate the documents were sampled
225 sample_rate = _share(len(sample), len(documents)) if documents else 1.0
226 other_chars = [r.chars for r in readings if r.content_type != PDF_CONTENT_TYPE]
227 return CorpusSignals(
228 files_total=len(types),
229 documents_total=len(documents),
230 files_read=len(readings),
231 cap=cap,
232 failed=tuple(failures),
233 file_types=dict(type_counts.most_common()),
234 code_share=_share(type_counts[CODE_CONTENT_TYPE], len(types)),
235 pdf=_pdf_signals(readings, image_files * sample_rate),
236 median_chars=_median(other_chars),
237 languages=_language_shares(readings),
238 image_files=image_files,
239 )
242def _stop_if_cancelled(cancel: CancelSignal | None) -> None:
243 if cancel is not None and cancel.is_set():
244 raise TaskCancelledError
247async def _extract_unless_cancelled(
248 batch: list[str],
249 files: Mapping[str, Path],
250 types: Mapping[str, str],
251 cancel: CancelSignal | None,
252) -> tuple[list[_Reading], list[FileFailure]]:
253 """Run one batch, dropping it as soon as *cancel* is set.
255 The token has no wait of its own, so the batch is polled against it, as a fleet chat is.
256 The extract call is native: dropping the batch ends the wait, not the work already running.
257 """
258 work = asyncio.ensure_future(_extract_batch(batch, files, types))
259 try:
260 while not work.done():
261 _stop_if_cancelled(cancel)
262 await asyncio.wait({work}, timeout=_CANCEL_POLL_S)
263 return work.result()
264 finally:
265 work.cancel()
266 await asyncio.gather(work, return_exceptions=True)
269async def collect_signals(
270 files: Mapping[str, Path],
271 *,
272 on_progress: DetailedProgressCallback = noop_callback,
273 cancel: CancelSignal | None = None,
274) -> CorpusSignals:
275 """Classify every file, then extract a sample of the documents natively with OCR off.
277 Raises ``TaskCancelledError`` when *cancel* is set, also while a batch is extracting.
278 """
279 config = active_config()
280 cap = config.analyze_max_files
281 types = await asyncio.to_thread(_classify, files)
282 counted_only = _counted_only()
283 documents = [key for key, kind in types.items() if kind not in counted_only]
284 to_extract = sample_keys(documents, cap)
285 size = config.batch_extraction_size
286 readings: list[_Reading] = []
287 failures: list[FileFailure] = []
288 for start in range(0, len(to_extract), size):
289 _stop_if_cancelled(cancel)
290 batch = to_extract[start : start + size]
291 batch_readings, batch_failures = await _extract_unless_cancelled(
292 batch, files, types, cancel
293 )
294 readings += batch_readings
295 failures += batch_failures
296 done = start + len(batch)
297 event = AnalyzeEvent(done=done, total=len(to_extract), file=batch[-1])
298 on_progress(EventType.ANALYZE, event)
299 _stop_if_cancelled(cancel)
300 return _signals(types, documents, to_extract, cap, readings, failures)
303def ocr_language_supported(code: str) -> bool:
304 """Whether the Tesseract backend has the language *code* installed."""
305 from xberg import ocr_backend_supports_language
307 return ocr_backend_supports_language(OcrBackendName.TESSERACT, code)