diff --git a/.agents/skills/review/SKILL.md b/.agents/skills/review/SKILL.md new file mode 100644 index 0000000000..2d42bd30a2 --- /dev/null +++ b/.agents/skills/review/SKILL.md @@ -0,0 +1,77 @@ +--- +name: review +description: Review or re-review a Docling pull request with reproducible findings and explicit validation results. Use for code, tests, dependencies, documentation, or agent guidance. +--- + +# Review a Docling pull request + +Use these steps for every review. Human reviewers can use the same steps. +Explain the affected contract so that a reviewer need not know Docling internals. +Read [AGENTS.md](../../../AGENTS.md) and load other applicable +[task routes](../skill-router.json). For document conversion without repository +changes, use the [package usage skill](../../../docling/.agents/skills/docling/SKILL.md). + +## 1. Fix the review scope + +Read the live PR description, diff, reviews, and CI results. Record the base and +head commit IDs. Use an isolated worktree; preserve other changes and staging. +For a re-review, read each open finding and the author's response. + +**Exit:** The changed files and affected contracts are known. The reviewed head +is exact, and prior findings are listed for verification. + +## 2. Check behavior and tests + +Follow input through the affected code to the user-visible result. + +- For conversion changes, check source content, `DoclingDocument` ownership and + order, and the affected exports. Check nested and non-text content where relevant. + Text presence or item counts alone do not prove correct structure. +- For a bug fix, run a small case on base and head when practical. The regression + test must fail for the original defect and pass with the fix. State any limit. +- Check public types, Python 3.10 support, defaults, error paths, and optional + dependencies where affected. Prefer explicit contracts to attribute probing. +- Inspect test assertions and changed reference data. Confirm that they detect + the defect rather than accept the new output without a reason. +- For loops over document content, check how work grows with input size. Compare + base and head on the same input when needed. Do not use noisy timing thresholds + as a correctness gate. + +**Exit:** Each affected contract has evidence, or an explicit unchecked limit. +Every prior finding is confirmed fixed or still reproducible on the new head. + +## 3. Run the applicable checks + +Use the commands in `AGENTS.md`. For changes you make, run `make validate`, inspect +hook edits, and repeat until it passes. For a read-only review, use `make check` +and targeted tests. Do not regenerate reference data during a read-only review. +Inspect CI on the exact reviewed head and explain failures that affect the verdict. + +`python3 .github/scripts/check_skill_routes.py` checks contributor skill routes +and Codex/Claude links. Existing Ruff, ty, Tach, and lock checks check code and +dependency rules. These checks do not prove conversion quality, test adequacy, +or full ASD-STE100 compliance. Those decisions require review. + +**Exit:** Record commands, results, reviewed commit, and checks not run. Never +describe skipped, pending, or failed checks as passed. + +## 4. Give a precise verdict + +Use one finding per defect. Include its location, trigger, actual and expected +behavior, user impact, and reproduction or test evidence. Link the affected code. +Use these severity terms consistently: + +| Severity | Meaning | +| --- | --- | +| Blocker | Security defect, data loss, incorrect supported behavior, or a required check failure. Request changes. | +| Suggestion | A useful improvement with no demonstrated contract failure. Do not block on preference. | +| Question | Evidence is missing. Ask for the specific evidence; do not assert a defect. | + +Approve only when no blocker remains and relevant validation is complete. +If required evidence is unavailable, state the limit and withhold approval. +Post reviews or comments only when the user has authorized that action. Before +posting, confirm that the PR head still matches the reviewed commit. If it changed, +check the new diff and repeat affected validation. Avoid duplicate findings. + +**Exit:** The verdict identifies the reviewed head, blockers, validation results, +and limits. Follow the communication rule in `AGENTS.md`. diff --git a/.agents/skills/skill-router.json b/.agents/skills/skill-router.json new file mode 100644 index 0000000000..eac07ad9ca --- /dev/null +++ b/.agents/skills/skill-router.json @@ -0,0 +1,5 @@ +{ + "review": "Review or re-review a pull request. Start here, then load other applicable skills.", + "dignified-python": "Write, review, or change Python code. Repository rules take priority.", + "building-pydantic-ai-agents": "Build or review code that uses Pydantic AI." +} diff --git a/.claude/skills/review b/.claude/skills/review new file mode 120000 index 0000000000..d19862ae43 --- /dev/null +++ b/.claude/skills/review @@ -0,0 +1 @@ +../../.agents/skills/review \ No newline at end of file diff --git a/.codex/skills/review b/.codex/skills/review new file mode 120000 index 0000000000..d19862ae43 --- /dev/null +++ b/.codex/skills/review @@ -0,0 +1 @@ +../../.agents/skills/review \ No newline at end of file diff --git a/.github/PULL_REQUEST_TEMPLATE.md b/.github/PULL_REQUEST_TEMPLATE.md index 58a6bf0afd..a1a26acf1f 100644 --- a/.github/PULL_REQUEST_TEMPLATE.md +++ b/.github/PULL_REQUEST_TEMPLATE.md @@ -17,3 +17,9 @@ Resolves # - [ ] Documentation has been updated, if necessary. - [ ] Examples have been added, if necessary. - [ ] Tests have been added, if necessary. + +**Review evidence:** + + diff --git a/.github/scripts/check_skill_routes.py b/.github/scripts/check_skill_routes.py new file mode 100644 index 0000000000..254dd9d3f8 --- /dev/null +++ b/.github/scripts/check_skill_routes.py @@ -0,0 +1,64 @@ +# SPDX-FileCopyrightText: The Docling Contributors +# SPDX-License-Identifier: MIT + +from __future__ import annotations + +import json +import re +import sys +from pathlib import Path + + +def check_routes(repo_root: Path) -> list[str]: + skills_root = repo_root / ".agents/skills" + router = skills_root / "skill-router.json" + try: + routes = json.loads(router.read_text(encoding="utf-8")) + except (OSError, ValueError) as error: + return [f"Cannot read {router}: {error}"] + + if not isinstance(routes, dict) or not routes: + return ["Skill routes must be a non-empty object of skill names and tasks."] + if any( + not re.fullmatch(r"[a-z0-9]+(?:-[a-z0-9]+)*", name) + or not isinstance(task, str) + or not task.strip() + for name, task in routes.items() + ): + return ["Each route needs a skill directory name and a non-empty task."] + + errors: list[str] = [] + discovered = {path.name for path in skills_root.iterdir() if path.is_dir()} + for name in sorted(discovered - routes.keys()): + errors.append(f"Contributor skill has no task route: {name}") + for name in sorted(routes): + skill = skills_root / name + if not (skill / "SKILL.md").is_file(): + errors.append(f"Route has no SKILL.md: {name}") + for agent in ("codex", "claude"): + link = repo_root / f".{agent}/skills" / name + if not link.is_symlink() or link.resolve() != skill.resolve(): + errors.append(f".{agent}/skills/{name} must link to {skill}") + + for agent in ("codex", "claude"): + links = repo_root / f".{agent}/skills" + if links.is_dir(): + for link in sorted(links.iterdir()): + if link.name not in routes: + errors.append( + f"Agent skill has no task route: .{agent}/skills/{link.name}" + ) + return errors + + +def main() -> int: + errors = check_routes(Path(__file__).resolve().parents[2]) + if errors: + print("\n".join(errors), file=sys.stderr) + return 1 + print("Contributor skill routes and agent links are valid.") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 87c4e312d4..52860630b7 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -86,6 +86,10 @@ jobs: with: filters: | run_ci: + - "AGENTS.md" + - ".agents/skills/**" + - ".codex/skills/**" + - ".claude/skills/**" - ".github/workflows/**" - ".github/scripts/**" - ".github/actions/**" diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 0fe5de49f5..fc3e1e202f 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -37,6 +37,12 @@ repos: files: '^(docling|tests|docs/examples|perfs).*\.(py|ipynb)$' - repo: local hooks: + - id: skill-routes + name: Skill routing and link integrity + entry: python3 .github/scripts/check_skill_routes.py + pass_filenames: false + language: system + files: '^(AGENTS\.md|\.agents/skills/|\.(codex|claude)/skills/|\.github/scripts/check_skill_routes\.py|tests/test_check_skill_routes\.py)' - id: ty name: ty entry: uv run --no-sync ty check diff --git a/AGENTS.md b/AGENTS.md index 8d38d91267..7139365478 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -32,6 +32,9 @@ scripts/ # project maintenance scripts - **Development skills** (for contributors working *on* Docling) live in [`.agents/skills/`](.agents/skills/) at the repo root, e.g. `dignified-python` and `building-pydantic-ai-agents`. + Select the applicable skills from the + [task routes](.agents/skills/skill-router.json). For PR reviews and re-reviews, + use [review](.agents/skills/review/SKILL.md) before you assess the changes. - **Usage skills** (for agents *using* Docling to convert documents) are shipped inside the package at [`docling/.agents/skills/docling/`](docling/.agents/skills/docling/SKILL.md). diff --git a/Makefile b/Makefile index ea28958608..39f160be4f 100644 --- a/Makefile +++ b/Makefile @@ -18,6 +18,7 @@ check-all: ## Run all read-only local checks. uv run --no-sync tach check python3 scripts/check_tach_module_coverage.py python3 scripts/check_max_lines.py + python3 .github/scripts/check_skill_routes.py uv run --no-sync dprint check --config .github/dprint.json --config-discovery=false uv lock --locked diff --git a/docling/backend/afp_backend.py b/docling/backend/afp_backend.py index f714d70789..f4a4311e9c 100644 --- a/docling/backend/afp_backend.py +++ b/docling/backend/afp_backend.py @@ -33,6 +33,7 @@ DeclarativeDocumentBackend, PaginatedDocumentBackend, ) +from docling.datamodel.backend_options import BackendOptions from docling.datamodel.base_models import InputFormat from docling.datamodel.document import InputDocument from docling.exceptions import DocumentLoadError @@ -216,8 +217,9 @@ def __init__( self, in_doc: InputDocument, path_or_stream: BytesIO | Path, + options: BackendOptions | None = None, ) -> None: - super().__init__(in_doc, path_or_stream) + super().__init__(in_doc, path_or_stream, options) self.page_range = in_doc.limits.page_range try: self.content = ( diff --git a/docling/backend/asciidoc_backend.py b/docling/backend/asciidoc_backend.py index 56862f7c36..3b2009f289 100644 --- a/docling/backend/asciidoc_backend.py +++ b/docling/backend/asciidoc_backend.py @@ -7,7 +7,7 @@ from dataclasses import dataclass from io import BytesIO from pathlib import Path -from typing import Final, Optional, Union +from typing import Final, Optional, TypedDict, Union from docling_core.types.doc import ( DocItemLabel, @@ -23,6 +23,7 @@ TableData, TextItem, ) +from typing_extensions import override from docling.backend.abstract_backend import DeclarativeDocumentBackend from docling.backend.utils.image_resource_loader import ImageResourceLoader @@ -44,6 +45,15 @@ _LIST_ITEM_PATTERN: Final = r"^(\s*)(\*+|-|\.+|\d+\.|\w+\.)\s+(.*)" +class _ParsedBlock(TypedDict, total=False): + type: str + text: str + level: int + indent: int + marker: str + numbered: bool + + @dataclass(frozen=True) class _LiteralBlock: text: str @@ -84,20 +94,25 @@ def __init__( ) from e return + @override def is_valid(self) -> bool: return self.valid @classmethod + @override def supports_pagination(cls) -> bool: return False + @override def unload(self): return @classmethod + @override def supported_formats(cls) -> set[InputFormat]: return {InputFormat.ASCIIDOC} + @override def convert(self) -> DoclingDocument: """ Parses the ASCII into a structured document model. @@ -115,7 +130,7 @@ def convert(self) -> DoclingDocument: return doc - def _parse(self, doc: DoclingDocument): + def _parse(self, doc: DoclingDocument) -> DoclingDocument: """Orchestrate parsing and populate `doc` from the source lines. Handles titles, section headers, text paragraphs, lists, tables, @@ -145,9 +160,9 @@ def _parse(self, doc: DoclingDocument): list_continuation = False # parents: dict[int, Union[DocItem, GroupItem, None]] = {} - parents: dict[int, Union[GroupItem, None]] = {} + parents: dict[int, NodeItem | None] = {} # indents: dict[int, Union[DocItem, GroupItem, None]] = {} - indents: dict[int, Union[GroupItem, None]] = {} + indents: dict[int, int | None] = {} for i in range(10): parents[i] = None @@ -171,10 +186,12 @@ def _parse(self, doc: DoclingDocument): ) caption: Optional[TextItem] = None if caption_data: - caption = doc.add_text( + caption_item = doc.add_text( text=" ".join(caption_data), label=DocItemLabel.CAPTION, ) + assert isinstance(caption_item, TextItem) + caption = caption_item caption_data = [] code_language = ( detect_code_language(block.text, hint=block.language) @@ -244,6 +261,7 @@ def _parse(self, doc: DoclingDocument): level = self._get_current_level(parents) + current_indent = indents[level] if not in_list: in_list = True # GroupItem has no caption slot; emit the pending block @@ -262,14 +280,18 @@ def _parse(self, doc: DoclingDocument): ) indents[level + 1] = item["indent"] - elif in_list and item["indent"] > indents[level]: + elif current_indent is not None and item["indent"] > current_indent: parents[level + 1] = doc.add_group( parent=parents[level], name="list", label=GroupLabel.LIST ) indents[level + 1] = item["indent"] - elif in_list and item["indent"] < indents[level]: - while level > 0 and item["indent"] < indents[level]: + elif current_indent is not None and item["indent"] < current_indent: + while ( + level > 0 + and current_indent is not None + and item["indent"] < current_indent + ): # Only pop the current level if there is an outer group # to fall back to; otherwise keep it as the list root. if indents[level - 1] is None: @@ -277,13 +299,16 @@ def _parse(self, doc: DoclingDocument): parents[level] = None indents[level] = None level -= 1 + current_indent = indents[level] - last_list_item = doc.add_list_item( + list_item = doc.add_list_item( item["text"], enumerated=item["numbered"], marker=(item["marker"] if item["marker"][:-1].isdigit() else None), parent=self._get_current_parent(parents), ) + assert isinstance(list_item, ListItem) + last_list_item = list_item list_continuation = False elif in_list and stripped_line == "+": @@ -377,7 +402,7 @@ def _parse(self, doc: DoclingDocument): return doc @staticmethod - def _get_current_level(parents): + def _get_current_level(parents: dict[int, NodeItem | None]) -> int: for k, v in parents.items(): if v is None and k > 0: return k - 1 @@ -385,14 +410,16 @@ def _get_current_level(parents): return 0 @staticmethod - def _get_current_parent(parents): + def _get_current_parent(parents: dict[int, NodeItem | None]) -> NodeItem | None: for k, v in parents.items(): if v is None and k > 0: return parents[k - 1] return None - _SOURCE_ATTR_RE = re.compile(r"^\[source(?:,\s*([\w+#.-]+))?[^\]]*\]$") + _SOURCE_ATTR_RE: re.Pattern[str] = re.compile( + r"^\[source(?:,\s*([\w+#.-]+))?[^\]]*\]$" + ) _CONTENT_BLOCK_DELIMITERS = ("====", "****", "____", "--", "+++") @staticmethod @@ -458,7 +485,9 @@ def _iter_blocks(cls, lines: list[str]) -> Iterator[str | _LiteralBlock]: if source_attr is not None and i + 1 < n and lines[i + 1].strip() == "----": block_data = [] block_delimiter = "----" - block_language = source_attr.group(1) + language = source_attr.group(1) + assert language is None or isinstance(language, str) + block_language = language i += 2 continue @@ -474,7 +503,7 @@ def _close_list_if_needed( *, line: str, in_list: bool, - parents: dict[int, GroupItem | None], + parents: dict[int, NodeItem | None], last_list_item: ListItem | None, list_continuation: bool, is_continuation_block: bool, @@ -496,7 +525,7 @@ def _flush_text_data( *, doc: DoclingDocument, text_data: list[str], - parent: GroupItem | None, + parent: NodeItem | None, ) -> list[str]: if len(text_data) > 0: doc.add_text( @@ -508,22 +537,24 @@ def _flush_text_data( # ========= Title @staticmethod - def _is_title(line): + def _is_title(line: str) -> re.Match[str] | None: return re.match(r"^= ", line) @staticmethod - def _parse_title(line): + def _parse_title(line: str) -> _ParsedBlock: return {"type": "title", "text": line[2:].strip(), "level": 0} # ========= Section headers @staticmethod - def _is_section_header(line): + def _is_section_header(line: str) -> re.Match[str] | None: return re.match(r"^==+\s+", line) @staticmethod - def _parse_section_header(line): + def _parse_section_header(line: str) -> _ParsedBlock: match = re.match(r"^(=+)\s+(.*)", line) + if match is None: + raise ValueError(f"Invalid AsciiDoc section header: {line!r}") marker = match.group(1) # The list marker (e.g., "*", "-", "1.") text = match.group(2) # The actual text of the list item @@ -536,11 +567,11 @@ def _parse_section_header(line): # ========= Lists @staticmethod - def _is_list_item(line): + def _is_list_item(line: str) -> re.Match[str] | None: return re.match(_LIST_ITEM_PATTERN, line) @staticmethod - def _parse_list_item(line): + def _parse_list_item(line: str) -> _ParsedBlock: """Extract the item marker (number or bullet symbol) and the text of the item.""" match = re.match(_LIST_ITEM_PATTERN, line) @@ -580,11 +611,11 @@ def _parse_list_item(line): # ========= Tables @staticmethod - def _is_table_line(line): + def _is_table_line(line: str) -> re.Match[str] | None: return re.match(rf"^{_CELL_SPEC}\|.*\|", line) @staticmethod - def _parse_table_line(line): + def _parse_table_line(line: str) -> list[str]: # Drop cell specifiers glued to a "|" (e.g. "^.^h"); anchored to # whitespace so content ending in a style letter (e.g. "Eth") survives. line = re.sub(rf"(^|\s){_CELL_SPEC}(?=\|)", r"\1", line) @@ -613,7 +644,7 @@ def _add_table_if_nonempty( doc.add_table(data=data, parent=parent, caption=caption) @staticmethod - def _populate_table_as_grid(table_data): + def _populate_table_as_grid(table_data: list[list[str]]) -> TableData: num_rows = len(table_data) # Adjust the table data into a grid format @@ -645,19 +676,22 @@ def _populate_table_as_grid(table_data): # ========= Pictures @staticmethod - def _is_picture(line): + def _is_picture(line: str) -> re.Match[str] | None: return re.match(r"^image::", line) @staticmethod - def _parse_picture(line): + def _parse_picture(line: str) -> dict[str, str]: """ Parse an image macro, extracting its path and attributes. Syntax: image::path/to/image.png[Alt Text, width=200, height=150, align=center] """ mtch = re.match(r"^image::(.+)\[(.*)\]$", line) if mtch: - picture_path = mtch.group(1).strip() - attributes = mtch.group(2).split(",") + path = mtch.group(1) + attrs = mtch.group(2) + assert isinstance(path, str) and isinstance(attrs, str) + picture_path = path.strip() + attributes = attrs.split(",") picture_info = {"type": "picture", "uri": picture_path} # Extract optional attributes (alt text, width, height, alignment) @@ -677,11 +711,11 @@ def _parse_picture(line): # ========= Captions @staticmethod - def _is_caption(line): + def _is_caption(line: str) -> re.Match[str] | None: return re.match(r"^\.(\S.*)", line) @staticmethod - def _parse_caption(line): + def _parse_caption(line: str) -> _ParsedBlock: mtch = re.match(r"^\.(.+)", line) if mtch: text = mtch.group(1) @@ -691,5 +725,5 @@ def _parse_caption(line): # ========= Plain text @staticmethod - def _parse_text(line): + def _parse_text(line: str) -> _ParsedBlock: return {"type": "text", "text": line.strip()} diff --git a/docling/backend/boxnote_backend.py b/docling/backend/boxnote_backend.py index d0e5b412a4..38d32c8fb7 100644 --- a/docling/backend/boxnote_backend.py +++ b/docling/backend/boxnote_backend.py @@ -29,6 +29,7 @@ clamp_span, table_width, ) +from docling.datamodel.backend_options import BackendOptions from docling.datamodel.base_models import InputFormat from docling.datamodel.document import InputDocument from docling.exceptions import DocumentLoadError @@ -61,8 +62,13 @@ class BoxNoteDocumentBackend(DeclarativeDocumentBackend): """ @override - def __init__(self, in_doc: InputDocument, path_or_stream: BytesIO | Path): - super().__init__(in_doc, path_or_stream) + def __init__( + self, + in_doc: InputDocument, + path_or_stream: BytesIO | Path, + options: BackendOptions | None = None, + ): + super().__init__(in_doc, path_or_stream, options) self.data: dict[str, Any] = {} # utf-8-sig drops a leading BOM. Kept, it reaches json.loads, which @@ -122,7 +128,7 @@ def convert(self) -> DoclingDocument: return doc def _add_blocks( - self, nodes: list[dict], doc: DoclingDocument, parent: NodeItem | None + self, nodes: list[dict[str, Any]], doc: DoclingDocument, parent: NodeItem | None ) -> None: for node in nodes: self._add_block(node, doc, parent) @@ -170,7 +176,10 @@ def _add_block( self._add_blocks(content, doc, parent) def _add_paragraph( - self, content: list[dict], doc: DoclingDocument, parent: NodeItem | None + self, + content: list[dict[str, Any]], + doc: DoclingDocument, + parent: NodeItem | None, ) -> None: runs = self._runs(content) if not runs: @@ -198,7 +207,7 @@ def _add_paragraph( def _add_list( self, list_type: str, - items: list[dict], + items: list[dict[str, Any]], doc: DoclingDocument, parent: NodeItem | None, ) -> None: @@ -253,14 +262,14 @@ def _add_check_item( def _split_item( self, item: dict[str, Any] - ) -> tuple[str, Formatting | None, AnyUrl | None, list[dict]]: + ) -> tuple[str, Formatting | None, AnyUrl | None, list[dict[str, Any]]]: """Split a list item into its own text and any nested blocks. The leading paragraph becomes the item text; everything else (nested lists, extra paragraphs) is returned to be added as children. """ text, formatting, hyperlink = "", None, None - nested: list[dict] = [] + nested: list[dict[str, Any]] = [] for child in item.get("content", []): if not text and child.get("type") == "paragraph": text, formatting, hyperlink = self._collapse(child.get("content", [])) @@ -269,7 +278,7 @@ def _split_item( return text, formatting, hyperlink, nested def _add_table( - self, rows: list[dict], doc: DoclingDocument, parent: NodeItem | None + self, rows: list[dict[str, Any]], doc: DoclingDocument, parent: NodeItem | None ) -> None: rows = [row for row in rows if row.get("type") == "table_row"] if not rows: @@ -359,7 +368,7 @@ def _add_table( table.data.num_cols = num_cols - def _cell_is_rich(self, blocks: list[dict]) -> bool: + def _cell_is_rich(self, blocks: list[dict[str, Any]]) -> bool: """A cell is rich when it cannot be one plain-text TableCell. That covers images, multiple blocks, any non-paragraph block (lists, @@ -391,7 +400,7 @@ def _add_image( caption = doc.add_text(label=DocItemLabel.CAPTION, text=label) doc.add_picture(caption=caption, parent=parent) - def _collapse(self, content: list[dict]) -> _Run: + def _collapse(self, content: list[dict[str, Any]]) -> _Run: """Reduce inline content to a single run. Headings and list items take one formatting and one hyperlink, so a @@ -403,7 +412,7 @@ def _collapse(self, content: list[dict]) -> _Run: return runs[0] return "".join(text for text, _, _ in runs), None, None - def _runs(self, content: list[dict]) -> list[_Run]: + def _runs(self, content: list[dict[str, Any]]) -> list[_Run]: runs: list[_Run] = [] for node in content or []: node_type = node.get("type") @@ -416,7 +425,9 @@ def _runs(self, content: list[dict]) -> list[_Run]: runs.append((" ", None, None)) return runs - def _marks(self, marks: list[dict]) -> tuple[Formatting | None, AnyUrl | None]: + def _marks( + self, marks: list[dict[str, Any]] + ) -> tuple[Formatting | None, AnyUrl | None]: formatting: Formatting | None = None hyperlink: AnyUrl | None = None for mark in marks or []: @@ -439,7 +450,7 @@ def _marks(self, marks: list[dict]) -> tuple[Formatting | None, AnyUrl | None]: hyperlink = self._as_url(href) return formatting, hyperlink - def _plain_text(self, nodes: list[dict]) -> str: + def _plain_text(self, nodes: list[dict[str, Any]]) -> str: parts: list[str] = [] for node in nodes or []: node_type = node.get("type") @@ -451,7 +462,7 @@ def _plain_text(self, nodes: list[dict]) -> str: parts.append(self._plain_text(node["content"])) return "".join(parts) - def _cell_text(self, blocks: list[dict]) -> str: + def _cell_text(self, blocks: list[dict[str, Any]]) -> str: texts = (self._plain_text(block.get("content", [])).strip() for block in blocks) return " ".join(text for text in texts if text) diff --git a/docling/backend/csv_backend.py b/docling/backend/csv_backend.py index b7ab0ba490..ccb9e862cb 100644 --- a/docling/backend/csv_backend.py +++ b/docling/backend/csv_backend.py @@ -10,6 +10,7 @@ from typing import Final, Optional, Set, Union from docling_core.types.doc import DoclingDocument, DocumentOrigin, TableCell, TableData +from typing_extensions import override from docling.backend.abstract_backend import DeclarativeDocumentBackend from docling.datamodel.backend_options import CsvBackendOptions @@ -76,22 +77,27 @@ def __init__( ) from e return + @override def is_valid(self) -> bool: return self.valid @classmethod + @override def supports_pagination(cls) -> bool: return False + @override def unload(self): if isinstance(self.path_or_stream, BytesIO): self.path_or_stream.close() self.path_or_stream = None @classmethod + @override def supported_formats(cls) -> Set[InputFormat]: return {InputFormat.CSV} + @override def convert(self) -> DoclingDocument: """Parse the CSV content into a DoclingDocument. diff --git a/docling/backend/docling_parse_backend.py b/docling/backend/docling_parse_backend.py index 129fbc6618..f95b927eb2 100644 --- a/docling/backend/docling_parse_backend.py +++ b/docling/backend/docling_parse_backend.py @@ -17,6 +17,7 @@ TextCell, ) from PIL import Image +from typing_extensions import override from docling.backend.pdf_backend import PdfDocumentBackend, PdfPageBackend from docling.datamodel.accelerator_options import AcceleratorOptions @@ -130,15 +131,19 @@ def __init__(self, result: PageParseResult, rendered: bool = True): self._seg_page: Optional[SegmentedPdfPage] = None @property + @override def page_no(self) -> int: return self._result.page_number + @override def is_valid(self) -> bool: return self._result.success + @override def get_error_message(self) -> str: return self._result.error_message + @override def get_text_in_rect(self, bbox: BoundingBox) -> str: segmented_page = self.get_segmented_page() if segmented_page is None: @@ -155,6 +160,7 @@ def get_text_in_rect(self, bbox: BoundingBox) -> str: return text_piece + @override def get_segmented_page(self) -> Optional[SegmentedPdfPage]: if not self.is_valid(): return None @@ -170,22 +176,25 @@ def get_segmented_page(self) -> Optional[SegmentedPdfPage]: self._seg_page = seg_page return self._seg_page + @override def get_text_cells(self) -> Iterable[TextCell]: segmented_page = self.get_segmented_page() if segmented_page is None: - return [] + return list[TextCell]() return segmented_page.textline_cells + @override def get_visible_text_cells(self) -> Optional[list[TextCell]]: segmented_page = self.get_segmented_page() if segmented_page is None: return [] return _visible_text_cells(segmented_page.textline_cells) + @override def get_bitmap_rects(self, scale: float = 1) -> Iterable[BoundingBox]: segmented_page = self.get_segmented_page() if segmented_page is None: - return [] + return list[BoundingBox]() page_height = self.get_size().height cropboxes: list[BoundingBox] = [] @@ -197,6 +206,7 @@ def get_bitmap_rects(self, scale: float = 1) -> Iterable[BoundingBox]: cropboxes.append(cropbox.scaled(scale=scale)) return cropboxes + @override def has_content_in( self, *, @@ -211,6 +221,7 @@ def has_content_in( bbox=bbox, chars=chars, shapes=shapes, bitmaps=bitmaps ) + @override def get_shape_lines( self, *, @@ -229,6 +240,7 @@ def get_shape_lines( ) ] + @override def get_connected_shape_bounding_boxes( self, *, tolerance: float = 0.0 ) -> Optional[list[BoundingBox]]: @@ -243,6 +255,7 @@ def get_connected_shape_bounding_boxes( ) ] + @override def get_page_image( self, scale: float = 1, cropbox: Optional[BoundingBox] = None ) -> Image.Image: @@ -253,9 +266,11 @@ def get_page_image( ) return self._result.get_image(scale=scale, cropbox=cropbox).convert("RGB") + @override def get_size(self) -> Size: return Size(width=self._result.page_width, height=self._result.page_height) + @override def unload(self) -> None: return None @@ -328,7 +343,7 @@ def __init__( if isinstance(self.path_or_stream, BytesIO): self.path_or_stream.seek(0) self.doc_key = self.parser.load( - self.path_or_stream, + path_or_stream, password=password, page_range=in_doc.limits.page_range, ) @@ -344,29 +359,35 @@ def __init__( f"docling-parse could not load document {self.document_hash}." ) from e + @override def is_valid(self) -> bool: return not self._closed and self.page_count() > 0 + @override def page_count(self) -> int: return self.parser.page_count(self.doc_key) + @override def get_document_outline(self) -> list[_PdfOutlineItem]: """Extract the outline from the threaded parser's document annotations.""" annotations = self.parser.get_annotations(self.doc_key) toc = annotations.table_of_contents if annotations is not None else None return extract_outline_from_docling_parse(toc) + @override def load_page(self, page_no: int) -> PdfPageBackend: raise NotImplementedError( "ThreadedDoclingParseDocumentBackend only supports iter_pages()." ) + @override def iter_pages(self) -> Iterator[ThreadedDoclingParsePageBackend]: self._iterating = True for result in self.parser.iterate_results(): yield ThreadedDoclingParsePageBackend(result, rendered=self._render_pages) self._iterating = False + @override def unload(self) -> None: if self._closed: return diff --git a/docling/backend/docx/drawingml/utils.py b/docling/backend/docx/drawingml/utils.py index 1646552be6..707968d705 100644 --- a/docling/backend/docx/drawingml/utils.py +++ b/docling/backend/docx/drawingml/utils.py @@ -142,7 +142,7 @@ def _build_soffice_command( ] -def _kill_soffice_process_group(proc: subprocess.Popen) -> None: +def _kill_soffice_process_group(proc: subprocess.Popen[bytes]) -> None: """Best-effort SIGKILL of the whole ``soffice`` process group. ``soffice`` is a thin wrapper that forks ``soffice.bin``; killing only @@ -328,7 +328,7 @@ def convert_to_modern_format( shutil.rmtree(tmp_dir, ignore_errors=True) -def get_docx_to_pdf_converter() -> Optional[Callable]: +def get_docx_to_pdf_converter() -> Optional[Callable[[Path, Path], None]]: """ Detects the best available DOCX to PDF tool and returns a conversion function. The returned function accepts (input_path, output_path). @@ -411,7 +411,7 @@ def crop_whitespace(image: Image.Image, bg_color=None, padding=0) -> Image.Image def get_pil_from_dml_docx( - docx: Document, converter: Optional[Callable] + docx: Document, converter: Optional[Callable[[Path, Path], None]] ) -> Optional[Image.Image]: if converter is None: return None diff --git a/docling/backend/docx/latex/omml.py b/docling/backend/docx/latex/omml.py index 049cac8e41..880857d8f0 100644 --- a/docling/backend/docx/latex/omml.py +++ b/docling/backend/docx/latex/omml.py @@ -18,6 +18,7 @@ import lxml.etree as ET from lxml.etree import _Element +from typing_extensions import override try: # pragma: no cover - import-time guard from pylatexenc.latexencode import UnicodeToLatexEncoder @@ -109,7 +110,7 @@ def escape_latex(strs: str) -> str: return BLANK.join(new_chr) -def get_val(key: str | None, default: str, store: dict | None = CHR) -> str: +def get_val(key: str | None, default: str, store: dict[str, str] | None = CHR) -> str: """Get a value from a dictionary store or return the key/default. When default is provided, this function always returns a non-None string. @@ -145,7 +146,9 @@ def call_method(self, elm: _Element, stag: str | None = None) -> Any | None: """ getmethod = self.tag2meth.get if stag is None: - stag = elm.tag.replace(OMML_NS, "") + tag = elm.tag + assert isinstance(tag, str) + stag = str(tag).replace(OMML_NS, "") method = getmethod(stag) if method: return method(self, elm) @@ -231,7 +234,7 @@ class Pr(Tag2Method): text: str __val_tags: tuple[str, ...] - __innerdict: dict[str, Any] + __innerdict: dict[str, str | None] def __init__(self, elm: _Element): """Initialize properties processor. @@ -243,6 +246,7 @@ def __init__(self, elm: _Element): self.__innerdict = {} self.text = self.process_children(elm) + @override def __str__(self) -> str: """Return string representation.""" return self.text @@ -251,7 +255,7 @@ def __unicode__(self) -> str: """Return unicode representation.""" return self.__str__() - def __getattr__(self, name: str) -> Any | None: + def __getattr__(self, name: str) -> str | None: """Get attribute from internal dictionary. Args: @@ -318,6 +322,7 @@ def __init__(self, element: _Element): ) self._latex = self.process_children(element) + @override def __str__(self) -> str: """Return LaTeX string with normalized spacing.""" return self.latex.replace(" ", " ") @@ -326,6 +331,7 @@ def __unicode__(self) -> str: """Return unicode representation.""" return self.__str__() + @override def process_unknow(self, elm: _Element, stag: str) -> Any | None: """Handle unknown element types. @@ -383,6 +389,7 @@ def do_bar(self, elm: _Element) -> str: """ c_dict = self.process_children_dict(elm) pr = c_dict["barPr"] + assert isinstance(pr, Pr) latex_s = get_val(pr.pos, default=POS_DEFAULT["BAR_VAL"], store=POS) # If latex_s contains %s, format it; otherwise return as-is (unmapped character) if "%s" in latex_s: @@ -401,6 +408,7 @@ def do_d(self, elm: _Element) -> str: """ c_dict = self.process_children_dict(elm) pr = c_dict["dPr"] + assert isinstance(pr, Pr) null = D_DEFAULT["null"] s_val = get_val(pr.begChr, default=D_DEFAULT["left"], store=T) @@ -466,10 +474,11 @@ def do_ssub(self, elm: _Element) -> str: """ c_dict = self.process_children_dict(elm, include=("e", "sub", "sSubPr")) base = c_dict.get("e", "").rstrip() + assert isinstance(base, str) sub = self._unwrap_script(c_dict.get("sub", ""), "_") if self._needs_grouping(base): base = "{" + base + "}" - return base + (SUB % sub) + return f"{base}_{{{sub}}}" def do_ssup(self, elm: _Element) -> str: """Process superscript element. @@ -482,10 +491,11 @@ def do_ssup(self, elm: _Element) -> str: """ c_dict = self.process_children_dict(elm, include=("e", "sup", "sSupPr")) base = c_dict.get("e", "").rstrip() + assert isinstance(base, str) sup = self._unwrap_script(c_dict.get("sup", ""), "^") if self._needs_grouping(base): base = "{" + base + "}" - return base + (SUP % sup) + return f"{base}^{{{sup}}}" def do_ssubsup(self, elm: _Element) -> str: """Process combined sub-superscript element. @@ -500,11 +510,12 @@ def do_ssubsup(self, elm: _Element) -> str: elm, include=("e", "sub", "sup", "sSubSupPr") ) base = c_dict.get("e", "").rstrip() + assert isinstance(base, str) sub = self._unwrap_script(c_dict.get("sub", ""), "_") sup = self._unwrap_script(c_dict.get("sup", ""), "^") if self._needs_grouping(base): base = "{" + base + "}" - return base + (SUB % sub) + (SUP % sup) + return f"{base}_{{{sub}}}^{{{sup}}}" def do_sub(self, elm: _Element) -> str: """Process standalone subscript content. @@ -547,6 +558,7 @@ def do_f(self, elm: _Element) -> str: "num": c_dict.get("num"), "den": c_dict.get("den"), } + assert isinstance(pr, Pr) latex_s = get_val(pr.type, default=F_DEFAULT, store=F) return pr.text + ( latex_s % {"num": c_dict.get("num"), "den": c_dict.get("den")} @@ -563,6 +575,7 @@ def do_func(self, elm: _Element) -> str: """ c_dict = self.process_children_dict(elm) func_name = c_dict.get("fName", "") + assert isinstance(func_name, str) return func_name.replace(FUNC_PLACE, c_dict.get("e", "")) def do_fname(self, elm: _Element) -> str: @@ -602,6 +615,7 @@ def do_groupchr(self, elm: _Element) -> str: """ c_dict = self.process_children_dict(elm) pr = c_dict["groupChrPr"] + assert isinstance(pr, Pr) latex_s = get_val(pr.chr, default=CHR_DEFAULT["GROUPCHR_VAL"], store=CHR) # If latex_s contains %s, format it; otherwise return as-is (unmapped character) if "%s" in latex_s: @@ -780,7 +794,7 @@ def process_unicode(self, s: str) -> str: if s in self._MATH_CHAR_MAP: return self._MATH_CHAR_MAP[s] - out_latex_str = self.u.unicode_to_latex(s) + out_latex_str = str(self.u.unicode_to_latex(s)) if ( s.startswith("{") is False diff --git a/docling/backend/ebcdic_backend.py b/docling/backend/ebcdic_backend.py index 5642083064..8f9862a357 100644 --- a/docling/backend/ebcdic_backend.py +++ b/docling/backend/ebcdic_backend.py @@ -32,6 +32,7 @@ from docling.backend.abstract_backend import DeclarativeDocumentBackend from docling.datamodel.backend_options import ( + BaseBackendOptions, EbcdicBackendOptions, EbcdicField, EbcdicFieldType, @@ -221,10 +222,12 @@ def __init__( self, in_doc: InputDocument, path_or_stream: Union[BytesIO, Path], - options: Union[EbcdicBackendOptions, None] = None, + options: Union[BaseBackendOptions, None] = None, ) -> None: if options is None: options = EbcdicBackendOptions() + if not isinstance(options, EbcdicBackendOptions): + raise TypeError("Expected EbcdicBackendOptions.") super().__init__(in_doc, path_or_stream, options) self.layout = self._resolve_layout() diff --git a/docling/backend/email_backend.py b/docling/backend/email_backend.py index 9e5ea2e0c0..0a1577a5a1 100644 --- a/docling/backend/email_backend.py +++ b/docling/backend/email_backend.py @@ -12,6 +12,7 @@ from pathlib import Path from docling_core.types.doc import DocItemLabel, DoclingDocument, DocumentOrigin +from typing_extensions import override from docling.backend.abstract_backend import DeclarativeDocumentBackend from docling.backend.html_backend import HTMLDocumentBackend @@ -200,14 +201,17 @@ def _msg_to_rfc822_bytes(data: bytes) -> bytes: return email_message.as_bytes() + @override def is_valid(self) -> bool: return self.valid @classmethod + @override def supports_pagination(cls) -> bool: return False @classmethod + @override def supported_formats(cls) -> set[InputFormat]: return {InputFormat.EMAIL} @@ -316,6 +320,7 @@ def _get_attachment_labels(self) -> list[str]: labels.append(f"{filename} ({content_type})" if content_type else filename) return labels + @override def convert(self) -> DoclingDocument: if not self.is_valid() or self.mail is None: raise RuntimeError( diff --git a/docling/backend/epub_backend.py b/docling/backend/epub_backend.py index ec8d89b1ed..dc1a8dde30 100644 --- a/docling/backend/epub_backend.py +++ b/docling/backend/epub_backend.py @@ -28,6 +28,8 @@ class EpubDocumentBackend(DeclarativeDocumentBackend): + options: EpubBackendOptions + """Backend for converting EPUB files to DoclingDocument format. EPUB files are essentially ZIP archives containing XHTML content files. @@ -458,7 +460,7 @@ def convert(self) -> DoclingDocument: if self.temp_dir: dummy_file = self.temp_dir / "combined.html" - html_options.source_uri = str(dummy_file) + html_options.source_uri = dummy_file _log.debug(f"Set HTML backend source_uri to: {html_options.source_uri}") in_doc = InputDocument( diff --git a/docling/backend/html_backend.py b/docling/backend/html_backend.py index 31dbca22a8..6e15da8f8e 100644 --- a/docling/backend/html_backend.py +++ b/docling/backend/html_backend.py @@ -24,6 +24,7 @@ import requests import urllib3 +import urllib3.exceptions from docling_core.types.doc import ( BoundingBox, CodeLanguageLabel, @@ -39,6 +40,8 @@ GraphLinkLabel, GroupItem, GroupLabel, + ListItem, + NodeItem, PictureClassificationLabel, PictureClassificationMetaField, PictureClassificationPrediction, @@ -56,7 +59,7 @@ from docling_core.types.doc.document import ContentLayer, Formatting, ImageRef, Script from PIL import Image, UnidentifiedImageError from pydantic import AnyUrl, BaseModel, ValidationError -from typing_extensions import Self, override +from typing_extensions import override from docling.backend.abstract_backend import ( DeclarativeDocumentBackend, @@ -68,7 +71,7 @@ clamp_span, table_width, ) -from docling.datamodel.backend_options import HTMLBackendOptions +from docling.datamodel.backend_options import BaseBackendOptions, HTMLBackendOptions from docling.datamodel.base_models import FormatToMimeType, InputFormat from docling.datamodel.document import InputDocument from docling.exceptions import DocumentLoadError, OperationNotAllowed @@ -325,7 +328,7 @@ class AnnotatedText(BaseModel): source_tag_id: Optional[str] = None -class AnnotatedTextList(list): +class AnnotatedTextList(list[AnnotatedText]): def to_single_text_element(self) -> AnnotatedText: current_h = None current_text = "" @@ -423,7 +426,7 @@ def simplify_text_elements(self) -> AnnotatedTextList: ) return simplified - def split_by_newline(self) -> list[Self]: + def split_by_newline(self) -> list[AnnotatedTextList]: """Split text elements on multiple consecutive line breaks (from
tags). Single
tags are converted to newline characters (\n) within the same paragraph. @@ -431,7 +434,7 @@ def split_by_newline(self) -> list[Self]: Regular newlines from HTML source formatting have already been normalized to spaces during text extraction. """ - super_list: list[Self] = [] + super_list: list[AnnotatedTextList] = [] active_annotated_text_list = AnnotatedTextList() double_sentinel = _BR_SENTINEL + _BR_SENTINEL @@ -467,35 +470,40 @@ def __init__( self, in_doc: InputDocument, path_or_stream: Union[BytesIO, Path], - options: Optional[HTMLBackendOptions] = None, + options: Optional[BaseBackendOptions] = None, ): if not _BS4_AVAILABLE: raise ImportError(_INSTALL_HINT) from _BS4_IMPORT_ERROR if options is None: options = HTMLBackendOptions() - super().__init__(in_doc, path_or_stream, options) + typed_options = cast(HTMLBackendOptions, options) + super().__init__(in_doc, path_or_stream, typed_options) self.options: HTMLBackendOptions self.soup: Optional[BeautifulSoup] = None - self.path_or_stream: Union[BytesIO, Path] = path_or_stream + self.path_or_stream: Union[BytesIO, Path, None] = path_or_stream configured_base_path: Optional[str] = ( - str(options.source_uri) if options.source_uri is not None else None + str(typed_options.source_uri) + if typed_options.source_uri is not None + else None ) self.base_path = configured_base_path self._mhtml_resources: dict[str, bytes] | None = None self._image_loader = ImageResourceLoader( - enable_local_fetch=options.enable_local_fetch, - enable_remote_fetch=options.enable_remote_fetch, - max_image_data_base64_bytes=options.max_image_data_base64_bytes, - max_remote_image_bytes=options.max_remote_image_bytes, - max_redirects=options.max_redirects, - headers=options.headers, - header_origins=self._get_header_origins(options, configured_base_path), + enable_local_fetch=typed_options.enable_local_fetch, + enable_remote_fetch=typed_options.enable_remote_fetch, + max_image_data_base64_bytes=typed_options.max_image_data_base64_bytes, + max_remote_image_bytes=typed_options.max_remote_image_bytes, + max_redirects=typed_options.max_redirects, + headers=typed_options.headers, + header_origins=self._get_header_origins( + typed_options, configured_base_path + ), ) # Initialize the parents for the hierarchy self.max_levels = 10 self.level = 0 - self.parents: dict[int, Optional[Union[DocItem, GroupItem]]] = {} + self.parents: dict[int, NodeItem | None] = {} self.ctx = _Context() self._disable_inline_group_depth: int = 0 for i in range(self.max_levels): @@ -522,7 +530,7 @@ def __init__( else Path(path_or_stream).read_bytes() ) if self.input_format == InputFormat.MHTML: - if options.render_page: + if typed_options.render_page: raise DocumentLoadError( "Browser rendering is not supported for MHTML input." ) @@ -871,6 +879,8 @@ def convert(self) -> DoclingDocument: if not self.is_valid(): raise RuntimeError("Invalid HTML document.") + assert self.input_format is not None + origin = DocumentOrigin( filename=self.file.name or "file", mimetype=FormatToMimeType[self.input_format][0], @@ -878,13 +888,13 @@ def convert(self) -> DoclingDocument: ) doc = DoclingDocument(name=self.file.stem or "file", origin=origin) - if cast(HTMLBackendOptions, self.options).render_page: + if self.options.render_page: self._render_with_browser() if self._rendered_html: self.soup = BeautifulSoup(self._rendered_html, "html.parser") if self._rendered_page_images and self._rendered_page_size: - render_dpi = cast(HTMLBackendOptions, self.options).render_dpi + render_dpi = self.options.render_dpi for page_no, page_image in enumerate(self._rendered_page_images, start=1): doc.add_page( page_no=page_no, @@ -947,6 +957,7 @@ def _get_header_origins( options: HTMLBackendOptions, source_uri: Optional[str] ) -> list[str]: """Return the origins that receive ``options.headers``.""" + origins: list[str] if options.headers_allowed_origins is not None: origins = list(options.headers_allowed_origins) elif source_uri is not None and ImageResourceLoader.is_remote_url(source_uri): @@ -958,7 +969,7 @@ def _get_header_origins( return origins def _get_render_page_size(self) -> tuple[int, int]: - options = cast(HTMLBackendOptions, self.options) + options = self.options width = options.render_page_width height = options.render_page_height if options.render_page_orientation == "landscape": @@ -1079,12 +1090,12 @@ def _pad_image(self, image: Image.Image, width: int, height: int) -> Image.Image return canvas def _render_with_browser(self) -> None: - options = cast(HTMLBackendOptions, self.options) + options = self.options if not options.render_page: return try: - from playwright.sync_api import sync_playwright # type: ignore + from playwright.sync_api import sync_playwright except ImportError as exc: raise RuntimeError( "Playwright is required for HTML rendering. " @@ -1313,7 +1324,7 @@ def _route_request(route, request) -> None: def _capture_page_images( self, page, - render_data: dict, + render_data: dict[str, Any], page_width: int, page_height: int, full_page: bool, @@ -1349,7 +1360,7 @@ def _capture_page_images( return page_images def _build_bbox_mapping( - self, render_data: dict, page_height: int, full_page: bool + self, render_data: dict[str, Any], page_height: int, full_page: bool ) -> dict[str, _RenderedBBox]: boxes = render_data.get("boxes", {}) or {} scroll_height = float(render_data.get("scrollHeight", page_height)) @@ -1462,8 +1473,6 @@ def _has_rendered_presence(self, tag: Tag) -> bool: has_visible_descendant = False for descendant in tag.find_all(True): - if not isinstance(descendant, Tag): - continue if ( self._get_rendered_text_bbox_for_tag(descendant) is not None or self._get_rendered_bbox_for_tag(descendant) is not None @@ -1530,8 +1539,6 @@ def _make_text_prov( if render_box is None and isinstance(tag, Tag): descendant_boxes: list[_RenderedBBox] = [] for descendant in [tag, *tag.find_all(True)]: - if not isinstance(descendant, Tag): - continue descendant_box = self._get_rendered_text_bbox_for_tag(descendant) if descendant_box is not None: descendant_boxes.append(descendant_box) @@ -1623,7 +1630,7 @@ def _are_source_tag_ids_inline_neighbors( horizontal_gap = right_box.l - left_box.r max_gap = max(8.0, 1.5 * max_height) - min_gap = -0.5 * min(left_box.width, right_box.width) + min_gap = -0.5 * float(min(left_box.width, right_box.width)) return min_gap <= horizontal_gap <= max_gap def _compact_adjacent_single_char_parts( @@ -1737,8 +1744,11 @@ def _start_para(): nonlocal current_p if current_p is None: current_p = soup.new_tag("p") - if p.get(_DATA_DOCLING_ID_ATTR): - current_p[_DATA_DOCLING_ID_ATTR] = p.get(_DATA_DOCLING_ID_ATTR) + source_id = HTMLDocumentBackend._get_attr_as_string( + p, _DATA_DOCLING_ID_ATTR + ) + if source_id: + current_p[_DATA_DOCLING_ID_ATTR] = source_id new_nodes.append(current_p) def _flush_para_if_empty(): @@ -1828,10 +1838,10 @@ def group_cell_elements( group_element.children.append(prov) pr_item = prov.resolve(doc) item_parent = pr_item.parent.resolve(doc) - if pr_item.get_ref() in item_parent.children: - item_parent.children.remove(pr_item.get_ref()) - pr_item.parent = group_element.get_ref() - ref_for_rich_cell = group_element.get_ref() + if RefItem(cref=pr_item.self_ref) in item_parent.children: + item_parent.children.remove(RefItem(cref=pr_item.self_ref)) + pr_item.parent = RefItem(cref=group_element.self_ref) + ref_for_rich_cell = RefItem(cref=group_element.self_ref) return ref_for_rich_cell @staticmethod @@ -1877,12 +1887,10 @@ def _is_rich_table_cell(self, table_cell: Tag) -> bool: direct_block_text_children = [ child for child in table_cell.find_all(recursive=False) - if isinstance(child, Tag) and child.name in {"p", "div", "li"} + if child.name in {"p", "div", "li"} ] has_nested_form_semantic_id = any( - self._is_form_semantic_tag(child) - for child in children - if isinstance(child, Tag) + self._is_form_semantic_tag(child) for child in children ) if has_nested_form_semantic_id: return True @@ -1927,7 +1935,9 @@ def parse_table_data( t.unwrap() _log.debug(f"The table has {num_rows} rows and {num_cols} cols.") - grid: list = [[None for _ in range(num_cols)] for _ in range(num_rows)] + grid: list[list[str | None]] = [ + [None for _ in range(num_cols)] for _ in range(num_rows) + ] data = TableData(num_rows=num_rows, num_cols=num_cols, table_cells=[]) # Iterate over the rows in the table @@ -1936,8 +1946,6 @@ def parse_table_data( # We don't want this recursive to support nested tables for row in element("tr", recursive=False): - if not isinstance(row, Tag): - continue row_classes = { class_name.lower() for class_name in self._get_tag_classes(row) } @@ -1949,13 +1957,12 @@ def parse_table_data( col_header = True row_header = True for html_cell in cells: - if isinstance(html_cell, Tag): - _, row_span = HTMLDocumentBackend._get_cell_spans(html_cell) - if html_cell.name == "td": - col_header = False - row_header = False - elif row_span == 1: - row_header = False + _, row_span = HTMLDocumentBackend._get_cell_spans(html_cell) + if html_cell.name == "td": + col_header = False + row_header = False + elif row_span == 1: + row_header = False if not row_header: row_idx += 1 start_row_span = 0 @@ -1965,8 +1972,6 @@ def parse_table_data( # Extract the text content of each cell col_idx = 0 for html_cell in cells: - if not isinstance(html_cell, Tag): - continue cell_classes = { class_name.lower() for class_name in self._get_tag_classes(html_cell) @@ -1981,6 +1986,7 @@ def parse_table_data( formula.replace_with(NavigableString(math_formula)) provs_in_cell: list[RefItem] = [] + ref_for_rich_cell: RefItem | None = None rich_table_cell = self._is_rich_table_cell(html_cell) if rich_table_cell: # Parse table cell sub-tree for Rich Cells content: @@ -2021,6 +2027,7 @@ def parse_table_data( grid[row_idx + r][col_idx + c] = text if rich_table_cell: + assert ref_for_rich_cell is not None rich_cell = RichTableCell( text=text, bbox=cell_bbox, @@ -2109,7 +2116,9 @@ def _flush_buffer() -> None: prov=prov, ) if inline_ref is None: - added_refs.append(docling_code2.get_ref()) + added_refs.append( + RefItem(cref=docling_code2.self_ref) + ) else: prov = self._make_text_prov_for_source_tag_ids( text=seg_clean, @@ -2127,7 +2136,9 @@ def _flush_buffer() -> None: prov=prov, ) if inline_ref is None: - added_refs.append(docling_text2.get_ref()) + added_refs.append( + RefItem(cref=docling_text2.self_ref) + ) if inline_ref is not None: added_refs.append(inline_ref) @@ -2439,7 +2450,7 @@ def _get_tag_name_for_docling_id(self, source_tag_id: str) -> Optional[str]: return tag_name or None if self.soup is None: return None - tag = self.soup.find(attrs={_DATA_DOCLING_ID_ATTR: source_tag_id}) + tag = self.soup.find(True, attrs={_DATA_DOCLING_ID_ATTR: source_tag_id}) tag_name = tag.name if isinstance(tag, Tag) else "" self._tag_name_by_docling_id_cache[source_tag_id] = tag_name return tag_name or None @@ -2508,7 +2519,7 @@ def _use_inline_group( self.parents[self.level + 1] = inline_fmt self.level += 1 try: - yield inline_fmt.get_ref() + yield RefItem(cref=inline_fmt.self_ref) finally: self.parents[self.level] = None self.level -= 1 @@ -2591,7 +2602,7 @@ def _use_table_cell_context(self): self.parents = original_parents @contextmanager - def _use_list_item_context(self, parent_item: RefItem | None) -> Iterator[None]: + def _use_list_item_context(self, parent_item: NodeItem | None) -> Iterator[None]: """Set up context for processing nested content within a list item. Args: @@ -2613,7 +2624,7 @@ def _use_list_item_context(self, parent_item: RefItem | None) -> Iterator[None]: yield def _handle_heading(self, tag: Tag, doc: DoclingDocument) -> list[RefItem]: - added_ref = [] + added_ref: list[RefItem] = [] tag_name = tag.name.lower() # set default content layer to BODY as soon as we encounter a heading self.content_layer = ContentLayer.BODY @@ -2642,7 +2653,7 @@ def _handle_heading(self, tag: Tag, doc: DoclingDocument) -> list[RefItem]: ) p1 = self.parents[self.level + 1] if p1 is not None: - added_ref = [p1.get_ref()] + added_ref = [RefItem(cref=p1.self_ref)] # the other levels need to be lowered by 1 if a title was set else: level -= 1 @@ -2676,13 +2687,12 @@ def _handle_heading(self, tag: Tag, doc: DoclingDocument) -> list[RefItem]: ) p2 = self.parents[self.level + 1] if p2 is not None: - added_ref = [p2.get_ref()] + added_ref = [RefItem(cref=p2.self_ref)] self.level += 1 for img_tag in tag("img"): - if isinstance(img_tag, Tag): - im_ref = self._emit_image(img_tag, doc) - if im_ref: - added_ref.append(im_ref) + im_ref = self._emit_image(img_tag, doc) + if im_ref: + added_ref.append(im_ref) return added_ref def _has_list_ancestor(self, elem: Tag, boundary: Tag) -> bool: @@ -2697,7 +2707,7 @@ def _has_list_ancestor(self, elem: Tag, boundary: Tag) -> bool: """ parent = elem.parent while parent and parent != boundary: - if isinstance(parent, Tag) and parent.name in {"ul", "ol", "dl"}: + if parent.name in {"ul", "ol", "dl"}: return True parent = parent.parent return False @@ -2707,7 +2717,7 @@ def _process_nested_element( elem, li: Tag, doc: DoclingDocument, - processed_elements: set, + processed_elements: set[int], ) -> None: """Process a single nested element within a list item. @@ -2746,7 +2756,7 @@ def _process_list_item_nested_content( self, li: Tag, doc: DoclingDocument, - processed_elements: set, + processed_elements: set[int], ) -> None: """Process nested content (images, lists, etc.) within a list item in DOM order. @@ -2762,11 +2772,11 @@ def _add_list_item_with_content( self, tag: Tag, doc: DoclingDocument, - parent: RefItem, + parent: NodeItem, enumerated: bool = False, marker: str = "", extra_formatting: Optional[Formatting] = None, - ) -> Optional[RefItem]: + ) -> Optional[ListItem]: """Helper method to add a list item with its content. Handles both simple and complex content with inline groups. @@ -2842,6 +2852,7 @@ def _add_list_item_with_content( self.parents[self.level] = None self.level -= 1 + assert isinstance(list_item, ListItem) return list_item else: # Simple content - single text element @@ -2873,6 +2884,7 @@ def _add_list_item_with_content( hyperlink=annotated_text.hyperlink, prov=prov, ) + assert isinstance(list_item, ListItem) return list_item def _emit_task_list_inputs( @@ -2880,8 +2892,7 @@ def _emit_task_list_inputs( ) -> None: """Emit the checkbox items of a task-list
  • under the list group.""" for input_tag in inputs_in_li: - if isinstance(input_tag, Tag): - self._emit_input(input_tag, doc) + self._emit_input(input_tag, doc) def _is_task_list_item( self, li: Tag, inputs_in_li: list[Tag], custom_checkboxes_in_li: list[Tag] @@ -2900,7 +2911,7 @@ def _description_list_children(dl: Tag) -> list[PageElement]: (HTML allows wrapping each group of
    and
    elements in a
    ).""" children: list[PageElement] = [] for child in dl.find_all(["dt", "dd", "div"], recursive=False): - if isinstance(child, Tag) and child.name == "div": + if child.name == "div": children.extend(child.find_all(["dt", "dd"], recursive=False)) else: children.append(child) @@ -3007,13 +3018,10 @@ def _handle_list(self, tag: Tag, doc: DoclingDocument) -> RefItem: self.parents[self.level + 1] = None self.level -= 1 - return list_group.get_ref() + return RefItem(cref=list_group.self_ref) # For each top-level
  • in this list (ul/ol) for li in tag.find_all({"li", "ul", "ol"}, recursive=False): - if not isinstance(li, Tag): - continue - # sub-list items should be indented under main list items, but temporarily # addressing invalid HTML (docling-core/issues/357) if li.name in {"ul", "ol"}: @@ -3076,11 +3084,9 @@ def _handle_list(self, tag: Tag, doc: DoclingDocument) -> RefItem: # Handle inputs and checkboxes if inputs_in_li or custom_checkboxes_in_li: for input_tag in inputs_in_li: - if isinstance(input_tag, Tag): - self._emit_input(input_tag, doc) + self._emit_input(input_tag, doc) for checkbox_tag in custom_checkboxes_in_li: - if isinstance(checkbox_tag, Tag): - self._emit_custom_checkbox(checkbox_tag, doc) + self._emit_custom_checkbox(checkbox_tag, doc) # 4) Process nested content (images, lists, etc.) in DOM order processed_elements = set() @@ -3090,16 +3096,14 @@ def _handle_list(self, tag: Tag, doc: DoclingDocument) -> RefItem: else: # No content, but check for nested lists (including those wrapped in divs) for sublist in li({"ul", "ol", "dl"}): - if isinstance(sublist, Tag): - # Check if this list has a ul/ol/dl ancestor within the current li - has_list_ancestor = self._has_list_ancestor(sublist, li) + has_list_ancestor = self._has_list_ancestor(sublist, li) - if not has_list_ancestor: - self._handle_block(sublist, doc) + if not has_list_ancestor: + self._handle_block(sublist, doc) self.parents[self.level + 1] = None self.level -= 1 - return list_group.get_ref() + return RefItem(cref=list_group.self_ref) @staticmethod def get_html_table_row_col(tag: Tag) -> tuple[int, int]: @@ -3111,12 +3115,8 @@ def get_html_table_row_col(tag: Tag) -> tuple[int, int]: for row in tag("tr", recursive=False): cell_spans: list[tuple[int, int]] = [] is_row_header = True - if not isinstance(row, Tag): - continue for cell in row(["td", "th"], recursive=False): - if not isinstance(row, Tag): - continue - cell_tag = cast(Tag, cell) + cell_tag = cell col_span, row_span = HTMLDocumentBackend._get_cell_spans(cell_tag) cell_spans.append((col_span, row_span)) if cell_tag.name == "td" or row_span == 1: @@ -3167,7 +3167,7 @@ def _handle_block(self, tag: Tag, doc: DoclingDocument) -> list[RefItem]: # noq if cap_item is not None and added_refs: first_item = added_refs[0].resolve(doc) if isinstance(first_item, TableItem): - first_item.captions.append(cap_item.get_ref()) + first_item.captions.append(RefItem(cref=cap_item.self_ref)) elif tag_name in {"h1", "h2", "h3", "h4", "h5", "h6"}: heading_refs = self._handle_heading(tag, doc) @@ -3204,7 +3204,9 @@ def _handle_block(self, tag: Tag, doc: DoclingDocument) -> list[RefItem]: # noq prov=prov, ) if inline_ref is None: - added_refs.append(docling_code.get_ref()) + added_refs.append( + RefItem(cref=docling_code.self_ref) + ) else: prov = self._make_text_prov_for_source_tag_ids( text=seg_clean, @@ -3222,27 +3224,26 @@ def _handle_block(self, tag: Tag, doc: DoclingDocument) -> list[RefItem]: # noq prov=prov, ) if inline_ref is None: - added_refs.append(docling_text.get_ref()) + added_refs.append( + RefItem(cref=docling_text.self_ref) + ) if inline_ref is not None: added_refs.append(inline_ref) for img_tag in tag("img"): - if isinstance(img_tag, Tag): - self._emit_image(img_tag, doc) + self._emit_image(img_tag, doc) for input_tag in tag("input"): - if isinstance(input_tag, Tag): - input_ref = self._emit_input(input_tag, doc) - if input_ref is not None: - added_refs.append(input_ref) + input_ref = self._emit_input(input_tag, doc) + if input_ref is not None: + added_refs.append(input_ref) for checkbox_tag in tag.find_all( lambda item: ( isinstance(item, Tag) and self._is_custom_checkbox_tag(item) ) ): - if isinstance(checkbox_tag, Tag): - checkbox_ref = self._emit_custom_checkbox(checkbox_tag, doc) - if checkbox_ref is not None: - added_refs.append(checkbox_ref) + checkbox_ref = self._emit_custom_checkbox(checkbox_tag, doc) + if checkbox_ref is not None: + added_refs.append(checkbox_ref) elif tag_name == "table": num_rows, num_cols = self.get_html_table_row_col(tag) @@ -3261,17 +3262,18 @@ def _handle_block(self, tag: Tag, doc: DoclingDocument) -> list[RefItem]: # noq content_layer=self.content_layer, caption=cap_item, ) - added_refs.append(docling_table.get_ref()) + added_refs.append(RefItem(cref=docling_table.self_ref)) self.parse_table_data(tag, doc, docling_table, num_rows, num_cols) elif tag_name in {"stamp", "signature"}: _class_name = PictureClassificationLabel.STAMP.value if tag_name == "signature": _class_name = PictureClassificationLabel.SIGNATURE.value - placeholder: PictureItem = doc.add_picture( + placeholder = doc.add_picture( parent=self.parents[self.level], content_layer=self.content_layer, ) + assert isinstance(placeholder, PictureItem) placeholder.meta = PictureMeta( classification=PictureClassificationMetaField( predictions=[ @@ -3318,7 +3320,7 @@ def _handle_block(self, tag: Tag, doc: DoclingDocument) -> list[RefItem]: # noq prov=prov, ) if inline_ref is None: - added_refs.append(docling_code2.get_ref()) + added_refs.append(RefItem(cref=docling_code2.self_ref)) if inline_ref is not None: added_refs.append(inline_ref) @@ -3341,7 +3343,7 @@ def _is_form_container(tag: Tag) -> bool: def _nearest_form_container_ancestor(self, tag: Tag) -> Optional[Tag]: for parent in tag.parents: - if isinstance(parent, Tag) and self._is_form_container(parent): + if self._is_form_container(parent): return parent return None @@ -3384,7 +3386,7 @@ def _emit_flattened_text_tag( content_layer=self.content_layer, prov=prov, ) - return text_item.get_ref() + return RefItem(cref=text_item.self_ref) def _ensure_tag_html_id(self, tag: Tag) -> str: existing = self._get_html_id(tag) @@ -3413,13 +3415,9 @@ def _dom_distance_between_tags(left_tag: Tag, right_tag: Tag) -> int: return 0 left_chain: list[Tag] = [left_tag] - left_chain.extend( - parent for parent in left_tag.parents if isinstance(parent, Tag) - ) + left_chain.extend(parent for parent in left_tag.parents) right_chain: list[Tag] = [right_tag] - right_chain.extend( - parent for parent in right_tag.parents if isinstance(parent, Tag) - ) + right_chain.extend(parent for parent in right_tag.parents) left_positions = {id(tag): idx for idx, tag in enumerate(left_chain)} best_distance: Optional[int] = None @@ -3464,14 +3462,7 @@ def _select_form_value_entries( ranked_entries = sorted( entries, key=lambda entry: ( - ( - 0 - if ( - key_tag is not None - and self._is_value_in_key_scope(key_tag, entry[2]) - ) - else 1 - ) + (0 if (self._is_value_in_key_scope(key_tag, entry[2])) else 1) if key_tag is not None else 0, ( @@ -3631,7 +3622,7 @@ def _get_tag_classes(tag: Tag) -> set[str]: return set() if isinstance(classes, str): return {classes} - return {str(value) for value in classes if isinstance(value, str)} + return set(classes) @staticmethod def _code_language_hint(tag: Tag) -> str | None: @@ -3659,7 +3650,7 @@ def _has_inline_display_style(tag: Tag) -> bool: display_match = re.search(r"display\s*:\s*([^;]+)", style_attr, flags=re.I) if display_match is None: return False - display_value = display_match.group(1).strip().lower() + display_value = str(display_match.group(1)).strip().lower() return display_value.startswith("inline") or display_value == "contents" def _should_buffer_tag_text_inline(self, tag: Tag) -> bool: @@ -3738,9 +3729,7 @@ def _extract(node: PageElement) -> list[str]: @staticmethod def _has_direct_checkbox_like_child(tag: Tag) -> bool: for child in tag.find_all(recursive=False): - if isinstance(child, Tag) and HTMLDocumentBackend._is_checkbox_like_tag( - child - ): + if HTMLDocumentBackend._is_checkbox_like_tag(child): return True return False @@ -3836,8 +3825,7 @@ def _label_texts(tags: list[Tag]) -> list[str]: for label_tag in self.soup.find_all( "label", attrs={"for": input_id} ): - if isinstance(label_tag, Tag): - _add_label_tag(label_tag) + _add_label_tag(label_tag) # Backward-compatible local search (same parent container). if parent is not None: for sibling in parent.find_all("label", recursive=False): @@ -3872,7 +3860,7 @@ def _label_texts(tags: list[Tag]) -> list[str]: if "checkbox-container" in parent_classes: label_texts: list[str] = [] for sibling in parent.find_all(recursive=False): - if not isinstance(sibling, Tag) or sibling is checkbox_tag: + if sibling is checkbox_tag: continue if "checkbox-label" not in self._get_tag_classes(sibling): continue @@ -3895,7 +3883,7 @@ def _label_texts(tags: list[Tag]) -> list[str]: for child in parent.contents ) for sibling in parent.find_all(recursive=False): - if not isinstance(sibling, Tag) or sibling is checkbox_tag: + if sibling is checkbox_tag: continue if self._is_checkbox_like_tag(sibling): continue @@ -4091,7 +4079,7 @@ def _emit_custom_checkbox( hyperlink=self.hyperlink, prov=prov, ) - return checkbox_item.get_ref() + return RefItem(cref=checkbox_item.self_ref) @contextmanager def _suppress_tag_ids(self, tag_ids: set[str]): @@ -4277,9 +4265,7 @@ def _is_simple_extra_text_candidate(self, tag: Tag) -> bool: ): return False - direct_tag_children = [ - child for child in tag.find_all(recursive=False) if isinstance(child, Tag) - ] + direct_tag_children = list(tag.find_all(recursive=False)) has_direct_text = any( isinstance(child, NavigableString) and bool(str(child).strip()) for child in tag.contents @@ -4299,7 +4285,7 @@ def _add_field_item_from_extracted( self, field: _ExtractedFormField, doc: DoclingDocument, - parent: Optional[Union[DocItem, GroupItem]], + parent: NodeItem | None, ) -> list[RefItem]: refs: list[RefItem] = [] doc_with_fields = cast(Any, doc) @@ -4307,7 +4293,7 @@ def _add_field_item_from_extracted( parent=parent, content_layer=self.content_layer, ) - refs.append(field_item.get_ref()) + refs.append(RefItem(cref=field_item.self_ref)) parts: list[tuple[int, Literal["key", "marker", "value", "text"], Any]] = [] if field.key_tag is not None and field.key_text: @@ -4328,7 +4314,7 @@ def _add_field_item_from_extracted( parent=field_item, content_layer=self.content_layer, ) - refs.append(field_key.get_ref()) + refs.append(RefItem(cref=field_key.self_ref)) elif part_type == "marker": marker = cast(_ExtractedFormMarker, payload) marker_item = doc.add_text( @@ -4339,7 +4325,7 @@ def _add_field_item_from_extracted( parent=field_item, content_layer=self.content_layer, ) - refs.append(marker_item.get_ref()) + refs.append(RefItem(cref=marker_item.self_ref)) elif part_type == "text": extra_text = cast(_ExtractedFormText, payload) text_item = doc.add_text( @@ -4350,7 +4336,7 @@ def _add_field_item_from_extracted( parent=field_item, content_layer=self.content_layer, ) - refs.append(text_item.get_ref()) + refs.append(RefItem(cref=text_item.self_ref)) else: value = cast(_ExtractedFormValue, payload) if value.checkbox_label is not None: @@ -4362,7 +4348,7 @@ def _add_field_item_from_extracted( parent=field_item, content_layer=self.content_layer, ) - refs.append(checkbox_item.get_ref()) + refs.append(RefItem(cref=checkbox_item.self_ref)) else: field_value = doc_with_fields.add_field_value( text=value.text, @@ -4372,7 +4358,7 @@ def _add_field_item_from_extracted( content_layer=self.content_layer, kind=value.kind, ) - refs.append(field_value.get_ref()) + refs.append(RefItem(cref=field_value.self_ref)) return refs @@ -4430,10 +4416,9 @@ def _extract_form_region(self, form_tag: Tag) -> Optional[_ExtractedFormRegion]: table_bboxes: list[BoundingBox] = [] if self._rendered_bbox_by_id: for table_tag in form_tag.find_all("table"): - if isinstance(table_tag, Tag): - rendered = self._get_rendered_bbox_for_tag(table_tag) - if rendered is not None: - table_bboxes.append(rendered.bbox) + rendered = self._get_rendered_bbox_for_tag(table_tag) + if rendered is not None: + table_bboxes.append(rendered.bbox) key_ids_in_order = sorted( first_order_by_key.keys(), key=lambda key_id: first_order_by_key[key_id] @@ -4637,11 +4622,7 @@ def _extract_form_region(self, form_tag: Tag) -> Optional[_ExtractedFormRegion]: for parent_tag in unique_parent_tags_to_scan: component_direct_child_indices: list[int] = [] - direct_children = [ - child - for child in parent_tag.find_all(recursive=False) - if isinstance(child, Tag) - ] + direct_children = list(parent_tag.find_all(recursive=False)) for idx, child in enumerate(direct_children): child_obj_id = id(child) if child_obj_id in component_tag_obj_ids: @@ -4917,7 +4898,7 @@ def _handle_form_container(self, tag: Tag, doc: DoclingDocument) -> list[RefItem parent=self.parents[self.level], ) field_region.content_layer = self.content_layer - added_refs.append(field_region.get_ref()) + added_refs.append(RefItem(cref=field_region.self_ref)) with self._use_form_container(field_region): with self._use_form_fields_by_key_id(fields_by_key_id): @@ -4939,7 +4920,7 @@ def _handle_form_container(self, tag: Tag, doc: DoclingDocument) -> list[RefItem parent=self.parents[self.level], ) form_item.content_layer = self.content_layer - added_refs.append(form_item.get_ref()) + added_refs.append(RefItem(cref=form_item.self_ref)) if form_graph is not None: kv_item = doc.add_key_values( @@ -4948,7 +4929,7 @@ def _handle_form_container(self, tag: Tag, doc: DoclingDocument) -> list[RefItem parent=form_item, ) kv_item.content_layer = self.content_layer - added_refs.append(kv_item.get_ref()) + added_refs.append(RefItem(cref=kv_item.self_ref)) with self._use_form_container(form_item): if tag.name.lower() == "table": @@ -4974,7 +4955,7 @@ def _emit_caption( tag=caption_tag, source_tag_id=cap_anno.source_tag_id, ) - return doc.add_text( + caption = doc.add_text( label=DocItemLabel.CAPTION, text=cap_text, orig=cap_anno.text, @@ -4985,6 +4966,9 @@ def _emit_caption( parent=self.parents[self.level], ) + assert isinstance(caption, TextItem) + return caption + def _emit_image(self, img_tag: Tag, doc: DoclingDocument) -> Optional[RefItem]: figure = img_tag.find_parent("figure") caption: AnnotatedTextList = AnnotatedTextList() @@ -5002,7 +4986,7 @@ def get_img_hyperlink(img_tag): return None if img_hyperlink := get_img_hyperlink(img_tag): - img_text = img_tag.get("alt") or "" + img_text = self._get_attr_as_string(img_tag, "alt") caption.append(AnnotatedText(text=img_text, hyperlink=img_hyperlink)) caption_prov_tag = img_tag @@ -5014,7 +4998,9 @@ def get_img_hyperlink(img_tag): ) caption_prov_tag = caption_tag if not caption and img_tag.get("alt"): - caption = AnnotatedTextList([AnnotatedText(text=img_tag.get("alt"))]) + caption = AnnotatedTextList( + [AnnotatedText(text=self._get_attr_as_string(img_tag, "alt"))] + ) caption_prov_tag = img_tag caption_anno_text = caption.to_single_text_element() @@ -5029,7 +5015,7 @@ def get_img_hyperlink(img_tag): tag=caption_prov_tag or img_tag, source_tag_id=caption_anno_text.source_tag_id, ) - caption_item = doc.add_text( + emitted_caption = doc.add_text( label=DocItemLabel.CAPTION, text=text_clean, orig=caption_anno_text.text, @@ -5039,17 +5025,21 @@ def get_img_hyperlink(img_tag): prov=prov, ) + assert isinstance(emitted_caption, TextItem) + caption_item = emitted_caption + src_loc: str = self._get_attr_as_string(img_tag, "src") pic_prov = self._make_prov(text="", tag=img_tag) - if not cast(HTMLBackendOptions, self.options).fetch_images or not src_loc: + if not self.options.fetch_images or not src_loc: # Do not fetch the image, just add a placeholder - placeholder: PictureItem = doc.add_picture( + placeholder = doc.add_picture( caption=caption_item, parent=parent, content_layer=self.content_layer, prov=pic_prov, ) - return placeholder.get_ref() + assert isinstance(placeholder, PictureItem) + return RefItem(cref=placeholder.self_ref) src_loc = self._resolve_relative_path(src_loc) img_ref = self._create_image_ref(src_loc) @@ -5061,7 +5051,7 @@ def get_img_hyperlink(img_tag): content_layer=self.content_layer, prov=pic_prov, ) - return docling_pic.get_ref() + return RefItem(cref=docling_pic.self_ref) def _emit_input(self, input_tag: Tag, doc: DoclingDocument) -> Optional[RefItem]: if self._is_suppressed_tag(input_tag): @@ -5071,6 +5061,7 @@ def _emit_input(self, input_tag: Tag, doc: DoclingDocument) -> Optional[RefItem] return None label = DocItemLabel.TEXT + checkbox_label_tags: list[Tag] = [] checkbox_label = self._get_checkbox_label_for_tag(input_tag) if checkbox_label is not None: label = checkbox_label @@ -5120,7 +5111,7 @@ def _emit_input(self, input_tag: Tag, doc: DoclingDocument) -> Optional[RefItem] hyperlink=self.hyperlink, prov=prov, ) - return input_item.get_ref() + return RefItem(cref=input_item.self_ref) def _create_image_ref(self, src_url: str) -> Optional[ImageRef]: if self._mhtml_resources is not None: diff --git a/docling/backend/image_backend.py b/docling/backend/image_backend.py index 0deb1866d0..4adb7b4a1f 100644 --- a/docling/backend/image_backend.py +++ b/docling/backend/image_backend.py @@ -17,6 +17,7 @@ TextCell, ) from PIL import Image, ImageOps +from typing_extensions import override from docling.backend.abstract_backend import AbstractDocumentBackend from docling.backend.pdf_backend import PdfDocumentBackend, PdfPageBackend @@ -93,16 +94,20 @@ def __init__( self.valid: bool = self._image is not None @property + @override def page_no(self) -> int: return self._page_no + 1 + @override def is_valid(self) -> bool: return self.valid + @override def get_text_in_rect(self, bbox: BoundingBox) -> str: # No text extraction from raw images without OCR return "" + @override def get_segmented_page(self) -> SegmentedPdfPage: # Return empty segmented page with proper dimensions for raw images assert self._image is not None @@ -134,10 +139,12 @@ def get_segmented_page(self) -> SegmentedPdfPage: has_lines=False, ) + @override def get_text_cells(self) -> Iterable[TextCell]: # No text cells on raw images - return [] + return list[TextCell]() + @override def get_bitmap_rects(self, scale: float = 1) -> Iterable[BoundingBox]: # For raw images, the entire page is a bitmap assert self._image is not None @@ -153,6 +160,7 @@ def get_bitmap_rects(self, scale: float = 1) -> Iterable[BoundingBox]: full_page_bbox = full_page_bbox.scaled(scale=scale) yield full_page_bbox + @override def get_page_image( self, scale: float = 1, cropbox: BoundingBox | None = None ) -> Image.Image: @@ -194,6 +202,7 @@ def get_page_image( box=source_box, ) + @override def get_size(self) -> Size: assert self._image is not None return Size( @@ -201,6 +210,7 @@ def get_size(self) -> Size: height=self._image.height * _POINTS_PER_INCH / self._dpi[1], ) + @override def unload(self): # Help GC and free memory self._image = None @@ -238,7 +248,7 @@ def __init__( self._frames: list[Image.Image] = [] self._frame_dpi: list[tuple[float, float]] = [] try: - with Image.open(self.path_or_stream) as img: # type: ignore[arg-type] + with Image.open(path_or_stream) as img: # Handle multi-frame and single-frame images # - multiframe formats: TIFF, GIF, ICO # - singleframe formats: JPEG (.jpg, .jpeg), PNG (.png), BMP, WEBP (unless animated), HEIC @@ -263,12 +273,15 @@ def __init__( f"Could not load image for document {self.file}" ) from e + @override def is_valid(self) -> bool: return len(self._frames) > 0 + @override def page_count(self) -> int: return len(self._frames) + @override def load_page(self, page_no: int) -> _ImagePageBackend: if not (0 <= page_no < len(self._frames)): raise IndexError(f"Page index out of range: {page_no}") @@ -277,14 +290,17 @@ def load_page(self, page_no: int) -> _ImagePageBackend: ) @classmethod + @override def supported_formats(cls) -> set[InputFormat]: # Only IMAGE here; PDF handling remains in PDF-oriented backends return {InputFormat.IMAGE} @classmethod + @override def supports_pagination(cls) -> bool: return True + @override def unload(self): for frame in self._frames: frame.close() diff --git a/docling/backend/iwork/cells.py b/docling/backend/iwork/cells.py index 1915c97695..85dd77e6c4 100644 --- a/docling/backend/iwork/cells.py +++ b/docling/backend/iwork/cells.py @@ -263,7 +263,7 @@ def iwa_current_cell(storage: bytes, start: int, values: CellValues) -> Cell | N decoded = decoded._replace( text=values.strings.get(read_uint32(storage, offset)) ) - elif flag == CELL_FLAG_RICH_TEXT: + else: decoded = decoded._replace( text=values.rich_text.get(read_uint32(storage, offset)) ) diff --git a/docling/backend/iwork/keynote_iwa.py b/docling/backend/iwork/keynote_iwa.py index a2345ebab2..675a0b17fa 100644 --- a/docling/backend/iwork/keynote_iwa.py +++ b/docling/backend/iwork/keynote_iwa.py @@ -19,6 +19,7 @@ import zipfile from docling_core.types.doc import DocItemLabel +from typing_extensions import override from docling.backend.iwork.archives import ( SHAPE_STORAGE_FIELD, @@ -348,6 +349,7 @@ def _blocks( for block in self._drawable_blocks(identifier) ], [] + @override def _drawable_blocks(self, identifier: int) -> list[Block]: """Read a drawable, reading a chart here and deferring anything else. diff --git a/docling/backend/iwork/legacy.py b/docling/backend/iwork/legacy.py index 117bab6434..f8c6fe44d6 100644 --- a/docling/backend/iwork/legacy.py +++ b/docling/backend/iwork/legacy.py @@ -653,7 +653,7 @@ def legacy_formatting(style: Element) -> Formatting | None: label = SF_PROPERTY_LABELS.get(element.tag) if label is not None: active.add(label) - elif element.tag == SF_SUPERSCRIPT and number is not None: + elif element.tag == SF_SUPERSCRIPT: script = SCRIPTS.get(as_int(number)) return build_formatting(active, script) diff --git a/docling/backend/iwork/numbers_xml.py b/docling/backend/iwork/numbers_xml.py index 95b4a50ae5..1f43363d58 100644 --- a/docling/backend/iwork/numbers_xml.py +++ b/docling/backend/iwork/numbers_xml.py @@ -17,6 +17,7 @@ import logging import plistlib import zipfile +from typing import cast from xml.etree.ElementTree import Element from docling.backend.iwork.content import Chart, ChartKind, ChartSeries @@ -177,7 +178,9 @@ def read_content( return sheets -def read_chart_shares(archive: zipfile.ZipFile, max_file_bytes: int) -> dict[str, dict]: +def read_chart_shares( + archive: zipfile.ZipFile, max_file_bytes: int +) -> dict[str, dict[str, object]]: """Read the cached data of every chart in the container. Args: @@ -189,7 +192,7 @@ def read_chart_shares(archive: zipfile.ZipFile, max_file_bytes: int) -> dict[str is missing or unreadable is left out rather than failing the document, since the chart still has its names without it. """ - shares: dict[str, dict] = {} + shares: dict[str, dict[str, object]] = {} for info in archive.infolist(): if not info.filename.endswith(CHART_SHARE_SUFFIX): continue @@ -206,7 +209,9 @@ def read_chart_shares(archive: zipfile.ZipFile, max_file_bytes: int) -> dict[str _log.debug("Could not read chart data '%s'.", info.filename) continue if isinstance(share, dict): - shares[info.filename[: -len(CHART_SHARE_SUFFIX)]] = share + shares[info.filename[: -len(CHART_SHARE_SUFFIX)]] = cast( + dict[str, object], share + ) return shares @@ -426,7 +431,9 @@ def formula_text(cell: Element) -> str | None: return None -def read_chart(info: Element, shares: dict[str, dict]) -> PlacedChart | None: +def read_chart( + info: Element, shares: dict[str, dict[str, object]] +) -> PlacedChart | None: """Build one chart from an ``sf:chart-info``, and say where it sits. A chart bound to a table leaves ``sf:chart-data`` empty and keeps the values @@ -456,10 +463,18 @@ def read_chart(info: Element, shares: dict[str, dict]) -> PlacedChart | None: if share is not None: # The share is what the chart last drew, so where the two disagree — a # stale name left behind in the element, say — the share wins. - shared_names = [str(name) for name in share.get(SHARE_COLUMNS_KEY, [])] + columns = share.get(SHARE_COLUMNS_KEY, []) + shared_names = ( + [str(name) for name in columns] if isinstance(columns, list) else [] + ) if shared_names: names = shared_names - rows = share.get(SHARE_ROWS_KEY) or [] + raw_rows = share.get(SHARE_ROWS_KEY) or [] + rows = ( + [cast(dict[str, object], row) for row in raw_rows if isinstance(row, dict)] + if isinstance(raw_rows, list) + else [] + ) if rows: categories = [str(row.get(SHARE_ROW_NAME_KEY, "")) for row in rows] by_category = [points(row) for row in rows] @@ -508,13 +523,16 @@ def _point_of( return row[series] if series < len(row) else None -def points(row: dict) -> list[float | None]: +def points(row: dict[str, object]) -> list[float | None]: """Read one category's plotted values, leaving the gaps empty.""" + values = row.get(SHARE_ROW_VALUES_KEY, []) + if not isinstance(values, list): + return [] return [ float(value) if isinstance(value, (int, float)) and not isinstance(value, bool) else None - for value in row.get(SHARE_ROW_VALUES_KEY, []) + for value in values ] diff --git a/docling/backend/iwork_backend.py b/docling/backend/iwork_backend.py index dde1023bd5..1e47e8eba9 100644 --- a/docling/backend/iwork_backend.py +++ b/docling/backend/iwork_backend.py @@ -86,7 +86,7 @@ ) from docling.backend.iwork.iwa import is_encrypted from docling.backend.iwork.keynote_content import Presentation, Slide -from docling.datamodel.backend_options import IWorkBackendOptions +from docling.datamodel.backend_options import BaseBackendOptions, IWorkBackendOptions from docling.datamodel.base_models import InputFormat from docling.datamodel.document import InputDocument from docling.exceptions import DocumentLoadError @@ -292,10 +292,12 @@ def __init__( self, in_doc: InputDocument, path_or_stream: BytesIO | Path, - options: IWorkBackendOptions | None = None, + options: BaseBackendOptions | None = None, ): if options is None: options = IWorkBackendOptions() + if not isinstance(options, IWorkBackendOptions): + raise TypeError("Expected IWorkBackendOptions.") super().__init__(in_doc, path_or_stream, options) self.options: IWorkBackendOptions = options @@ -422,10 +424,12 @@ def __init__( self, in_doc: InputDocument, path_or_stream: BytesIO | Path, - options: IWorkBackendOptions | None = None, + options: BaseBackendOptions | None = None, ): if options is None: options = IWorkBackendOptions() + if not isinstance(options, IWorkBackendOptions): + raise TypeError("Expected IWorkBackendOptions.") super().__init__(in_doc, path_or_stream, options) self.options: IWorkBackendOptions = options @@ -924,14 +928,18 @@ def _add_paragraph( return _add_list_item(doc, paragraph, paragraph.list_label, lists, prov) if paragraph.label == DocItemLabel.TITLE: - return doc.add_title(text=paragraph.text, parent=parent, prov=prov) + item = doc.add_title(text=paragraph.text, parent=parent, prov=prov) + assert isinstance(item, TextItem) + return item if paragraph.label == DocItemLabel.SECTION_HEADER: - return doc.add_heading( + item = doc.add_heading( text=paragraph.text, level=paragraph.level or 1, parent=parent, prov=prov, ) + assert isinstance(item, TextItem) + return item return _add_runs(doc, paragraph, paragraph.label, parent=parent, prov=prov) @@ -966,7 +974,7 @@ def _add_runs( runs = [run for run in paragraph.runs if run.text] if _uniform(runs): first = runs[0] if runs else Run("", None) - return doc.add_text( + item = doc.add_text( label=label, text=paragraph.text, formatting=first.formatting, @@ -975,6 +983,8 @@ def _add_runs( parent=parent, prov=prov, ) + assert isinstance(item, TextItem) + return item group = doc.add_inline_group(content_layer=content_layer, parent=parent) items = [ @@ -989,7 +999,9 @@ def _add_runs( ) for run in runs ] - return items[0] + first_item = items[0] + assert isinstance(first_item, TextItem) + return first_item def _add_list_item( @@ -1003,7 +1015,7 @@ def _add_list_item( group = lists.group_for(label.depth) runs = [run for run in paragraph.runs if run.text] uniform = runs[0] if runs and _uniform(runs) else Run("", None) - return doc.add_list_item( + item = doc.add_list_item( text=paragraph.text, enumerated=label.enumerated, marker=label.marker, @@ -1012,6 +1024,8 @@ def _add_list_item( hyperlink=_hyperlink(uniform.hyperlink), prov=prov, ) + assert isinstance(item, TextItem) + return item def _uniform(runs: list[Run]) -> bool: @@ -1087,10 +1101,12 @@ def __init__( self, in_doc: InputDocument, path_or_stream: BytesIO | Path, - options: IWorkBackendOptions | None = None, + options: BaseBackendOptions | None = None, ): if options is None: options = IWorkBackendOptions() + if not isinstance(options, IWorkBackendOptions): + raise TypeError("Expected IWorkBackendOptions.") super().__init__(in_doc, path_or_stream, options) self.options: IWorkBackendOptions = options self.page_range = in_doc.limits.page_range @@ -1203,7 +1219,7 @@ def convert(self) -> DoclingDocument: for drawable in sorted(drawn, key=numbers_content.sheet_order): if isinstance(drawable, numbers_content.Table): _add_sheet_table(doc, drawable, parent=group, page_no=index) - elif isinstance(drawable, numbers_content.PlacedChart): + else: _add_chart( doc, drawable.chart, diff --git a/docling/backend/latex/backend.py b/docling/backend/latex/backend.py index 2fedeaaee6..69875728cc 100644 --- a/docling/backend/latex/backend.py +++ b/docling/backend/latex/backend.py @@ -10,6 +10,7 @@ from docling_core.types.doc import DocItemLabel, DoclingDocument, NodeItem from docling_core.types.doc.document import Formatting +from typing_extensions import override from docling.backend.abstract_backend import DeclarativeDocumentBackend from docling.backend.latex.handlers.environments import EnvironmentHandlerMixin @@ -78,9 +79,10 @@ def __init__( self._input_stack: set[str] = set() self._tectonic_engine: TectonicEngine | None = None self._tikz_executor: concurrent.futures.ThreadPoolExecutor | None = None - self._tikz_futures: list[concurrent.futures.Future] = [] + self._tikz_futures: list[concurrent.futures.Future[None]] = [] self.latex_text = decode_latex_content(self.path_or_stream) + @override def is_valid(self) -> bool: text = self.latex_text.strip() if not text: @@ -88,10 +90,12 @@ def is_valid(self) -> bool: return True @classmethod + @override def supports_pagination(cls) -> bool: return False @classmethod + @override def supported_formats(cls) -> set[InputFormat]: return {InputFormat.LATEX} @@ -127,7 +131,7 @@ def _do_parse_and_process(self, doc: DoclingDocument) -> DoclingDocument: workers = max(1, (os.cpu_count() or 2) - 1) self._tikz_executor = concurrent.futures.ThreadPoolExecutor(max_workers=workers) - self._tikz_futures = [] # type: ignore + self._tikz_futures = [] try: self._extract_custom_macros(nodes) @@ -150,6 +154,7 @@ def _do_parse_and_process(self, doc: DoclingDocument) -> DoclingDocument: return doc + @override def convert(self) -> DoclingDocument: doc = DoclingDocument(name=self.file.stem) opts = cast(LatexBackendOptions, self.options) @@ -191,6 +196,7 @@ def _worker_fn(): return doc + @override def _process_nodes( self, nodes, diff --git a/docling/backend/latex/engines/tectonic.py b/docling/backend/latex/engines/tectonic.py index bbaa2c3853..4b4216b611 100644 --- a/docling/backend/latex/engines/tectonic.py +++ b/docling/backend/latex/engines/tectonic.py @@ -12,6 +12,7 @@ from docling_core.types.doc.document import ImageRef from PIL import Image, ImageChops +from typing_extensions import override from docling.backend.latex.engines.base import RenderEngine from docling.backend.latex.utils.encoding import decode_latex_content @@ -88,6 +89,7 @@ def __init__( self._is_available = False self.install() + @override def is_available(self) -> bool: return self._is_available @@ -283,6 +285,7 @@ def _build_command(self, tex_file: Path) -> list[str]: cmd.append(str(tex_file)) return cmd + @override def render( self, tikz_code: str, preamble: str = "", source_root: Path | None = None ) -> ImageRef | None: diff --git a/docling/backend/latex/handlers/environments.py b/docling/backend/latex/handlers/environments.py index b9d72f0d01..2cf808b05d 100644 --- a/docling/backend/latex/handlers/environments.py +++ b/docling/backend/latex/handlers/environments.py @@ -32,7 +32,7 @@ ) try: # pragma: no cover - import-time guard - from pylatexenc.latexwalker import LatexEnvironmentNode, LatexMacroNode + from pylatexenc.latexwalker import LatexEnvironmentNode, LatexMacroNode, LatexNode except ImportError: pass # guarded by LatexDocumentBackend.__init__ @@ -250,7 +250,7 @@ def set_code_fallback(): ) def _extract_tikzpicture_atomic(self, node: LatexEnvironmentNode) -> str | None: - raw = node.latex_verbatim() + raw = str(node.latex_verbatim()) if _TIKZ_END_PATTERN.search(raw) is None: return None if not self._validate_tikz_nodelist(node.nodelist, 0): @@ -265,7 +265,7 @@ def _validate_tikz_nodelist(self, nodes, depth: int = 0) -> bool: for node in nodes: if isinstance(node, LatexEnvironmentNode) and node.envname == "tikzpicture": - nested_raw = node.latex_verbatim() + nested_raw = str(node.latex_verbatim()) if _TIKZ_END_PATTERN.search(nested_raw) is None: return False @@ -309,7 +309,7 @@ def _process_list( list_group = doc.add_group(parent=parent, name="list", label=GroupLabel.LIST) items = [] - current_item: list = [] + current_item: list[LatexNode] = [] if node.nodelist is not None: for n in node.nodelist: @@ -347,7 +347,7 @@ def _process_bibliography( ) items = [] - current_item: list = [] + current_item: list[LatexNode] = [] current_key = "" if node.nodelist is not None: diff --git a/docling/backend/latex/handlers/macros.py b/docling/backend/latex/handlers/macros.py index cc6de91b96..148821b030 100644 --- a/docling/backend/latex/handlers/macros.py +++ b/docling/backend/latex/handlers/macros.py @@ -47,6 +47,8 @@ LatexEnvironmentNode, LatexGroupNode, LatexMacroNode, + LatexMathNode, + LatexNode, LatexWalker, LatexWalkerParseError, ) @@ -544,7 +546,7 @@ def _extract_macro_arg(self, node: LatexMacroNode) -> str: if arg: if hasattr(arg, "nodelist"): return self._nodes_to_text(arg.nodelist) - return arg.latex_verbatim().strip("{} ") + return str(arg.latex_verbatim()).strip("{} ") return "" def _extract_macro_arg_by_index(self, node: LatexMacroNode, index: int) -> str: @@ -554,15 +556,27 @@ def _extract_macro_arg_by_index(self, node: LatexMacroNode, index: int) -> str: if arg: if hasattr(arg, "nodelist"): return self._nodes_to_text(arg.nodelist) - return arg.latex_verbatim().strip("{} ") + return str(arg.latex_verbatim()).strip("{} ") return "" - def _extract_macro_arg_nodes(self, node: LatexMacroNode, index: int) -> list: + def _extract_macro_arg_nodes( + self, node: LatexMacroNode, index: int + ) -> list[LatexNode]: if node.nodeargd and node.nodeargd.argnlist: if 0 <= index < len(node.nodeargd.argnlist): arg = node.nodeargd.argnlist[index] - if arg and hasattr(arg, "nodelist"): - return arg.nodelist + if isinstance( + arg, (LatexGroupNode, LatexEnvironmentNode, LatexMathNode) + ): + nodes = arg.nodelist + if nodes is None: + return [] + assert isinstance(nodes, list) + typed_nodes: list[LatexNode] = [] + for child in nodes: + assert isinstance(child, LatexNode) + typed_nodes.append(child) + return typed_nodes return [] def _extract_all_macro_args_inline(self, node: LatexMacroNode) -> str: diff --git a/docling/backend/latex/handlers/math.py b/docling/backend/latex/handlers/math.py index 5e0c76e741..97ed923aea 100644 --- a/docling/backend/latex/handlers/math.py +++ b/docling/backend/latex/handlers/math.py @@ -50,7 +50,7 @@ def _clean_math(self, latex_str: str, env_name: str) -> str: pattern = rf"\\begin\{{{re.escape(env_name)}\}}(.*?)\\end\{{{re.escape(env_name)}\}}" match = re.search(pattern, latex_str, re.DOTALL) if match: - latex_str = match.group(1) + latex_str = str(match.group(1)) latex_str = latex_str.strip() diff --git a/docling/backend/latex/utils/encoding.py b/docling/backend/latex/utils/encoding.py index 1d7bd39214..6646269fb6 100644 --- a/docling/backend/latex/utils/encoding.py +++ b/docling/backend/latex/utils/encoding.py @@ -26,7 +26,7 @@ def decode_latex_content(path_or_stream: Union[BytesIO, Path]) -> str: if not latex_text: _log.warning("Failed to decode LaTeX content, using replacement mode") latex_text = raw_bytes.decode("utf-8", errors="replace") - elif isinstance(path_or_stream, Path): + else: for encoding in _ENCODINGS: try: with open(path_or_stream, encoding=encoding) as f: diff --git a/docling/backend/latex/utils/table.py b/docling/backend/latex/utils/table.py index ed23f64b87..fcc230a70a 100644 --- a/docling/backend/latex/utils/table.py +++ b/docling/backend/latex/utils/table.py @@ -6,6 +6,7 @@ from typing import TYPE_CHECKING, Callable, List, Optional from docling_core.types.doc.document import TableCell, TableData +from pydantic import PrivateAttr from docling.backend.latex.constants import ( MACROS_ESCAPED, @@ -22,6 +23,7 @@ LatexCharsNode, LatexEnvironmentNode, LatexMacroNode, + LatexNode, LatexWalker, LatexWalkerParseError, ) @@ -29,6 +31,12 @@ pass # guarded by LatexDocumentBackend.__init__ +class _ParsedTableCell(TableCell): + _col_span: int = PrivateAttr(1) + _row_span: int = PrivateAttr(1) + _is_placeholder: bool = PrivateAttr(False) + + class TableHelperMixin: if TYPE_CHECKING: @@ -118,11 +126,11 @@ def _process_table_macro_node( def _parse_table(self, node: LatexEnvironmentNode) -> TableData | None: rows = [] current_row = [] - current_cell_nodes: list = [] + current_cell_nodes: list[LatexNode] = [] source_latex = node.latex_verbatim() - def parse_brace_args(text: str) -> list: + def parse_brace_args(text: str) -> list[str]: args = [] i = 0 while i < len(text): @@ -143,27 +151,27 @@ def parse_brace_args(text: str) -> list: def finish_cell(col_span: int = 1, row_span: int = 1): text = self._nodes_to_text(current_cell_nodes).strip() - cell = TableCell( + cell = _ParsedTableCell( text=text, start_row_offset_idx=0, end_row_offset_idx=0, start_col_offset_idx=0, end_col_offset_idx=0, ) - cell._col_span = col_span # type: ignore[attr-defined] - cell._row_span = row_span # type: ignore[attr-defined] + cell._col_span = col_span + cell._row_span = row_span current_row.append(cell) current_cell_nodes.clear() for _ in range(col_span - 1): - placeholder = TableCell( + placeholder = _ParsedTableCell( text="", start_row_offset_idx=0, end_row_offset_idx=0, start_col_offset_idx=0, end_col_offset_idx=0, ) - placeholder._is_placeholder = True # type: ignore[attr-defined] + placeholder._is_placeholder = True current_row.append(placeholder) def finish_row(): @@ -216,10 +224,10 @@ def finish_row(): for j in range(num_cols): if j < len(row): cell = row[j] - if getattr(cell, "_is_placeholder", False): + if cell._is_placeholder: continue else: - cell = TableCell( + cell = _ParsedTableCell( text="", start_row_offset_idx=0, end_row_offset_idx=0, @@ -230,8 +238,8 @@ def finish_row(): cell.start_row_offset_idx = i cell.start_col_offset_idx = j - col_span = getattr(cell, "_col_span", 1) - row_span = getattr(cell, "_row_span", 1) + col_span = cell._col_span + row_span = cell._row_span cell.end_row_offset_idx = i + row_span cell.end_col_offset_idx = j + col_span diff --git a/docling/backend/latex/utils/text.py b/docling/backend/latex/utils/text.py index 113ad75671..3cb390d110 100644 --- a/docling/backend/latex/utils/text.py +++ b/docling/backend/latex/utils/text.py @@ -117,7 +117,9 @@ def _extract_verbatim_content(self, latex_str: str, env_name: str) -> str: return re.sub(r"\A(?:[ \t]*\r?\n)+", "", match.group(1)).rstrip() return latex_str - def _macro_node_to_text(self, node: LatexMacroNode, following_nodes) -> tuple: + def _macro_node_to_text( + self, node: LatexMacroNode, following_nodes + ) -> tuple[str, int]: """Return ``(text, consumed_following)`` for a single macro node.""" consumed = 0 if node.macroname in (MACROS_TEXT_FORMATTING | MACROS_TEXT_STYLE): @@ -130,7 +132,7 @@ def _macro_node_to_text(self, node: LatexMacroNode, following_nodes) -> tuple: return (self._nodes_to_text(text_arg.nodelist), consumed) return ("", consumed) if node.macroname in MACROS_CITATION: - return (node.latex_verbatim(), consumed) + return (str(node.latex_verbatim()), consumed) if node.macroname == "\\": return ("\n", consumed) if node.macroname in ["~"]: @@ -143,7 +145,7 @@ def _macro_node_to_text(self, node: LatexMacroNode, following_nodes) -> tuple: return (f"{opt_text}: ", consumed) return ("", consumed) if node.macroname in MACROS_ESCAPED: - return (node.macroname, consumed) + return (str(node.macroname), consumed) if node.macroname in self._custom_macros: expansion, consumed = self._expand_custom_macro_invocation( node, following_nodes diff --git a/docling/backend/managed_pdfium_backend.py b/docling/backend/managed_pdfium_backend.py index 7b95f8dbe1..5cc6acb109 100644 --- a/docling/backend/managed_pdfium_backend.py +++ b/docling/backend/managed_pdfium_backend.py @@ -8,6 +8,8 @@ from pathlib import Path from typing import TYPE_CHECKING, Optional, Union +from typing_extensions import override + from docling.backend.pdf_backend import PdfDocumentBackend, PdfPageBackend from docling.datamodel.backend_options import PdfBackendOptions @@ -33,6 +35,7 @@ def __init__( def _close_native_document(self) -> None: pass + @override def unload(self) -> None: if self._closed: return @@ -51,6 +54,7 @@ def __init__(self) -> None: def _close_native_page(self) -> None: pass + @override def unload(self) -> None: if self._closed: return diff --git a/docling/backend/md_backend.py b/docling/backend/md_backend.py index 969529e9bf..01d1f350d0 100644 --- a/docling/backend/md_backend.py +++ b/docling/backend/md_backend.py @@ -36,6 +36,7 @@ from docling.backend.html_backend import HTMLDocumentBackend from docling.backend.utils.image_resource_loader import ImageResourceLoader from docling.datamodel.backend_options import ( + BaseBackendOptions, HTMLBackendOptions, MarkdownBackendOptions, ) @@ -102,7 +103,7 @@ class _ListItemCreationPayload(BaseModel): ] -def _only_plain_line_breaks(children: list) -> bool: +def _only_plain_line_breaks(children: list[marko.element.Element]) -> bool: """Return True when children consist solely of RawText/Literal runs separated by LineBreak nodes (soft or hard), with at least one break present. @@ -280,7 +281,7 @@ def __init__( self, in_doc: InputDocument, path_or_stream: Union[BytesIO, Path], - options: Optional[MarkdownBackendOptions] = None, + options: Optional[BaseBackendOptions] = None, ): # Raised first so a missing optional dependency gives an actionable # message rather than a NameError when marko is dereferenced below. @@ -288,6 +289,8 @@ def __init__( raise ImportError(_INSTALL_HINT) from _MARKO_IMPORT_ERROR if options is None: options = MarkdownBackendOptions() + if not isinstance(options, MarkdownBackendOptions): + raise TypeError("Expected MarkdownBackendOptions.") super().__init__(in_doc, path_or_stream, options) _log.debug("Starting MarkdownDocumentBackend...") @@ -416,7 +419,7 @@ def _create_list_item( marker: str = "", formatting: Optional[Formatting] = None, hyperlink: Optional[Union[AnyUrl, Path]] = None, - ): + ) -> ListItem: item = doc.add_list_item( text=text, enumerated=enumerated, @@ -425,6 +428,7 @@ def _create_list_item( formatting=formatting, hyperlink=hyperlink, ) + assert isinstance(item, ListItem) return item def _create_heading_item( @@ -435,7 +439,7 @@ def _create_heading_item( level: int, formatting: Optional[Formatting] = None, hyperlink: Optional[Union[AnyUrl, Path]] = None, - ): + ) -> TextItem: if level == 1: item = doc.add_title( text=text, @@ -451,6 +455,7 @@ def _create_heading_item( formatting=formatting, hyperlink=hyperlink, ) + assert isinstance(item, TextItem) return item def _flush_creation_stack( @@ -487,12 +492,12 @@ def _flush_creation_stack( hyperlink=hyperlink, ) if parent_ref: - list_last_item_by_ref[parent_ref] = cast(ListItem, parent_item) + list_last_item_by_ref[parent_ref] = parent_item list_item_counter_by_ref[parent_ref] = ( list_item_counter_by_ref.get(parent_ref, 0) + 1 ) - elif isinstance(to_create, _HeadingCreationPayload): + else: # Not keeping as parent_item as logic for correctly tracking # that not implemented yet (section components not captured # as heading children in marko) @@ -621,7 +626,7 @@ def _iterate_elements( # noqa: C901 hyperlink=hyperlink, ) if parent_ref: - list_last_item_by_ref[parent_ref] = cast(ListItem, parent_item) + list_last_item_by_ref[parent_ref] = parent_item list_item_counter_by_ref[parent_ref] = ( list_item_counter_by_ref.get(parent_ref, 0) + 1 ) @@ -637,13 +642,16 @@ def _iterate_elements( # noqa: C901 fig_caption: Optional[TextItem] = None if element.title is not None and element.title != "": title = unescape(element.title) - fig_caption = doc.add_text( + caption_item = doc.add_text( label=DocItemLabel.CAPTION, text=title, formatting=formatting, hyperlink=hyperlink, ) + assert isinstance(caption_item, TextItem) + fig_caption = caption_item + image_ref = self._load_image_ref(element.dest) doc.add_picture(parent=parent_item, image=image_ref, caption=fig_caption) @@ -659,9 +667,9 @@ def _iterate_elements( # noqa: C901 elif isinstance(element, marko.inline.Link): _log.debug(" - Link: %s", element.children) - hyperlink = TypeAdapter(Optional[Union[AnyUrl, Path]]).validate_python( - element.dest - ) + hyperlink = TypeAdapter[AnyUrl | Path | None]( + AnyUrl | Path | None + ).validate_python(element.dest) elif isinstance(element, marko.inline.RawText | marko.inline.Literal): _log.debug(" - RawText/Literal: %s", element.children) @@ -865,10 +873,10 @@ def _iterate_elements( # noqa: C901 ) # Iterate through the element's children (if any) - if hasattr(element, "children") and not isinstance( + if isinstance(element_children, list) and not isinstance( element, processed_block_types ): - for child in element.children: + for child in element_children: if ( isinstance(element, marko.block.ListItem) and isinstance(child, marko.block.List) @@ -927,22 +935,27 @@ def _load_image_ref(self, dest: str) -> Optional[ImageRef]: ) return self._get_image_loader().load_image_ref(dest, base_path) + @override def is_valid(self) -> bool: return self.valid + @override def unload(self): if isinstance(self.path_or_stream, BytesIO): self.path_or_stream.close() self.path_or_stream = None @classmethod + @override def supports_pagination(cls) -> bool: return False @classmethod + @override def supported_formats(cls) -> set[InputFormat]: return {InputFormat.MD} + @override def convert(self) -> DoclingDocument: _log.debug("converting Markdown...") diff --git a/docling/backend/mets_gbs_backend.py b/docling/backend/mets_gbs_backend.py index 5b793b5e43..a73730540f 100644 --- a/docling/backend/mets_gbs_backend.py +++ b/docling/backend/mets_gbs_backend.py @@ -23,6 +23,7 @@ from lxml import etree from PIL import Image from PIL.Image import Image as PILImage +from typing_extensions import override from docling.backend.pdf_backend import PdfDocumentBackend, PdfPageBackend from docling.datamodel.backend_options import MetsGbsBackendOptions @@ -68,9 +69,11 @@ def __init__( self.valid = parsed_page is not None and page_im is not None @property + @override def page_no(self) -> int: return self._page_no + 1 + @override def is_valid(self) -> bool: return self.valid @@ -102,6 +105,7 @@ def _require_image(self) -> PILImage: ) return self._im + @override def get_text_in_rect(self, bbox: BoundingBox) -> str: dpage = self._require_page() # Find intersecting cells on the page @@ -128,12 +132,15 @@ def get_text_in_rect(self, bbox: BoundingBox) -> str: return text_piece + @override def get_segmented_page(self) -> SegmentedPdfPage | None: return self._dpage + @override def get_text_cells(self) -> Iterable[TextCell]: return self._require_page().textline_cells + @override def get_bitmap_rects(self, scale: float = 1) -> Iterable[BoundingBox]: dpage = self._require_page() AREA_THRESHOLD = 0 # 32 * 32 @@ -150,6 +157,7 @@ def get_bitmap_rects(self, scale: float = 1) -> Iterable[BoundingBox]: yield cropbox + @override def get_page_image( self, scale: float = 1, cropbox: BoundingBox | None = None ) -> Image.Image: @@ -175,10 +183,12 @@ def get_page_image( ).crop(cropbox.scaled(scale=scale).as_tuple()) return image + @override def get_size(self) -> Size: dpage = self._require_page() return Size(width=dpage.dimension.width, height=dpage.dimension.height) + @override def unload(self) -> None: self._im = None self._dpage = None @@ -521,31 +531,36 @@ def _parse_page( textline_cells=line_cells, char_cells=[], word_cells=word_cells, - has_textlines=True, has_words=True, has_chars=False, ) return page, im + @override def page_count(self) -> int: return len(self.page_map) + @override def load_page(self, page_no: int) -> MetsGbsPageBackend: # TODO: is this thread-safe? page, im = self._parse_page(page_no) return MetsGbsPageBackend(parsed_page=page, page_im=im, page_no=page_no) + @override def is_valid(self) -> bool: return self.root_mets is not None and self.page_count() > 0 @classmethod + @override def supported_formats(cls) -> set[InputFormat]: return {InputFormat.METS_GBS} @classmethod + @override def supports_pagination(cls) -> bool: return True + @override def unload(self) -> None: super().unload() self._tar.close() diff --git a/docling/backend/msexcel_backend.py b/docling/backend/msexcel_backend.py index 0743b6687b..c76f489428 100644 --- a/docling/backend/msexcel_backend.py +++ b/docling/backend/msexcel_backend.py @@ -52,7 +52,7 @@ crop_whitespace, get_docx_to_pdf_converter, ) -from docling.datamodel.backend_options import MsExcelBackendOptions +from docling.datamodel.backend_options import BaseBackendOptions, MsExcelBackendOptions from docling.datamodel.base_models import FormatToMimeType, InputFormat from docling.datamodel.document import InputDocument from docling.exceptions import DocumentLoadError @@ -225,14 +225,14 @@ def __init__(self, sheet: Worksheet) -> None: ), ) min_row = ( - merged_range.min_row + int(merged_range.min_row) if min_row is None - else min(min_row, merged_range.min_row) + else min(min_row, int(merged_range.min_row)) ) min_col = ( - merged_range.min_col + int(merged_range.min_col) if min_col is None - else min(min_col, merged_range.min_col) + else min(min_col, int(merged_range.min_col)) ) max_row = max(max_row, merged_range.max_row) max_col = max(max_col, merged_range.max_col) @@ -337,7 +337,7 @@ def __init__( self, in_doc: InputDocument, path_or_stream: BytesIO | Path, - options: MsExcelBackendOptions | None = None, + options: BaseBackendOptions | None = None, ) -> None: """Initialize the MsExcelDocumentBackend object. @@ -355,6 +355,8 @@ def __init__( path_or_stream = convert_to_modern_format(path_or_stream, "xls", "xlsx") if options is None: options = MsExcelBackendOptions() + if not isinstance(options, MsExcelBackendOptions): + raise TypeError("Expected MsExcelBackendOptions.") super().__init__(in_doc, path_or_stream, options) self.page_range = in_doc.limits.page_range @@ -363,7 +365,7 @@ def __init__( self.parent: GroupItem | None = None # Lazy-initialized LibreOffice converter for EMF/WMF images - self.xlsx_to_pdf_converter: Callable | None = None + self.xlsx_to_pdf_converter: Callable[[Path, Path], None] | None = None self.xlsx_to_pdf_converter_init: bool = False self.workbook = None @@ -448,6 +450,8 @@ def _parse_threaded_comments( except Exception as e: _log.debug(f"Could not parse person.xml: {e}") + if self.workbook is None: + return {} sheet_num = next( ( i @@ -560,6 +564,8 @@ def convert(self) -> DoclingDocument: Returns: The DoclingDocument object representing the Excel workbook. """ + assert self.input_format is not None + origin = DocumentOrigin( filename=self.file.name or "file.xlsx", mimetype=FormatToMimeType[self.input_format][0], @@ -659,8 +665,7 @@ def _convert_sheet( doc = self._find_tables_in_sheet(doc, sheet, page_no) doc = self._find_images_in_sheet(doc, sheet, page_no) # Charts can be on both Worksheet and Chartsheet objects - if isinstance(sheet, (Worksheet, Chartsheet)): - doc = self._find_chart_in_sheet(doc, sheet, page_no) + doc = self._find_chart_in_sheet(doc, sheet, page_no) self._sort_sheet_children_by_position(doc, page_no) return doc @@ -691,7 +696,7 @@ def _top_row(ref: Any) -> float: return float("inf") for prov in getattr(item, "prov", []): if prov.page_no == page_no: - return prov.bbox.t + return float(prov.bbox.t) return float("inf") sheet_group.children.sort(key=_top_row) @@ -811,7 +816,9 @@ def _find_tables_in_sheet( # Extract comments and link them to cells for (row, col), thread in comment_map.items(): cell_item = self._find_cell_item(doc, page_no, row, col) - targets = [cell_item] if cell_item else None + targets: list[DocItem | tuple[DocItem, tuple[int, int]]] | None = ( + [cell_item] if cell_item else None + ) comment_group = doc.add_group( label=GroupLabel.COMMENT_SECTION, @@ -1209,7 +1216,7 @@ def _anchor_to_tuple(anchor: Any) -> tuple[int, int, int, int]: ) return (0, 0, 0, 0) - def _get_libreoffice_converter(self) -> Callable | None: + def _get_libreoffice_converter(self) -> Callable[[Path, Path], None] | None: """Lazily initialize and return a LibreOffice converter callable. The converter accepts ``(input_path: Path, output_path: Path)`` and @@ -1289,10 +1296,11 @@ def _find_unsupported_images_in_sheet( Returns: The updated DoclingDocument. """ - # drawing rels are stored on the worksheet object by openpyxl + # openpyxl initializes these real private fields; its stubs omit them. + # Drawing relationships are stored on the worksheet object. drawing_paths: list[str] = [ rel.target - for rel in sheet._rels.find(SpreadsheetDrawing._rel_type) # type: ignore[attr-defined] + for rel in sheet._rels.find(SpreadsheetDrawing._rel_type) # ty: ignore[unresolved-attribute] ] if not drawing_paths: return doc @@ -1302,8 +1310,11 @@ def _find_unsupported_images_in_sheet( if isinstance(self.path_or_stream, BytesIO): self.path_or_stream.seek(0) + path_or_stream = self.path_or_stream + if path_or_stream is None: + return doc try: - with ZipFile(self.path_or_stream, "r") as zf: + with ZipFile(path_or_stream, "r") as zf: if _has_unsafe_zip_paths(zf.namelist()): _log.warning( "Skipping EMF/WMF scan: XLSX archive contains unsafe ZIP paths" @@ -1353,7 +1364,10 @@ def _process_drawing_for_unsupported_images( return doc deps = get_dependents(zf, rels_path) - for rel in drawing._blip_rels: + if drawing is None: + return doc + # openpyxl exposes this property; its stubs omit it. + for rel in drawing._blip_rels: # ty: ignore[unresolved-attribute] dep = deps.get(rel.embed) if dep.Type != IMAGE_NS or dep.target not in zf.namelist(): continue @@ -1410,13 +1424,15 @@ def _find_images_in_sheet( """ if self.workbook is not None: content_layer = self._get_sheet_content_layer(sheet) - # Images that PIL can read are already loaded by openpyxl into sheet._images - for item in sheet._images: # type: ignore[attr-defined] + # openpyxl initializes _images; its stubs omit this real field. + # These images can already be read by PIL. + for item in sheet._images: # ty: ignore[unresolved-attribute] try: image: Image = cast(Image, item) ref = image.ref + # Pillow also accepts read-only streams, which its IO stubs omit. pil_image = ( - ref if isinstance(ref, PILImage.Image) else PILImage.open(ref) + ref if isinstance(ref, PILImage.Image) else PILImage.open(ref) # ty: ignore[invalid-argument-type] ) doc.add_picture( parent=self.parent, @@ -1467,7 +1483,8 @@ def _find_chart_in_sheet( return doc content_layer = self._get_sheet_content_layer(sheet) - charts = sheet._charts # type: ignore[attr-defined] + # openpyxl initializes _charts on both sheet types; its stubs omit it. + charts = sheet._charts # ty: ignore[unresolved-attribute] if not charts: return doc @@ -1863,10 +1880,10 @@ def _ref_formula(data_source: Any) -> str | None: return None num_ref = getattr(data_source, "numRef", None) if num_ref is not None and num_ref.f: - return num_ref.f + return str(num_ref.f) str_ref = getattr(data_source, "strRef", None) if str_ref is not None and str_ref.f: - return str_ref.f + return str(str_ref.f) return None def _resolve_reference(self, ref: str) -> list[str]: @@ -1897,7 +1914,10 @@ def _resolve_reference(self, ref: str) -> list[str]: sheet_part = sheet_part[1:-1].replace("''", "'") sheet_name = sheet_part else: - sheet_name, cell_range = self.workbook.active.title, ref + active_sheet = self.workbook.active + if active_sheet is None: + return [] + sheet_name, cell_range = active_sheet.title, ref if sheet_name not in self.workbook.sheetnames: _log.debug("Chart references unknown sheet %r", sheet_name) return [] @@ -1910,6 +1930,8 @@ def _resolve_reference(self, ref: str) -> list[str]: _log.debug("Could not parse chart range %r", cell_range) return [] + if min_row is None or max_row is None or min_col is None or max_col is None: + return [] values: list[str] = [] for row in range(min_row, max_row + 1): for col in range(min_col, max_col + 1): @@ -1977,7 +1999,7 @@ def _find_cell_item( return None @staticmethod - def _get_sheet_content_layer(sheet: Worksheet) -> ContentLayer | None: + def _get_sheet_content_layer(sheet: Worksheet | Chartsheet) -> ContentLayer | None: return ( None if sheet.sheet_state == Worksheet.SHEETSTATE_VISIBLE diff --git a/docling/backend/mspowerpoint_backend.py b/docling/backend/mspowerpoint_backend.py index 1810356672..4dae31b564 100644 --- a/docling/backend/mspowerpoint_backend.py +++ b/docling/backend/mspowerpoint_backend.py @@ -10,7 +10,7 @@ from io import BytesIO from pathlib import Path from tempfile import mkdtemp -from typing import Any, Callable, Final, Iterable, Iterator, Optional, Union +from typing import Any, Callable, Final, Iterable, Iterator, Optional, TypedDict, Union from docling_core.types.doc import ( BoundingBox, @@ -22,6 +22,7 @@ ImageRef, ListGroup, ListItem, + NodeItem, PictureClassificationLabel, PictureClassificationMetaField, PictureClassificationPrediction, @@ -46,11 +47,22 @@ crop_whitespace, get_docx_to_pdf_converter, ) -from docling.datamodel.backend_options import MsPowerpointBackendOptions +from docling.datamodel.backend_options import ( + BaseBackendOptions, + MsPowerpointBackendOptions, +) from docling.datamodel.base_models import FormatToMimeType, InputFormat from docling.datamodel.document import InputDocument from docling.exceptions import DocumentLoadError + +class _ListMarker(TypedDict): + is_list: bool | None + kind: str | None + detail: str | None + level: int + + _log = logging.getLogger(__name__) _PPTX_AVAILABLE: bool = False @@ -201,7 +213,7 @@ def __init__( self, in_doc: InputDocument, path_or_stream: Union[BytesIO, Path], - options: Optional[MsPowerpointBackendOptions] = None, + options: BaseBackendOptions | None = None, ) -> None: if not _PPTX_AVAILABLE: raise ImportError(_INSTALL_HINT) from _PPTX_IMPORT_ERROR @@ -209,11 +221,13 @@ def __init__( path_or_stream = convert_to_modern_format(path_or_stream, "ppt", "pptx") if options is None: options = MsPowerpointBackendOptions() + if not isinstance(options, MsPowerpointBackendOptions): + raise TypeError("MsPowerpointBackendOptions are required.") super().__init__(in_doc, path_or_stream, options) - self.path_or_stream: Union[BytesIO, Path] = path_or_stream + self.path_or_stream: BytesIO | Path | None = path_or_stream self.page_range = in_doc.limits.page_range - self.pptx_to_pdf_converter: Optional[Callable] = None + self.pptx_to_pdf_converter: Callable[[Path, Path], None] | None = None self.pptx_to_pdf_converter_init: bool = False self._render_charts: bool = False self._metafile_hint_emitted: bool = False @@ -221,10 +235,10 @@ def __init__( self.pptx_obj: Optional[presentation.Presentation] = None self.valid: bool = False try: - if isinstance(self.path_or_stream, BytesIO): - self.pptx_obj = Presentation(self.path_or_stream) - elif isinstance(self.path_or_stream, Path): - self.pptx_obj = Presentation(str(self.path_or_stream)) + if isinstance(path_or_stream, BytesIO): + self.pptx_obj = Presentation(path_or_stream) + else: + self.pptx_obj = Presentation(str(path_or_stream)) self.valid = True except Exception as e: @@ -234,6 +248,7 @@ def __init__( return + @override def page_count(self) -> int: if self.is_valid(): assert self.pptx_obj is not None @@ -269,6 +284,8 @@ def convert(self) -> DoclingDocument: Returns: The parsed document. """ + assert self.input_format is not None + origin = DocumentOrigin( filename=self.file.name or "file", mimetype=FormatToMimeType[self.input_format][0], @@ -300,7 +317,7 @@ def _generate_prov( top = 0 width = slide_size.width height = slide_size.height - shape_bbox = [left, top, left + width, top + height] + shape_bbox = (float(left), float(top), float(left + width), float(top + height)) # python-pptx reports left and top as EMU from the slide's top-left, with # y growing downward, so the tuple above is in TOPLEFT order. Tagging it # BOTTOMLEFT does not convert it: BoundingBox.from_tuple unpacks @@ -387,12 +404,14 @@ def _parse_bullet_from_paragraph_properties( # Bullet character buChar = pPr.find("a:buChar", namespaces=self.NAMESPACES) if buChar is not None: - return (True, "buChar", buChar.get("char")) + detail = buChar.get("char") + return (True, "buChar", str(detail) if detail is not None else None) # Auto numbering buAuto = pPr.find("a:buAutoNum", namespaces=self.NAMESPACES) if buAuto is not None: - return (True, "buAutoNum", buAuto.get("type")) + detail = buAuto.get("type") + return (True, "buAutoNum", str(detail) if detail is not None else None) # Picture bullet buBlip = pPr.find("a:buBlip", namespaces=self.NAMESPACES) @@ -572,7 +591,7 @@ def _is_list_item(self, paragraph) -> tuple[bool, str]: else: return (False, "None") - def _get_effective_list_marker(self, shape, paragraph) -> dict: + def _get_effective_list_marker(self, shape, paragraph) -> _ListMarker: """Return a dictionary describing the effective list marker for a paragraph. List marker information can come from several sources: direct paragraph @@ -620,7 +639,7 @@ def _get_effective_list_marker(self, shape, paragraph) -> dict: } # 3) Layout placeholder lstStyle (if this is a placeholder) - layout_result = None + layout_result: _ListMarker | None = None if shape.is_placeholder: idx = shape.placeholder_format.idx layout = shape.part.slide.slide_layout @@ -1067,7 +1086,7 @@ def _chart_title_text(chart: Any) -> Optional[str]: """ if not chart.has_title: return None - text = chart.chart_title.text_frame.text.strip() + text = str(chart.chart_title.text_frame.text).strip() return text or None @staticmethod @@ -1254,7 +1273,7 @@ def _handle_chart(self, shape, parent_slide, slide_ind, doc, slide_size): ) return - def _get_libreoffice_converter(self) -> Optional[Callable]: + def _get_libreoffice_converter(self) -> Callable[[Path, Path], None] | None: """Lazily initialize and return a LibreOffice converter callable. The converter accepts ``(input_path, output_path)`` and converts the @@ -1454,6 +1473,7 @@ def _walk_linear( # Units of size in PPTX by default are EMU units (English Metric Units) slide_width = pptx_obj.slide_width slide_height = pptx_obj.slide_height + assert slide_width is not None and slide_height is not None self._render_charts = ( isinstance(self.options, MsPowerpointBackendOptions) @@ -1464,7 +1484,7 @@ def _walk_linear( self._render_charts = False max_levels = 10 - parents = {} # type: ignore + parents: dict[int, NodeItem | None] = {} for i in range(max_levels): parents[i] = None diff --git a/docling/backend/msword_backend.py b/docling/backend/msword_backend.py index dfcb50c99c..a6dcb7c890 100644 --- a/docling/backend/msword_backend.py +++ b/docling/backend/msword_backend.py @@ -11,7 +11,7 @@ from copy import deepcopy from io import BytesIO from pathlib import Path -from typing import Any, Callable, Final +from typing import Any, Callable, Final, TypedDict from urllib.parse import urlparse from docling_core.types.doc import ( @@ -50,7 +50,7 @@ get_pil_from_dml_docx, ) from docling.backend.docx.latex.omml import oMath2Latex -from docling.datamodel.backend_options import MsWordBackendOptions +from docling.datamodel.backend_options import BaseBackendOptions, MsWordBackendOptions from docling.datamodel.base_models import FormatToMimeType from docling.datamodel.document import InputDocument, InputFormat from docling.exceptions import DocumentLoadError, SecurityError @@ -566,7 +566,11 @@ def _remove_fragment_only_rels(content: bytes) -> bytes: parent = rel.getparent() if parent is not None: parent.remove(rel) - return etree.tostring(root, xml_declaration=True, encoding="UTF-8", standalone=True) + serialized = etree.tostring( + root, xml_declaration=True, encoding="UTF-8", standalone=True + ) + assert isinstance(serialized, bytes) + return serialized def _sanitize_docx(archive: zipfile.ZipFile) -> BytesIO: @@ -646,6 +650,13 @@ def _normalize_strict_ooxml(archive: zipfile.ZipFile) -> BytesIO: return normalized +class _ListHistory(TypedDict): + names: list[str | None] + levels: list[int | None] + numids: list[int | None] + indents: list[int | None] + + class MsWordDocumentBackend(DeclarativeDocumentBackend): """Backend for parsing Word documents (DOCX and DOC files). @@ -771,12 +782,14 @@ def __init__( self, in_doc: InputDocument, path_or_stream: BytesIO | Path, - options: MsWordBackendOptions | None = None, + options: BaseBackendOptions | None = None, ) -> None: if not _DOCX_AVAILABLE: raise ImportError(_INSTALL_HINT) from _DOCX_IMPORT_ERROR if options is None: options = MsWordBackendOptions() + if not isinstance(options, MsWordBackendOptions): + raise TypeError("Expected MsWordBackendOptions.") if in_doc.format in {InputFormat.DOC, InputFormat.RTF}: path_or_stream = convert_to_modern_format( path_or_stream, in_doc.format.value, "docx" @@ -789,7 +802,7 @@ def __init__( ) # self.initialise(path_or_stream) # Word file: - self.path_or_stream: BytesIO | Path = path_or_stream + self.path_or_stream: BytesIO | Path | None = path_or_stream self.valid: bool = False # Initialise the parents for the hierarchy self.max_levels: int = 10 @@ -800,7 +813,7 @@ def __init__( self.equation_bookends: str = "{EQ}" # Track processed textbox elements to avoid duplication self.processed_textbox_elements: set[etree._Element] = set() - self.docx_to_pdf_converter: Callable | None = None + self.docx_to_pdf_converter: Callable[[Path, Path], None] | None = None self.docx_to_pdf_converter_init = False self.display_drawingml_warning = True self._empty_docx_template: DocxDocument | None = None @@ -824,7 +837,7 @@ def __init__( # Set starting content layer self.content_layer = ContentLayer.BODY - self.history: dict[str, Any] = { + self.history: _ListHistory = { "names": [None], "levels": [None], "numids": [None], @@ -849,7 +862,7 @@ def __init__( self._default_paragraph_style: BaseStyle | None = None self.docx_obj = self.load_msword_file( - path_or_stream=self.path_or_stream, document_hash=self.document_hash + path_or_stream=path_or_stream, document_hash=self.document_hash ) if self.docx_obj: self.valid = True @@ -889,6 +902,8 @@ def convert(self) -> DoclingDocument: The parsed document. """ + assert self.input_format is not None + origin = DocumentOrigin( filename=self.file.name or "file", mimetype=FormatToMimeType[self.input_format][0], @@ -931,7 +946,7 @@ def load_msword_file( if has_fragment_rels: return Document(_sanitize_docx(archive)) return Document(str(path_or_stream)) - elif isinstance(path_or_stream, BytesIO): + else: with zipfile.ZipFile(path_or_stream) as archive: is_strict = _is_strict_ooxml(archive) has_fragment_rels = not is_strict and _has_fragment_only_rels( @@ -943,8 +958,6 @@ def load_msword_file( return Document(_sanitize_docx(archive)) path_or_stream.seek(0) return Document(path_or_stream) - else: - return None except SecurityError: raise except Exception as e: @@ -1014,21 +1027,22 @@ def _get_or_create_list_group( Returns: The list group to use (either reused or newly created). """ - if self._can_reuse_list_group(numid, parent): + cached_group = self.last_list_group + if cached_group is not None and self._can_reuse_list_group(numid, parent): # Reuse only if nothing but empty paragraphs (added when the list was # closed) follows the cached group in its parent. Otherwise the new # items would be placed before the intervening content, e.g. a table. container = parent if parent is not None else doc.body - trailing_empty: list[TextItem] = [] + trailing_empty: list[NodeItem] = [] for ref in reversed(container.children): item = ref.resolve(doc) if isinstance(item, TextItem) and not item.text.strip(): trailing_empty.append(item) continue - if item.self_ref == self.last_list_group.self_ref: + if item.self_ref == cached_group.self_ref: if trailing_empty: doc.delete_items(node_items=trailing_empty) - return self.last_list_group + return cached_group break list_gr = doc.add_list_group( @@ -1036,7 +1050,7 @@ def _get_or_create_list_group( parent=parent, content_layer=self.content_layer, ) - elem_ref.append(list_gr.get_ref()) + elem_ref.append(RefItem(cref=list_gr.self_ref)) # Update cache for potential future reuse self.last_list_group = list_gr @@ -1071,7 +1085,7 @@ def _isolated_list_context(self): This ensures that lists in different table cells are treated independently, even when they share the same numId. """ - saved_history = { + saved_history: _ListHistory = { "names": self.history["names"].copy(), "levels": self.history["levels"].copy(), "numids": self.history["numids"].copy(), @@ -1166,7 +1180,9 @@ def _in_textbox(elem): name="shape-text", content_layer=self.content_layer, ) - added_elements.append(shape_group.get_ref()) + added_elements.append( + RefItem(cref=shape_group.self_ref) + ) doc.add_text( label=DocItemLabel.TEXT, parent=shape_group, @@ -1291,15 +1307,16 @@ def _is_invisible_spacer(self, pil_image: Image.Image | None) -> bool: try: extrema = pil_image.getextrema() - if extrema is not None: - if pil_image.mode in ("RGBA", "LA"): + if pil_image.mode in ("RGBA", "LA"): + alpha = extrema[-1] + if isinstance(alpha, tuple): # extrema[-1] is the Alpha channel. If max alpha is 0, it is 100% invisible. - if extrema[-1][1] == 0: - return True - elif pil_image.mode == "RGB": - # If all channels are exactly 255, it is a pure white spacing box. - if extrema == ((255, 255), (255, 255), (255, 255)): + if alpha[1] == 0: return True + elif pil_image.mode == "RGB": + # If all channels are exactly 255, it is a pure white spacing box. + if extrema == ((255, 255), (255, 255), (255, 255)): + return True except Exception: pass # pragma: no cover @@ -1376,16 +1393,19 @@ def _style_numbering( if numId is None: numId_elem = style_numPr.find(f"{_W_NS_CLARK}numId") if numId_elem is not None: - numId = numId_elem.get(self.XML_KEY) + value = numId_elem.get(self.XML_KEY) + numId = str(value) if value is not None else None if ilvl is None: ilvl_elem = style_numPr.find(f"{_W_NS_CLARK}ilvl") if ilvl_elem is not None: - ilvl = ilvl_elem.get(self.XML_KEY) + value = ilvl_elem.get(self.XML_KEY) + ilvl = str(value) if value is not None else None if numId is not None and ilvl is not None: break # A malformed basedOn chain can hop to a style type that lacks # base_style; getattr keeps the walk safe. - style = getattr(style, "base_style", None) + base_style = getattr(style, "base_style", None) + style = base_style if isinstance(base_style, ParagraphStyle) else None depth += 1 # If numId is found but ilvl is not specified, default to level 0 @@ -1393,7 +1413,7 @@ def _style_numbering( ilvl = "0" return numId, ilvl - def _get_level_element(self, numid: int, ilvl: int) -> BaseOxmlElement | None: + def _get_level_element(self, numid: int, ilvl: int) -> etree._Element | None: """Find the level element from the numbering XML for a given numId and ilvl.""" try: numbering_part = None @@ -1431,9 +1451,11 @@ def _get_level_element(self, numid: int, ilvl: int) -> BaseOxmlElement | None: if abstract_num_element is None: return None - return abstract_num_element.find( + level = abstract_num_element.find( f".//w:lvl[@w:ilvl='{ilvl}']", namespaces=namespaces ) + assert level is None or isinstance(level, etree._Element) + return level except Exception as e: _log.debug(f"Error finding level element: {e}") return None @@ -1447,7 +1469,8 @@ def _get_level_num_fmt(self, numid: int, ilvl: int) -> str | None: num_fmt_element = lvl_element.find(".//w:numFmt", namespaces=namespaces) if num_fmt_element is None: return None - return num_fmt_element.get(self.XML_KEY) + value = num_fmt_element.get(self.XML_KEY) + return str(value) if value is not None else None def _get_start_value(self, numid: int, ilvl: int) -> int: """Read the start value from the abstractNum definition.""" @@ -1633,7 +1656,9 @@ def _last_child_item( container = parent if parent is not None else doc.body if not container.children: return None - return container.children[-1].resolve(doc) + child = container.children[-1].resolve(doc) + assert isinstance(child, NodeItem) + return child def _is_code_style(self, style: ParagraphStyle | None) -> bool: """Return True if a style marks its paragraphs as code. @@ -1657,7 +1682,8 @@ def _is_code_style(self, style: ParagraphStyle | None) -> bool: # A malformed basedOn chain can hop to a style type (e.g. a # numbering style) that lacks this attribute; getattr keeps # the walk safe. - style = getattr(style, "base_style", None) + base_style = getattr(style, "base_style", None) + style = base_style if isinstance(base_style, ParagraphStyle) else None depth += 1 return False @@ -1687,11 +1713,12 @@ def _effective_style_font(self, style: ParagraphStyle | None) -> str: return "" font_name = style.font.name if font_name: - return font_name.strip().lower() + return str(font_name).strip().lower() # A malformed basedOn chain can hop to a style type (e.g. a # numbering style) that lacks base_style; getattr keeps the # walk safe. - style = getattr(style, "base_style", None) + base_style = getattr(style, "base_style", None) + style = base_style if isinstance(base_style, ParagraphStyle) else None depth += 1 return "" @@ -1812,13 +1839,15 @@ def _get_label_and_level(self, paragraph: Paragraph) -> tuple[str, int | None]: if style is None: return "Normal", None - label: str = style.style_id - name: str = style.name or "" + label = str(style.style_id) + name = str(style.name or "") base_style_label: str | None = None base_style_name: str | None = None if isinstance(base_style := getattr(style, "base_style", None), ParagraphStyle): - base_style_label = base_style.style_id - base_style_name = base_style.name + base_style_label = str(base_style.style_id) + base_style_name = ( + str(base_style.name) if base_style.name is not None else None + ) if not label: return "Normal", None @@ -1928,16 +1957,15 @@ def _get_format_from_run( if not is_bold: try: - if run._element is not None: - b_tags = run._element.xpath(".//w:b") - for b in b_tags: - val = b.get(f"{_W_NS_CLARK}val") - if val not in ["0", "false"]: - is_bold = True - break + b_tags = run._element.xpath(".//w:b") + for b in b_tags: + val = b.get(f"{_W_NS_CLARK}val") + if val not in ["0", "false"]: + is_bold = True + break - if not is_bold and run._parent._element is not None: - pPr_b = run._parent._element.xpath("./w:pPr/w:rPr/w:b") + if not is_bold and isinstance(run._parent, Paragraph): + pPr_b = run._parent._p.xpath("./w:pPr/w:rPr/w:b") for b in pPr_b: val = b.get(f"{_W_NS_CLARK}val") if val not in ["0", "false"]: @@ -2067,7 +2095,7 @@ def _get_children_recursive(node): self._get_hyperlink_target(item), ) ) - elif isinstance(item, Run): + else: content.append( (item.text, self._get_format_from_run(item, paragraph), None) ) @@ -2157,7 +2185,7 @@ def _is_checkbox_checked(self, element: BaseOxmlElement) -> bool: if checked_elem is not None: val = checked_elem.get(f"{{{w14_ns}}}val") - return val == "1" + return str(val) == "1" return False @@ -2335,7 +2363,7 @@ def _collect_textbox_paragraphs( def _handle_textbox_content( self, - textbox_elements: list, + textbox_elements: list[etree._Element], doc: DoclingDocument, ) -> list[RefItem]: """Process textbox content and add it to the document structure.""" @@ -2348,7 +2376,7 @@ def _handle_textbox_content( name="textbox", content_layer=self.content_layer, ) - elem_ref.append(textbox_group.get_ref()) + elem_ref.append(RefItem(cref=textbox_group.self_ref)) # Set this as the current parent to ensure textbox content # is properly nested in document structure original_parent = self.parents[level] @@ -2508,7 +2536,7 @@ def _create_or_reuse_parent( *, doc: DoclingDocument, prev_parent: NodeItem | None, - paragraph_elements: list, + paragraph_elements: list[tuple[str, Formatting | None, AnyUrl | Path | None]], ) -> NodeItem | None: return ( doc.add_inline_group(parent=prev_parent, content_layer=self.content_layer) @@ -2629,7 +2657,7 @@ def _handle_text_elements( content_layer=self.content_layer, ) self.parents[0] = te - elem_ref.append(te.get_ref()) + elem_ref.append(RefItem(cref=te.self_ref)) elif "Heading" in p_style_id: # _add_heading clears the parents tail; reset list context so the # next list item opens fresh under this heading. @@ -2639,9 +2667,7 @@ def _handle_text_elements( elem_ref.extend(h1) elif len(equations) > 0: - if (paragraph.text is None or len(paragraph.text.strip()) == 0) and len( - text - ) > 0: + if len(paragraph.text.strip()) == 0 and len(text) > 0: # Standalone equation(s) — emit each as a separate formula level = self._get_level() parent = self.parents[level - 1] @@ -2655,7 +2681,7 @@ def _handle_text_elements( text=eq_text, content_layer=self.content_layer, ) - elem_ref.append(t1.get_ref()) + elem_ref.append(RefItem(cref=t1.self_ref)) else: t1 = doc.add_text( label=DocItemLabel.FORMULA, @@ -2663,14 +2689,14 @@ def _handle_text_elements( text=text.replace("", "").replace("", ""), content_layer=self.content_layer, ) - elem_ref.append(t1.get_ref()) + elem_ref.append(RefItem(cref=t1.self_ref)) else: # Inline equation level = self._get_level() inline_equation = doc.add_inline_group( parent=self.parents[level - 1], content_layer=self.content_layer ) - elem_ref.append(inline_equation.get_ref()) + elem_ref.append(RefItem(cref=inline_equation.self_ref)) self._add_inline_equations_to_parent( doc=doc, @@ -2713,7 +2739,7 @@ def _handle_text_elements( # Buffered: written only if more code follows, so a # block never ends in blank lines. self._pending_code_blank_lines += 1 - elem_ref.append(merge_target.get_ref()) + elem_ref.append(RefItem(cref=merge_target.self_ref)) self._force_new_code_block = False elif text: # Start a new block, but never on a leading blank paragraph. @@ -2725,7 +2751,7 @@ def _handle_text_elements( content_layer=self.content_layer, code_language=detect_code_language(code_text), ) - elem_ref.append(code_item.get_ref()) + elem_ref.append(RefItem(cref=code_item.self_ref)) self._force_new_code_block = False # A blank that neither starts nor extends a block leaves the # barrier armed. @@ -2756,7 +2782,7 @@ def _handle_text_elements( hyperlink=hyperlink, content_layer=self.content_layer, ) - elem_ref.append(text_item.get_ref()) + elem_ref.append(RefItem(cref=text_item.self_ref)) self._update_history(p_style_id, p_level, numid, ilevel) @@ -2786,7 +2812,7 @@ def _add_heading( label=GroupLabel.SECTION, name=f"header-{i}", ) - elem_ref.append(gr1.get_ref()) + elem_ref.append(RefItem(cref=gr1.self_ref)) self.parents[i] = gr1 elif curr_level < level: @@ -2836,13 +2862,13 @@ def _add_heading( level=add_level, ) self.parents[current_level] = hd - elem_ref.append(hd.get_ref()) + elem_ref.append(RefItem(cref=hd.self_ref)) return elem_ref def _add_formatted_list_item( self, doc: DoclingDocument, - elements: list, + elements: list[tuple[str, Formatting | None, AnyUrl | Path | None]], marker: str, enumerated: bool, level: int, @@ -2891,7 +2917,7 @@ def _add_formatted_list_item( def _add_list_item_with_marker( self, doc: DoclingDocument, - elements: list, + elements: list[tuple[str, Formatting | None, AnyUrl | Path | None]], numid: int, ilevel: int, is_numbered: bool, @@ -2946,7 +2972,7 @@ def _add_inline_equations_to_parent( content_layer=self.content_layer, ) if elem_ref is not None: - elem_ref.append(e1.get_ref()) + elem_ref.append(RefItem(cref=e1.self_ref)) e2 = doc.add_text( label=DocItemLabel.FORMULA, @@ -2955,7 +2981,7 @@ def _add_inline_equations_to_parent( content_layer=self.content_layer, ) if elem_ref is not None: - elem_ref.append(e2.get_ref()) + elem_ref.append(RefItem(cref=e2.self_ref)) if len(text_tmp) > 0: e3 = doc.add_text( @@ -2965,7 +2991,7 @@ def _add_inline_equations_to_parent( content_layer=self.content_layer, ) if elem_ref is not None: - elem_ref.append(e3.get_ref()) + elem_ref.append(RefItem(cref=e3.self_ref)) def _slot_for(self, word_ilevel: int) -> int: """Map a Word ``w:ilvl`` value to the internal parents-slot index. @@ -2980,6 +3006,8 @@ def _slot_for(self, word_ilevel: int) -> int: sits at level 1 and later returns to level 0) are clamped to the list base so they stay inside the current list instead of mapping below it. """ + assert self.level_at_new_list is not None + assert self.level_start_ilevel is not None return self.level_at_new_list + max(0, word_ilevel - self.level_start_ilevel) def _manage_list_structure( @@ -3057,7 +3085,7 @@ def _manage_list_structure( content_layer=self.content_layer, ) self.parents[i] = list_gr1 - elem_ref.append(list_gr1.get_ref()) + elem_ref.append(RefItem(cref=list_gr1.self_ref)) use_level = self._slot_for(ilevel) elif ( @@ -3118,7 +3146,7 @@ def _add_list_item( doc: DoclingDocument, numid: int, ilevel: int, - elements: list, + elements: list[tuple[str, Formatting | None, AnyUrl | Path | None]], is_numbered: bool = False, ) -> list[RefItem]: """Add a regular list item without inline equations. @@ -3234,10 +3262,10 @@ def _group_cell_elements( group_element.children.append(prov) pr_item = prov.resolve(doc) item_parent = pr_item.parent.resolve(doc) - if pr_item.get_ref() in item_parent.children: - item_parent.children.remove(pr_item.get_ref()) - pr_item.parent = group_element.get_ref() - ref_for_rich_cell = group_element.get_ref() + if RefItem(cref=pr_item.self_ref) in item_parent.children: + item_parent.children.remove(RefItem(cref=pr_item.self_ref)) + pr_item.parent = RefItem(cref=group_element.self_ref) + ref_for_rich_cell = RefItem(cref=group_element.self_ref) return ref_for_rich_cell @staticmethod @@ -3317,7 +3345,7 @@ def _handle_tables( docling_table = doc.add_table( data=data, parent=self.parents[level - 1], content_layer=self.content_layer ) - elem_ref.append(docling_table.get_ref()) + elem_ref.append(RefItem(cref=docling_table.self_ref)) open_cells: dict[int, TableCell] = {} for row_idx, row in enumerate(table.rows): @@ -3460,10 +3488,9 @@ def _is_rich_table_cell(self, cell: _Cell) -> bool: runs = list(para.iterchildren(f"{_W_NS_CLARK}r")) for rn in runs: item: Run = Run(rn, self.docx_obj) - if item is not None: - fm = MsWordDocumentBackend._get_format_from_run(item) - if fm != Formatting(): - return True + fm = MsWordDocumentBackend._get_format_from_run(item) + if fm != Formatting(): + return True # Walk a non-empty code-styled cell as rich content so it can emit a # CodeItem; the font fallback never fires inside cells. @@ -3502,7 +3529,9 @@ def _get_image_from_relationship( return None # Access the image part using the relationship ID image_part = rel.target_part - image_data = image_part.blob # Get the binary image data + blob = image_part.blob # Get the binary image data + assert isinstance(blob, bytes) + image_data = blob return image_data def _add_picture_to_doc( @@ -3537,7 +3566,7 @@ def _add_picture_to_doc( caption=None, content_layer=target_layer, ) - return p.get_ref() + return RefItem(cref=p.self_ref) def _convert_elements_via_docx( self, elements: Any | list[Any], element_tag: str | list[str] | None = None @@ -3570,6 +3599,7 @@ def _convert_elements_via_docx( try: # Loaded once per backend instance; del body[:] clears cheaply on every call. if self._empty_docx_template is None: + assert self.path_or_stream is not None self._empty_docx_template = self.load_msword_file( self.path_or_stream, self.document_hash ) @@ -3727,7 +3757,7 @@ def _handle_vml_pictures( ) pil_image = None - if pil_image is None and image_data is not None: + if pil_image is None: pil_image = self._convert_elements_via_docx( imagedata, ["object", "pict"] ) @@ -3914,7 +3944,7 @@ def _chart_title_text(self, chart_root: Any) -> str | None: if not text: cached = self._read_chart_cache(title) text = cached[0].strip() if cached else "" - return text or None + return str(text) if text else None def _chart_to_table_data(self, chart_root: Any) -> TableData | None: """Reconstruct a chart's underlying data grid as a TableData. @@ -4088,7 +4118,7 @@ def _handle_chart(self, doc: DoclingDocument, chart_el: Any) -> RefItem | None: else None ), ) - return picture.get_ref() + return RefItem(cref=picture.self_ref) def _add_header_footer(self, docx_obj: DocxDocument, doc: DoclingDocument) -> None: """Add section headers and footers. diff --git a/docling/backend/noop_backend.py b/docling/backend/noop_backend.py index dd58fc1140..344432e5e6 100644 --- a/docling/backend/noop_backend.py +++ b/docling/backend/noop_backend.py @@ -6,6 +6,8 @@ from pathlib import Path from typing import Set, Union +from typing_extensions import override + from docling.backend.abstract_backend import AbstractDocumentBackend from docling.datamodel.base_models import InputFormat from docling.datamodel.document import InputDocument @@ -42,13 +44,16 @@ def __init__(self, in_doc: "InputDocument", path_or_stream: Union[BytesIO, Path] _log.error(f"NoOpBackend validation failed: {e}") self.valid = False + @override def is_valid(self) -> bool: return self.valid @classmethod + @override def supports_pagination(cls) -> bool: return False @classmethod + @override def supported_formats(cls) -> Set[InputFormat]: return set(InputFormat) diff --git a/docling/backend/opendocument_backend.py b/docling/backend/opendocument_backend.py index 877118a52f..1b45862ce0 100644 --- a/docling/backend/opendocument_backend.py +++ b/docling/backend/opendocument_backend.py @@ -25,7 +25,7 @@ from dataclasses import dataclass from io import BytesIO from pathlib import Path -from typing import Any, cast +from typing import Any, TypedDict, cast from docling_core.types.doc import ( BoundingBox, @@ -61,7 +61,11 @@ PaginatedDocumentBackend, ) from docling.backend.utils.image_resource_loader import ImageResourceLoader -from docling.datamodel.backend_options import OdsBackendOptions +from docling.datamodel.backend_options import ( + BackendOptions, + BaseBackendOptions, + OdsBackendOptions, +) from docling.datamodel.base_models import InputFormat from docling.datamodel.document import InputDocument @@ -77,7 +81,9 @@ ListItem, Paragraph, Section, + Style, Table as OdfTable, + XmlPart, ) _ODFDO_AVAILABLE = True @@ -149,6 +155,23 @@ def _load_odf_document( ) from e +class _OdsTableRegion(TypedDict): + bounds: tuple[int, int, int, int] + data: TableData + + +class _OdfCellData(TypedDict): + text: str + row_span: int + col_span: int + start_row_offset_idx: int + end_row_offset_idx: int + start_col_offset_idx: int + end_col_offset_idx: int + column_header: bool + row_header: bool + + class _OdfBaseBackend(DeclarativeDocumentBackend): """Shared loading / validation logic for ODT, ODS and ODP backends.""" @@ -159,12 +182,12 @@ def __init__( self, in_doc: InputDocument, path_or_stream: BytesIO | Path, - options: OdsBackendOptions | None = None, + options: BackendOptions | None = None, ) -> None: if not _ODFDO_AVAILABLE: raise ImportError(_INSTALL_HINT) from _ODFDO_IMPORT_ERROR super().__init__(in_doc, path_or_stream, options) - self.path_or_stream: BytesIO | Path = path_or_stream + self.path_or_stream: BytesIO | Path | None = path_or_stream self.page_range = in_doc.limits.page_range self.valid: bool = False self.odf_obj: OdfDocument = _load_odf_document( @@ -315,7 +338,7 @@ def _formatting_from_odf_text_style( return _formatting_or_none(formatting) style = odf_obj.get_style("text", style_name) - if style is None: + if not isinstance(style, Style): return _formatting_or_none(formatting) props = style.get_properties() or {} @@ -331,7 +354,7 @@ def _formatting_from_odf_text_style( ), None, ) - if font_weight is not None: + if isinstance(font_weight, str): formatting.bold = _is_bold_weight(font_weight) font_style = next( @@ -360,7 +383,7 @@ def _formatting_from_odf_text_style( formatting.strikethrough = line_through != "none" text_position = props.get("style:text-position") - if text_position is not None: + if isinstance(text_position, str): if text_position.startswith("super"): formatting.script = Script.SUPER elif text_position.startswith("sub"): @@ -480,7 +503,7 @@ def _add_odf_text_runs( if not runs: return None if len(runs) == 1: - return doc.add_text( + item = doc.add_text( label=label, parent=parent, text=runs[0].text, @@ -489,6 +512,9 @@ def _add_odf_text_runs( hyperlink=runs[0].hyperlink, ) + assert isinstance(item, NodeItem) + return item + inline_group = doc.add_inline_group(parent=parent, content_layer=content_layer) for run in runs: doc.add_text( @@ -770,7 +796,7 @@ def _odf_list_level_style( return None style = odf_obj.get_style("list", style_name) - if style is None: + if not isinstance(style, Style): return None return style.get_level_style(level) @@ -785,7 +811,7 @@ def _odf_list_level_is_enumerated( level_style = _odf_list_level_style(odf_obj, odf_list, level) if level_style is None: return fallback - return level_style.tag == "text:list-level-style-number" + return bool(level_style.tag == "text:list-level-style-number") def _odf_list_start_value( @@ -901,11 +927,15 @@ def _image_ref_from_odf_image( image_data: bytes | None = None get_data = getattr(image, "get_data", None) if callable(get_data): - image_data = get_data() + data = get_data() + if isinstance(data, bytes): + image_data = data if image_data is None and odf_obj is not None and image_url: try: - image_data = odf_obj.get_part(image_url) + part = odf_obj.get_part(image_url) + if isinstance(part, bytes): + image_data = part except Exception: image_data = None @@ -923,15 +953,16 @@ def _image_ref_from_odf_image( def _odf_image_href(image: Any) -> str | None: - return getattr(image, "url", None) or getattr(image, "attributes", {}).get( + href = getattr(image, "url", None) or getattr(image, "attributes", {}).get( "xlink:href" ) + return str(href) if href is not None else None def _odf_image_can_be_bitmap(image: Any, image_url: str | None) -> bool: mime_type = getattr(image, "attributes", {}).get("draw:mime-type") if mime_type is not None: - return mime_type.startswith("image/") and mime_type != "image/svg+xml" + return bool(mime_type.startswith("image/") and mime_type != "image/svg+xml") if image_url is None: return True @@ -974,7 +1005,7 @@ def _strip_odf_image_reference_text(text: str, images: list[Any]) -> str: def _add_odf_images( doc: DoclingDocument, images: list[Any], - parent: NodeItem, + parent: NodeItem | None, content_layer: ContentLayer | None, odf_obj: OdfDocument | None, *, @@ -1143,7 +1174,7 @@ def _odf_chart_title(chart_content: Any) -> str | None: for title in chart_content.get_elements("descendant::chart:title"): text = title.text_content.strip() if text: - return text + return str(text) return None @@ -1163,6 +1194,8 @@ def _chart_data_from_frame( try: chart_content = odf_obj.get_part(_embedded_odf_content_path(object_href)) + if not isinstance(chart_content, XmlPart): + return None # odfdo resolves XML parts lazily, so a reference to a part that is missing # from the package only fails once the part is actually read. chart_classification = _odf_chart_classification(chart_content) @@ -1238,7 +1271,7 @@ def _add_odf_charts( def _add_odf_list( doc: DoclingDocument, odf_list: OdfList, - parent: NodeItem, + parent: NodeItem | None, content_layer: ContentLayer | None, odf_obj: OdfDocument | None, enumerated: bool = False, @@ -1283,7 +1316,7 @@ def _add_odf_list( ) return None - if should_continue and continued_state is not None: + if should_continue: list_group = continued_state.group current_enumerated = continued_state.enumerated counter = continued_state.counter @@ -1392,7 +1425,7 @@ def _add_odf_list( def _add_rich_cell_children( doc: DoclingDocument, cell: Any, - parent: NodeItem, + parent: NodeItem | None, content_layer: ContentLayer | None, odf_obj: OdfDocument | None, image_loader: ImageResourceLoader | None = None, @@ -1460,7 +1493,7 @@ def _add_table_from_odf( adjusted_row = row_idx - min_row adjusted_col = col_idx - min_col text = _odf_cell_text(cell) - cell_kwargs = { + cell_kwargs: _OdfCellData = { "text": text, "row_span": row_span, "col_span": col_span, @@ -1493,6 +1526,7 @@ def _add_table_from_odf( doc.add_table_cell(table_item=table_item, cell=table_cell) + assert isinstance(table_item, TableItem) return table_item @@ -1741,8 +1775,8 @@ def _is_slide_title_element(element: Any, is_first_text_content: bool) -> bool: attrs = getattr(element, "attributes", {}) if attrs.get("presentation:class") == "title": return True - return is_first_text_content and getattr(element, "tag", None) == ( - "draw:custom-shape" + return is_first_text_content and bool( + getattr(element, "tag", None) == "draw:custom-shape" ) def _walk_slide_frame( @@ -1762,6 +1796,8 @@ def _walk_slide_frame( ) for tbl in frame.get_elements("descendant::table:table"): + if not isinstance(tbl, OdfTable): + continue _add_table_from_odf( doc, tbl, @@ -1843,10 +1879,12 @@ def __init__( self, in_doc: InputDocument, path_or_stream: BytesIO | Path, - options: OdsBackendOptions | None = None, + options: BaseBackendOptions | None = None, ) -> None: if options is None: options = OdsBackendOptions() + if not isinstance(options, OdsBackendOptions): + raise TypeError("Expected OdsBackendOptions.") super().__init__(in_doc, path_or_stream, options) @classmethod @@ -2004,9 +2042,7 @@ def _convert_sheet_table( odf_obj=self.odf_obj, ) - def _find_data_tables_in_sheet( - self, table: OdfTable - ) -> list[dict[str, tuple[int, int, int, int] | TableData]]: + def _find_data_tables_in_sheet(self, table: OdfTable) -> list[_OdsTableRegion]: """Find all disconnected data tables in an ODS sheet using flood-fill. Returns a list of dictionaries, each containing: @@ -2032,7 +2068,7 @@ def _find_data_tables_in_sheet( return [] GAP_TOLERANCE = cast(OdsBackendOptions, self.options).gap_tolerance - tables: list[dict[str, tuple[int, int, int, int] | TableData]] = [] + tables: list[_OdsTableRegion] = [] visited: set[tuple[int, int]] = set() # Build a map of cell contents for quick lookup @@ -2126,15 +2162,15 @@ def _find_images_in_sheet( # Get all images in the table images = table.get_images() for img in images: + # Rich cells own their images; this pass handles sheet drawings. + if img.get_elements("ancestor::table:table-cell"): + continue try: # Get the image data - image_data = img.get_data() - if image_data: - # Convert to PIL Image - from io import BytesIO - - pil_image = PILImage.open(BytesIO(image_data)) - + image_ref = _image_ref_from_odf_image( + self.odf_obj, img, image_loader=self._image_loader + ) + if image_ref is not None: # Try to get position information # ODF images are typically anchored to cells # For now, use a default position @@ -2142,7 +2178,7 @@ def _find_images_in_sheet( doc.add_picture( parent=parent, - image=ImageRef.from_pil(image=pil_image, dpi=72), + image=image_ref, caption=None, prov=ProvenanceItem( page_no=page_no, @@ -2170,8 +2206,6 @@ def _find_page_size(doc: DoclingDocument, page_no: int) -> tuple[float, float]: if not isinstance(item, DocItem): continue for provenance in item.prov: - if provenance.bbox is None: - continue bbox = provenance.bbox left = min(left, bbox.l) if left != -1 else bbox.l right = max(right, bbox.r) if right != -1 else bbox.r diff --git a/docling/backend/pdf_backend.py b/docling/backend/pdf_backend.py index 4a93f3e84e..d196b51aca 100644 --- a/docling/backend/pdf_backend.py +++ b/docling/backend/pdf_backend.py @@ -10,6 +10,7 @@ from docling_core.types.doc import BoundingBox, Size from docling_core.types.doc.page import SegmentedPdfPage, TextCell from PIL import Image +from typing_extensions import override from docling.backend.abstract_backend import PaginatedDocumentBackend from docling.datamodel.backend_options import PdfBackendOptions @@ -138,6 +139,7 @@ def load_page(self, page_no: int) -> PdfPageBackend: pass @abstractmethod + @override def page_count(self) -> int: pass @@ -155,10 +157,12 @@ def get_document_outline(self) -> list[_PdfOutlineItem]: return [] @classmethod + @override def supported_formats(cls) -> Set[InputFormat]: return {InputFormat.PDF} @classmethod + @override def supports_pagination(cls) -> bool: return True diff --git a/docling/backend/pypdfium2_backend.py b/docling/backend/pypdfium2_backend.py index 2f252b87de..ff3431e524 100644 --- a/docling/backend/pypdfium2_backend.py +++ b/docling/backend/pypdfium2_backend.py @@ -23,6 +23,7 @@ from pypdfium2 import PdfTextPage from pypdfium2._helpers.misc import PdfiumError from rtree import index +from typing_extensions import override from docling.backend.managed_pdfium_backend import ( ManagedPdfiumDocumentBackend, @@ -277,10 +278,12 @@ def __init__( # the text or path objects. self._object_index: dict[int, List[tuple[BoundingBox, bool]]] = {} + @override def is_valid(self) -> bool: return self.valid @property + @override def page_no(self) -> int: return self._page_no + 1 @@ -451,7 +454,7 @@ def _get_object_bucket(self, obj_type: int) -> List[tuple[BoundingBox, bool]]: Returns: The cached ``(bbox, invisible_text)`` list for ``obj_type``. """ - bucket = self._object_index.get(obj_type) + bucket: list[tuple[BoundingBox, bool]] | None = self._object_index.get(obj_type) if bucket is not None: return bucket @@ -481,6 +484,7 @@ def _get_object_bucket(self, obj_type: int) -> List[tuple[BoundingBox, bool]]: bbox = BoundingBox.from_tuple( pos, origin=CoordOrigin.BOTTOMLEFT ).to_top_left_origin(page_height=page_size.height) + assert isinstance(bbox, BoundingBox) bucket.append((bbox, invisible)) self._object_index[obj_type] = bucket @@ -502,6 +506,7 @@ def _object_rects( continue yield bbox + @override def get_bitmap_rects(self, scale: float = 1) -> Iterable[BoundingBox]: AREA_THRESHOLD = 0 # 32 * 32 @@ -509,6 +514,7 @@ def get_bitmap_rects(self, scale: float = 1) -> Iterable[BoundingBox]: if cropbox.area() > AREA_THRESHOLD: yield cropbox.scaled(scale=scale) + @override def has_content_in( self, *, @@ -551,6 +557,7 @@ def has_content_in( return False + @override def get_connected_shape_bounding_boxes( self, *, tolerance: float = 0.0 ) -> Optional[List[BoundingBox]]: @@ -566,6 +573,7 @@ def get_connected_shape_bounding_boxes( list(self._object_rects(pdfium_c.FPDF_PAGEOBJ_PATH)), tolerance ) + @override def get_text_in_rect(self, bbox: BoundingBox) -> str: with pypdfium2_lock: page = self._require_page() @@ -585,8 +593,10 @@ def get_text_in_rect(self, bbox: BoundingBox) -> str: with pypdfium2_lock: text_piece = self.text_page.get_text_bounded(*pdf_rect) + assert isinstance(text_piece, str) return text_piece + @override def get_segmented_page(self) -> Optional[SegmentedPdfPage]: if not self.valid: return None @@ -605,15 +615,16 @@ def get_segmented_page(self) -> Optional[SegmentedPdfPage]: textline_cells=text_cells, char_cells=[], word_cells=[], - has_textlines=len(text_cells) > 0, has_words=False, has_chars=False, ) return self._seg_page + @override def get_text_cells(self) -> Iterable[TextCell]: return self._compute_text_cells() + @override def get_page_image( self, scale: float = 1, cropbox: Optional[BoundingBox] = None ) -> Image.Image: @@ -648,13 +659,16 @@ def get_page_image( size=(round(cropbox.width * scale), round(cropbox.height * scale)) ) + assert isinstance(image, Image.Image) return image + @override def get_size(self) -> Size: with pypdfium2_lock: page = self._require_page() return Size(width=page.get_width(), height=page.get_height()) + @override def _close_native_page(self) -> None: with pypdfium2_lock: if self.text_page is not None: @@ -690,23 +704,32 @@ def __init__( f"pypdfium could not load document with hash {self.document_hash}" ) from e + @override def page_count(self) -> int: with pypdfium2_lock: + if self._pdoc is None: + raise RuntimeError("PDF document is unloaded.") return len(self._pdoc) + @override def load_page(self, page_no: int) -> PyPdfiumPageBackend: with pypdfium2_lock: + if self._pdoc is None: + raise RuntimeError("PDF document is unloaded.") return PyPdfiumPageBackend(self._pdoc, self.document_hash, page_no) + @override def is_valid(self) -> bool: return self.page_count() > 0 + @override def get_document_outline(self) -> list[_PdfOutlineItem]: """Extract the PDF outline from the pypdfium2 document (title, depth, page, position).""" if self._pdoc is None: return [] return extract_outline_from_pdfium(self._pdoc) + @override def _close_native_document(self) -> None: with pypdfium2_lock: if self._pdoc is not None: diff --git a/docling/backend/utils/image_resource_loader.py b/docling/backend/utils/image_resource_loader.py index 5d498761a9..6d701500cd 100644 --- a/docling/backend/utils/image_resource_loader.py +++ b/docling/backend/utils/image_resource_loader.py @@ -27,6 +27,7 @@ import certifi import requests import urllib3 +import urllib3.exceptions from docling_core.types.doc.document import ImageRef from PIL import Image, UnidentifiedImageError from pydantic import ValidationError @@ -206,6 +207,7 @@ def _open_direct(url: str, headers: dict[str, str]) -> Iterator[urllib3.HTTPResp ) as e: errors.append(e) continue + assert isinstance(response, urllib3.HTTPResponse) with response: yield response return @@ -224,7 +226,9 @@ def _open_with_proxy( timeout=(5, 30), allow_redirects=False, ) as response: - yield response.raw + raw = response.raw + assert isinstance(raw, urllib3.HTTPResponse) + yield raw def url_origin(url: str) -> Optional[tuple[str, str, Optional[int]]]: diff --git a/docling/backend/webvtt_backend.py b/docling/backend/webvtt_backend.py index 2a0aa462f3..05a0d5bf62 100644 --- a/docling/backend/webvtt_backend.py +++ b/docling/backend/webvtt_backend.py @@ -27,6 +27,7 @@ from typing_extensions import override from docling.backend.abstract_backend import DeclarativeDocumentBackend +from docling.datamodel.backend_options import BackendOptions from docling.datamodel.base_models import InputFormat from docling.datamodel.document import InputDocument from docling.exceptions import DocumentLoadError @@ -40,7 +41,7 @@ class AnnotatedText: voice: str | None = None formatting: Formatting | None = None - def copy_meta(self, text): + def copy_meta(self, text: str) -> "AnnotatedText": return AnnotatedText( text=text, voice=self.voice, @@ -64,8 +65,13 @@ class WebVTTDocumentBackend(DeclarativeDocumentBackend): """ @override - def __init__(self, in_doc: InputDocument, path_or_stream: BytesIO | Path): - super().__init__(in_doc, path_or_stream) + def __init__( + self, + in_doc: InputDocument, + path_or_stream: BytesIO | Path, + options: BackendOptions | None = None, + ): + super().__init__(in_doc, path_or_stream, options) self.content: str = "" # utf-8-sig drops a leading BOM. The WebVTT file body grammar allows an diff --git a/docling/backend/xml/doclang_archive_backend.py b/docling/backend/xml/doclang_archive_backend.py index 742c5255ec..8ef3075531 100644 --- a/docling/backend/xml/doclang_archive_backend.py +++ b/docling/backend/xml/doclang_archive_backend.py @@ -11,6 +11,7 @@ from typing_extensions import override from docling.backend.abstract_backend import DeclarativeDocumentBackend +from docling.datamodel.backend_options import BackendOptions from docling.datamodel.base_models import InputFormat from docling.datamodel.document import InputDocument @@ -20,9 +21,12 @@ class DocLangArchiveBackend(DeclarativeDocumentBackend): @override def __init__( - self, in_doc: InputDocument, path_or_stream: Union[BytesIO, Path] + self, + in_doc: InputDocument, + path_or_stream: Union[BytesIO, Path], + options: BackendOptions | None = None, ) -> None: - super().__init__(in_doc, path_or_stream) + super().__init__(in_doc, path_or_stream, options) self._temp_dir: Path | None = None self._doc_or_err = self._get_doc_or_err() diff --git a/docling/backend/xml/doclang_backend.py b/docling/backend/xml/doclang_backend.py index 356c6a2254..b911ba3a52 100644 --- a/docling/backend/xml/doclang_backend.py +++ b/docling/backend/xml/doclang_backend.py @@ -10,6 +10,7 @@ from typing_extensions import override from docling.backend.abstract_backend import DeclarativeDocumentBackend +from docling.datamodel.backend_options import BackendOptions from docling.datamodel.base_models import InputFormat from docling.datamodel.document import InputDocument @@ -17,9 +18,12 @@ class DocLangDocumentBackend(DeclarativeDocumentBackend): @override def __init__( - self, in_doc: InputDocument, path_or_stream: Union[BytesIO, Path] + self, + in_doc: InputDocument, + path_or_stream: Union[BytesIO, Path], + options: BackendOptions | None = None, ) -> None: - super().__init__(in_doc, path_or_stream) + super().__init__(in_doc, path_or_stream, options) self._doc_or_err = self._get_doc_or_err() @override diff --git a/docling/backend/xml/jats_backend.py b/docling/backend/xml/jats_backend.py index aebf0b82e2..388ab4472a 100644 --- a/docling/backend/xml/jats_backend.py +++ b/docling/backend/xml/jats_backend.py @@ -51,7 +51,7 @@ from docling.backend.html_backend import HTMLDocumentBackend from docling.backend.utils.image_resource_loader import ImageResourceLoader from docling.backend.utils.table_spans import table_width -from docling.datamodel.backend_options import JatsBackendOptions +from docling.datamodel.backend_options import BaseBackendOptions, JatsBackendOptions from docling.datamodel.base_models import InputFormat from docling.datamodel.document import InputDocument from docling.exceptions import DocumentLoadError @@ -187,12 +187,14 @@ def __init__( self, in_doc: InputDocument, path_or_stream: BytesIO | Path, - options: JatsBackendOptions | None = None, + options: BaseBackendOptions | None = None, ) -> None: if not _BS4_AVAILABLE: raise ImportError(_INSTALL_HINT) from _BS4_IMPORT_ERROR if options is None: options = JatsBackendOptions() + if not isinstance(options, JatsBackendOptions): + raise TypeError("Expected JatsBackendOptions.") super().__init__(in_doc, path_or_stream, options) self.options: JatsBackendOptions self.path_or_stream = path_or_stream @@ -266,14 +268,14 @@ def supported_formats(cls) -> set[InputFormat]: @override def convert(self) -> DoclingDocument: + # Create the document before the recoverable parse block. + origin = DocumentOrigin( + filename=self.file.name or "file", + mimetype="application/xml", + binary_hash=self.document_hash, + ) + doc = DoclingDocument(name=self.file.stem or "file", origin=origin) try: - # Create empty document - origin = DocumentOrigin( - filename=self.file.name or "file", - mimetype="application/xml", - binary_hash=self.document_hash, - ) - doc = DoclingDocument(name=self.file.stem or "file", origin=origin) self.hlevel = 0 # Get metadata XML components @@ -300,7 +302,7 @@ def convert(self) -> DoclingDocument: def _get_text(node: etree._Element, sep: str | None = None) -> str: skip_tags = ["term", "disp-formula", "inline-formula"] text: str = ( - node.text.replace("\n", " ") + str(node.text).replace("\n", " ") if (node.tag not in skip_tags and node.text) else "" ) @@ -310,7 +312,7 @@ def _get_text(node: etree._Element, sep: str | None = None) -> str: text += JatsDocumentBackend._get_text(child, sep) if sep: text = text.rstrip(sep) + sep - text += child.tail.replace("\n", " ") if child.tail else "" + text += str(child.tail).replace("\n", " ") if child.tail else "" return text @@ -552,11 +554,12 @@ def _add_abstract( for section in sections: section_title = section["title"] if section_title: - section_parent: NodeItem = doc.add_heading( + section_parent = doc.add_heading( parent=abstract_heading, text=section_title, level=self.hlevel + 2, ) + assert isinstance(section_parent, NodeItem) else: section_parent = abstract_heading for paragraph in section["paragraphs"]: @@ -578,9 +581,9 @@ def _add_abstract( def _add_authors(self, doc: DoclingDocument, xml_components: XMLComponents) -> None: # TODO: once docling supports text formatting, add affiliation reference to # author names through superscripts - authors: list = [item["name"] for item in xml_components["authors"]] + authors: list[str] = [item["name"] for item in xml_components["authors"]] authors_str = ", ".join(authors) - affiliations: list = [ + affiliations: list[str] = [ item for author in xml_components["authors"] for item in author["affiliation_names"] @@ -659,19 +662,19 @@ def _parse_element_citation(self, node: etree._Element) -> str: ) # Journal, year, publisher name, publisher location, volume, elocation - fields: list[str] = [ - "source", - "year", - "publisher-name", - "publisher-loc", - "volume", - ] - for item in fields: - item_node = node.xpath(item) - if len(item_node) > 0: - citation[item.replace("-", "_")] = ( # type: ignore[literal-required] - JatsDocumentBackend._normalize_whitespace(item_node[0].text) - ) + def field_text(tag: str) -> str: + nodes = node.xpath(tag) + return ( + JatsDocumentBackend._normalize_whitespace(nodes[0].text) + if nodes + else "" + ) + + citation["source"] = field_text("source") + citation["year"] = field_text("year") + citation["publisher_name"] = field_text("publisher-name") + citation["publisher_loc"] = field_text("publisher-loc") + citation["volume"] = field_text("volume") # Publication identifier if len(node.xpath("pub-id")) > 0: @@ -751,7 +754,7 @@ def _add_equation( def _extract_tex_math(node: etree._Element) -> str | None: if not node.text: return None - text = node.text.strip() + text = str(node.text).strip() for delimiter in ("$$", "$"): if ( len(text) > 2 * len(delimiter) @@ -796,7 +799,7 @@ def _walk_inline_formula( current = JatsDocumentBackend._merge_formatting(formatting, node.tag) segments: list[InlineSegment] = [] if node.text: - text = node.text.replace("\n", " ") + text = str(node.text).replace("\n", " ") if text: segments.append( InlineSegment( @@ -826,7 +829,7 @@ def _walk_inline_formula( JatsDocumentBackend._walk_inline_formula(child, current, hyperlink) ) if child.tail: - tail = child.tail.replace("\n", " ") + tail = str(child.tail).replace("\n", " ") if tail: segments.append( InlineSegment( @@ -922,12 +925,14 @@ def _add_figure_captions( # TODO: format label vs caption once styling is supported fig_text: str = f"{label}{' ' if label and caption else ''}{caption}" - fig_caption: TextItem | None = ( + fig_caption = ( doc.add_text(label=DocItemLabel.CAPTION, text=fig_text) if fig_text else None ) + assert fig_caption is None or isinstance(fig_caption, TextItem) + doc.add_picture( parent=parent, caption=fig_caption, @@ -1034,12 +1039,8 @@ def parse_table_data(element: Tag) -> TableData | None: for row in element("tr"): cell_spans: list[tuple[int, int]] = [] is_row_header = True - if not isinstance(row, Tag): - continue for cell in row(["td", "th"]): - if not isinstance(row, Tag): - continue - cell_tag = cast(Tag, cell) + cell_tag = cell col_span, row_span = HTMLDocumentBackend._get_cell_spans(cell_tag) cell_spans.append((col_span, row_span)) if cell_tag.name == "td" or row_span == 1: @@ -1051,7 +1052,9 @@ def parse_table_data(element: Tag) -> TableData | None: _log.debug(f"The table has {num_rows} rows and {num_cols} cols.") - grid: list = [[None for _ in range(num_cols)] for _ in range(num_rows)] + grid: list[list[str | None]] = [ + [None for _ in range(num_cols)] for _ in range(num_rows) + ] data = TableData(num_rows=num_rows, num_cols=num_cols, table_cells=[]) @@ -1059,9 +1062,6 @@ def parse_table_data(element: Tag) -> TableData | None: start_row_span = 0 row_idx = -1 for row in element("tr"): - if not isinstance(row, Tag): - continue - # For each row, find all the column cells (both and ) cells = row(["td", "th"]) @@ -1069,13 +1069,12 @@ def parse_table_data(element: Tag) -> TableData | None: col_header = True row_header = True for html_cell in cells: - if isinstance(html_cell, Tag): - _, row_span = HTMLDocumentBackend._get_cell_spans(html_cell) - if html_cell.name == "td": - col_header = False - row_header = False - elif row_span == 1: - row_header = False + _, row_span = HTMLDocumentBackend._get_cell_spans(html_cell) + if html_cell.name == "td": + col_header = False + row_header = False + elif row_span == 1: + row_header = False if not row_header: row_idx += 1 start_row_span = 0 @@ -1085,9 +1084,6 @@ def parse_table_data(element: Tag) -> TableData | None: # Extract the text content of each cell col_idx = 0 for html_cell in cells: - if not isinstance(html_cell, Tag): - continue - # extract inline formulas for formula in html_cell("inline-formula"): math_parts = formula.text.split("$$") @@ -1142,11 +1138,13 @@ def _add_table( label = table_xml_component["label"] caption = table_xml_component["caption"] table_text: str = f"{label}{' ' if label and caption else ''}{caption}" - table_caption: TextItem | None = ( + table_caption = ( doc.add_text(label=DocItemLabel.CAPTION, text=table_text) if table_text else None ) + assert table_caption is None or isinstance(table_caption, TextItem) + if data is not None: doc.add_table(data=data, parent=parent, caption=table_caption) @@ -1245,7 +1243,7 @@ def _walk_linear( ) -> list[InlineSegment]: skip_tags = ["term"] flush_tags = ["ack", "sec", "list", "boxed-text", "disp-formula", "fig"] - new_parent: NodeItem = parent + new_parent = parent current = JatsDocumentBackend._merge_formatting(formatting, node.tag) current_hyperlink = hyperlink if node.tag == "ext-link": @@ -1267,6 +1265,7 @@ def _walk_linear( inline_segments = [] # add elements and decide whether to stop walking + section_has_heading = False if child.tag in ("sec", "ack"): header = child.xpath("title|label") text: str | None = None @@ -1275,14 +1274,17 @@ def _walk_linear( elif child.tag == "ack": text = DEFAULT_HEADER_ACKNOWLEDGMENTS if text: + section_has_heading = True self.hlevel += 1 new_parent = doc.add_heading( text=text, parent=parent, level=self.hlevel ) + assert isinstance(new_parent, NodeItem) elif child.tag == "list": new_parent = doc.add_group( label=GroupLabel.LIST, name="list", parent=parent ) + assert isinstance(new_parent, NodeItem) elif child.tag == "list-item": # TODO: address non-paragraph, non-list content inside list-item # (e.g. disp-formula, fig, table-wrap) @@ -1302,6 +1304,7 @@ def _walk_linear( text=text, parent=parent, ) + assert isinstance(new_parent, NodeItem) for nested in nested_lists: nested_group = doc.add_group( @@ -1334,9 +1337,11 @@ def _walk_linear( else DEFAULT_HEADER_REFERENCES ) new_parent = doc.add_heading(text=text, parent=parent) + assert isinstance(new_parent, NodeItem) new_parent = doc.add_group( parent=new_parent, label=GroupLabel.LIST, name="list" ) + assert isinstance(new_parent, NodeItem) elif child.tag == "element-citation": text = self._parse_element_citation(child) self._add_citation(doc, parent, text) @@ -1367,7 +1372,7 @@ def _walk_linear( JatsDocumentBackend._extend_segments( inline_segments, child_segments ) - if child.tag in ("sec", "ack") and text: + if section_has_heading: self.hlevel -= 1 # pick up the tail text diff --git a/docling/backend/xml/uspto_backend.py b/docling/backend/xml/uspto_backend.py index 3e0cf07bfe..60c8b3aabb 100644 --- a/docling/backend/xml/uspto_backend.py +++ b/docling/backend/xml/uspto_backend.py @@ -67,6 +67,7 @@ from typing_extensions import Self, TypedDict, override from docling.backend.abstract_backend import DeclarativeDocumentBackend +from docling.datamodel.backend_options import BackendOptions from docling.datamodel.base_models import InputFormat from docling.datamodel.document import InputDocument from docling.exceptions import DocumentLoadError @@ -132,22 +133,27 @@ class PatentHeading(Enum): CLAIMS = "CLAIMS", 2 @override - def __new__(cls, value: str, _) -> Self: + def __new__(cls, value: object, _: object = None) -> Self: obj = object.__new__(cls) obj._value_ = value return obj @override - def __init__(self, _, level: LevelNumber) -> None: + def __init__(self, _: object = None, level: LevelNumber = 2) -> None: self.level: LevelNumber = level class PatentUsptoDocumentBackend(DeclarativeDocumentBackend): @override - def __init__(self, in_doc: InputDocument, path_or_stream: BytesIO | Path) -> None: + def __init__( + self, + in_doc: InputDocument, + path_or_stream: BytesIO | Path, + options: BackendOptions | None = None, + ) -> None: if not _BS4_AVAILABLE: raise ImportError(_INSTALL_HINT) from _BS4_IMPORT_ERROR - super().__init__(in_doc, path_or_stream) + super().__init__(in_doc, path_or_stream, options) self.patent_content: str = "" self.parser: PatentUspto | None = None @@ -264,6 +270,7 @@ def __init__(self) -> None: """Build an instance of PatentUsptoIce class.""" self.handler = PatentUsptoIce.PatentHandler() + @override def parse(self, patent_content: str) -> DoclingDocument | None: try: parser = make_parser() @@ -335,13 +342,13 @@ class Element(Enum): MATHS = "maths", False # to avoid keeping formulas @override - def __new__(cls, value: str, _) -> Self: + def __new__(cls, value: object, _: object = None) -> Self: obj = object.__new__(cls) obj._value_ = value return obj @override - def __init__(self, _, is_text: bool) -> None: + def __init__(self, _: object = None, is_text: bool = False) -> None: self.is_text: bool = is_text @override @@ -364,13 +371,15 @@ def __init__(self) -> None: self.style_html = HtmlEntity() @override - def startElement(self, tag, attributes): + def startElement(self, name: str, attrs: AttributesImpl) -> None: """Signal the start of an element. Args: tag: The element tag. attributes: The element attributes. """ + tag = name + attributes = attrs if tag in ( self.APP_DOC_ELEMENT, self.GRANT_DOC_ELEMENT, @@ -414,12 +423,13 @@ def skippedEntity(self, name): self.text += unescaped @override - def endElement(self, tag): + def endElement(self, name: str) -> None: """Signal the end of an element. Args: tag: The element tag. """ + tag = name if tag in ( self.APP_DOC_ELEMENT, self.GRANT_DOC_ELEMENT, @@ -686,13 +696,13 @@ class Element(Enum): TABLE = ("table", False) # to keep track of table positions @override - def __new__(cls, value: str, _) -> Self: + def __new__(cls, value: object, _: object = None) -> Self: obj = object.__new__(cls) obj._value_ = value return obj @override - def __init__(self, _, is_text: bool) -> None: + def __init__(self, _: object = None, is_text: bool = False) -> None: self.is_text: bool = is_text @override @@ -715,13 +725,15 @@ def __init__(self) -> None: self.style_html = HtmlEntity() @override - def startElement(self, tag, attributes): + def startElement(self, name: str, attrs: AttributesImpl) -> None: """Signal the start of an element. Args: tag: The element tag. attributes: The element attributes. """ + tag = name + attributes = attrs if tag == self.GRANT_DOC_ELEMENT: self.doc = DoclingDocument(name="file") self.text = "" @@ -762,12 +774,13 @@ def skippedEntity(self, name): self.text += unescaped @override - def endElement(self, tag): + def endElement(self, name: str) -> None: """Signal the end of an element. Args: tag: The element tag. """ + tag = name if tag == self.GRANT_DOC_ELEMENT: self._clean_data() self._end_registered_element(tag) @@ -1018,9 +1031,7 @@ def get_last_text_item(self) -> TextItem | None: else: return None text_list: list[TextItem] = [ - item - for item in self.doc.texts - if isinstance(item, TextItem) and item.get_ref() in children + item for item in self.doc.texts if item.get_ref() in children ] if text_list: @@ -1144,6 +1155,7 @@ def store_content(self, section: str, field: str, value: str) -> None: parent=self.parents[self.level], ) + @override def parse(self, patent_content: str) -> DoclingDocument | None: self.doc = self.doc = DoclingDocument(name="file") section: str = "" @@ -1151,21 +1163,21 @@ def parse(self, patent_content: str) -> DoclingDocument | None: value: str = "" line_num = 0 for line in patent_content.splitlines(): - cols = re.split("\\s{2,}", line, maxsplit=1) + cols = re.split(r"\s{2,}", line, maxsplit=1) if key and value and (len(cols) == 1 or (len(cols) == 2 and cols[0])): self.store_content(section, key, value) key = "" value = "" if len(cols) == 1: # section title - section = cols[0] + section = str(cols[0]) self.store_section(section) _log.debug(f"Parsing section {section}") elif len(cols) == 2: # key value if cols[0]: # key present - key = cols[0] - value = cols[1] + key = str(cols[0]) + value = str(cols[1]) elif not re.match(r"^##STR\d+##$", cols[1]): # line continues - value += " " + cols[1] + value += " " + str(cols[1]) line_num += 1 if key and value: self.store_content(section, key, value) @@ -1260,13 +1272,13 @@ class Element(Enum): MATH = "math-cwu", False @override - def __new__(cls, value: str, _) -> Self: + def __new__(cls, value: object, _: object = None) -> Self: obj = object.__new__(cls) obj._value_ = value return obj @override - def __init__(self, _, is_text: bool) -> None: + def __init__(self, _: object = None, is_text: bool = False) -> None: self.is_text: bool = is_text @override @@ -1289,13 +1301,15 @@ def __init__(self) -> None: self.style_html = HtmlEntity() @override - def startElement(self, tag, attributes): + def startElement(self, name: str, attrs: AttributesImpl) -> None: """Signal the start of an element. Args: tag: The element tag. attributes: The element attributes. """ + tag = name + attributes = attrs if tag == self.APP_DOC_ELEMENT: self.doc = DoclingDocument(name="file") self.text = "" @@ -1336,12 +1350,13 @@ def skippedEntity(self, name): self.text += unescaped @override - def endElement(self, tag): + def endElement(self, name: str) -> None: """Signal the end of an element. Args: tag: The element tag. """ + tag = name if tag == self.APP_DOC_ELEMENT: self._clean_data() self._end_registered_element(tag) @@ -1533,11 +1548,11 @@ class XmlTable: class ColInfo(TypedDict): ncols: int - colinfo: list[dict] + colinfo: list[dict[str, str | list[str] | None]] class MinColInfoType(TypedDict): - offset: list[int] - colwidth: list[int] + offset: list[float] + colwidth: list[float] class ColInfoType(MinColInfoType): cell_range: list[int] @@ -1581,9 +1596,12 @@ def _create_tg_range(self, tgs: list[ColInfo]) -> dict[int, ColInfoType]: "cell_range": [], "cell_offst": [0], } - offst = 0 + offst: float = 0 for info in tg["colinfo"]: - cw = info["colwidth"] + width = info["colwidth"] + if not isinstance(width, str): + raise ValueError("Column width must be text.") + cw = width cw = re.sub("pt", "", cw, flags=re.I) cw = re.sub("mm", "", cw, flags=re.I) try: @@ -1669,8 +1687,6 @@ def _parse_table(self, table: Tag) -> TableData: tg_secs = table("tgroup") if tg_secs: for tg_sec in tg_secs: - if not isinstance(tg_sec, Tag): - continue col_val = tg_sec.get("cols") ncols = ( int(col_val) @@ -1681,8 +1697,6 @@ def _parse_table(self, table: Tag) -> TableData: cs_secs = tg_sec("colspec") if cs_secs: for cs_sec in cs_secs: - if not isinstance(cs_sec, Tag): - continue colname = cs_sec.get("colname") colwidth = cs_sec.get("colwidth") tg_align["colinfo"].append( @@ -1708,15 +1722,11 @@ def _parse_table(self, table: Tag) -> TableData: tg_secs = table("tgroup") if tg_secs: for itg, tg_sec in enumerate(tg_secs): - if not isinstance(tg_sec, Tag): - continue tg_range = tgs_range[itg] row_secs = tg_sec(["row", "tr"]) if row_secs: for row_sec in row_secs: - if not isinstance(row_sec, Tag): - continue entry_secs = row_sec(["entry", "td"]) is_header: bool = ( row_sec.parent is not None @@ -1729,8 +1739,6 @@ def _parse_table(self, table: Tag) -> TableData: if entry_secs: wrong_nbr_cols = False for ientry, entry_sec in enumerate(entry_secs): - if not isinstance(entry_sec, Tag): - continue text = entry_sec.get_text().strip() # start-end diff --git a/docling/backend/xml/xbrl_backend.py b/docling/backend/xml/xbrl_backend.py index 224cadada2..473a66ec87 100644 --- a/docling/backend/xml/xbrl_backend.py +++ b/docling/backend/xml/xbrl_backend.py @@ -42,7 +42,11 @@ from docling.backend.abstract_backend import DeclarativeDocumentBackend from docling.backend.html_backend import HTMLDocumentBackend -from docling.datamodel.backend_options import HTMLBackendOptions, XBRLBackendOptions +from docling.datamodel.backend_options import ( + BaseBackendOptions, + HTMLBackendOptions, + XBRLBackendOptions, +) from docling.datamodel.base_models import InputFormat from docling.datamodel.document import InputDocument from docling.exceptions import DocumentLoadError, OperationNotAllowed @@ -50,10 +54,10 @@ _XBRL_AVAILABLE: bool = False _XBRL_IMPORT_ERROR: ImportError | None = None try: - from arelle import Cntlr # type: ignore - from arelle.ModelDocument import Type # type: ignore - from arelle.ModelDtsObject import ModelConcept # type: ignore - from arelle.ModelXbrl import ModelXbrl # type: ignore + from arelle import Cntlr + from arelle.ModelDocument import Type + from arelle.ModelDtsObject import ModelConcept + from arelle.ModelXbrl import ModelXbrl _XBRL_AVAILABLE = True except ImportError as e: @@ -87,10 +91,12 @@ def __init__( self, in_doc: InputDocument, path_or_stream: BytesIO | Path, - options: XBRLBackendOptions | None = None, + options: BaseBackendOptions | None = None, ) -> None: if options is None: options = XBRLBackendOptions() + if not isinstance(options, XBRLBackendOptions): + raise TypeError("Expected XBRLBackendOptions.") # Check if arelle is available before proceeding if not _XBRL_AVAILABLE: raise ImportError( @@ -214,14 +220,16 @@ def _get_hierarchy_cell( concept: ModelConcept, ) -> int: """Get existing or create new cell for a concept node.""" - qname_str = str(concept.qname) + qname = concept.qname + assert qname is not None + qname_str = str(qname) if qname_str not in self._hierarchy_cell_ids: cell_id = self._kv_idx self._cells.append( GraphCell( label=GraphCellLabel.KEY, cell_id=cell_id, - text=concept.qname.localName, + text=qname.localName, orig=qname_str, ) ) @@ -278,7 +286,10 @@ def convert(self) -> DoclingDocument: if fact.qname.localName == "DocumentPeriodEndDate" and fact.value: doc_period = fact.value title = f"{doc_type} {doc_org} {doc_period}".strip() - title = title if title else self.model_xbrl.modelDocument.basename + model_document = self.model_xbrl.modelDocument + title = title or ( + model_document.basename if model_document is not None else self.file.name + ) doc.add_title(text=title) # Text blocks (as HTML) @@ -322,18 +333,21 @@ def convert(self) -> DoclingDocument: # period period_text = "" if fact.context is not None: - if fact.context.isInstantPeriod: - period_text = str(fact.context.instantDatetime.date()) - elif fact.context.isStartEndPeriod: - period_text = f"{fact.context.startDatetime.date()} - {fact.context.endDatetime.date()}" + instant = fact.context.instantDatetime + start = fact.context.startDatetime + end = fact.context.endDatetime + if fact.context.isInstantPeriod and instant is not None: + period_text = str(instant.date()) + elif ( + fact.context.isStartEndPeriod + and start is not None + and end is not None + ): + period_text = f"{start.date()} - {end.date()}" # unit unit_text = "" - if ( - fact.unit is not None - and fact.unit.measures - and fact.unit.measures[0] - ): + if fact.unit is not None and fact.unit.measures[0]: unit_text = fact.unit.measures[0][0].localName # decimals diff --git a/docling/cli/export_utils.py b/docling/cli/export_utils.py index c7d2539dd8..56ff1347a8 100644 --- a/docling/cli/export_utils.py +++ b/docling/cli/export_utils.py @@ -3,6 +3,7 @@ import re from pathlib import Path +from typing import TypedDict import typer from docling_core.types.doc import ImageRefMode @@ -32,7 +33,22 @@ def _should_generate_export_images( ) -def _export_flags_from_formats(to_formats: list[OutputFormat]) -> dict[str, bool]: +class _ExportFlags(TypedDict): + export_json: bool + export_yaml: bool + export_html: bool + export_html_split_page: bool + export_md: bool + export_txt: bool + export_doctags: bool + export_vtt: bool + export_doclang: bool + export_dclx: bool + export_chunks: bool + export_latex: bool + + +def _export_flags_from_formats(to_formats: list[OutputFormat]) -> _ExportFlags: """Expand requested output formats into the per-format export booleans. The returned keys match the ``export_*`` keyword arguments of diff --git a/docling/cli/main.py b/docling/cli/main.py index 25291eb357..c9c2728d12 100644 --- a/docling/cli/main.py +++ b/docling/cli/main.py @@ -2,6 +2,7 @@ # SPDX-License-Identifier: MIT import datetime +import json as _json import logging import re import sys @@ -21,6 +22,7 @@ try: import rich.table import typer + from typer.core import TyperGroup except ImportError as e: missing_package = str(e).split("'")[1] if "'" in str(e) else "typer or rich" print( @@ -47,6 +49,7 @@ from docling_core.utils.file import resolve_source_to_path from pydantic import SecretStr, TypeAdapter, ValidationError from rich.console import Console +from typing_extensions import override from docling.cli.export_utils import ( _export_flags_from_formats, @@ -128,6 +131,8 @@ OcrOptions, PdfBackend, PdfPipelineOptions, + PictureDescriptionVlmEngineOptions, + PictureDescriptionVlmOptions, PipelineOptions, ProcessingPipeline, TableFormerMode, @@ -145,17 +150,10 @@ # The local model stack (scipy, torch, …) is absent on lightweight installs # (docling-slim[service-client] / docling-client). Guard these imports so the # CLI module — and `convert-remote` — remain importable without them. -_local_model_stack_available: bool try: - from docling.models.factories import ( - get_layout_factory, - get_ocr_factory, - get_table_structure_factory, - ) - - _local_model_stack_available = True + from docling.models import factories as _model_factories except ImportError: - _local_model_stack_available = False + _model_factories = None if TYPE_CHECKING: from docling.models.factories.base_factory import BaseFactory @@ -283,14 +281,18 @@ def _expand_from_formats(from_formats: list[str] | None) -> list[InputFormat]: return list(dict.fromkeys(expanded_formats)) -if _local_model_stack_available: - ocr_factory_internal = get_ocr_factory(allow_external_plugins=False) +if _model_factories is not None: + ocr_factory_internal = _model_factories.get_ocr_factory( + allow_external_plugins=False + ) ocr_engines_enum_internal = ocr_factory_internal.get_enum() - layout_factory_internal = get_layout_factory(allow_external_plugins=False) + layout_factory_internal = _model_factories.get_layout_factory( + allow_external_plugins=False + ) layout_engines_enum_internal = layout_factory_internal.get_enum() - table_structure_factory_internal = get_table_structure_factory( + table_structure_factory_internal = _model_factories.get_table_structure_factory( allow_external_plugins=False ) table_structure_engines_enum_internal = table_structure_factory_internal.get_enum() @@ -340,7 +342,7 @@ def _expand_from_formats(from_formats: list[str] | None) -> list[InputFormat]: """ -class _DefaultCommandGroup(typer.core.TyperGroup): +class _DefaultCommandGroup(TyperGroup): """Route a bare ``docling `` invocation to the ``convert`` command. Historically the CLI exposed a single command, so Typer let users run @@ -353,6 +355,7 @@ class _DefaultCommandGroup(typer.core.TyperGroup): default_command = "convert" + @override def parse_args(self, ctx, args): if args and args[0] not in self.commands and args[0] not in ("--help", "-h"): args = [self.default_command, *args] @@ -459,11 +462,19 @@ def version_callback(value: bool): def show_external_plugins_callback(value: bool): if value: - ocr_factory_all = get_ocr_factory(allow_external_plugins=True) - layout_factory_all = get_layout_factory(allow_external_plugins=True) - table_factory_all = get_table_structure_factory(allow_external_plugins=True) + if _model_factories is None: + raise typer.BadParameter( + "External plugins require the local model dependencies" + ) + ocr_factory_all = _model_factories.get_ocr_factory(allow_external_plugins=True) + layout_factory_all = _model_factories.get_layout_factory( + allow_external_plugins=True + ) + table_factory_all = _model_factories.get_table_structure_factory( + allow_external_plugins=True + ) - def print_external_plugins(factory: BaseFactory, factory_name: str): + def print_external_plugins(factory: BaseFactory[Any], factory_name: str): table = rich.table.Table(title=f"Available {factory_name} engines") table.add_column("Name", justify="right") table.add_column("Plugin") @@ -534,10 +545,7 @@ def export_documents( # Initialize chunker once for all documents chunker_obj = None if export_chunks: - import json as _json - from docling_core.transforms.chunker.hierarchical_chunker import ( - DocChunk, HierarchicalChunker, ) from docling_core.transforms.chunker.hybrid_chunker import ( @@ -681,13 +689,19 @@ def export_documents( # Export Chunks format: if export_chunks and chunker_obj is not None: + from docling_core.transforms.chunker.hierarchical_chunker import ( + DocChunk, + ) + from docling_core.transforms.chunker.hybrid_chunker import HybridChunker + fname = output_file or output_dir / f"{doc_filename}.chunks.jsonl" _log.info(f"writing Chunks output to {fname}") with fname.open("w", encoding="utf-8") as fp: for i, chunk in enumerate( chunker_obj.chunk(dl_doc=conv_res.document) ): - doc_chunk = cast(DocChunk, chunk) + assert isinstance(chunk, DocChunk) + doc_chunk = chunk page_numbers = sorted( { prov.page_no @@ -799,12 +813,12 @@ def convert( # noqa: C901 help="PDF files to convert. Can be local file / directory paths or URL.", ), ], - from_formats: list[str] = typer.Option( + from_formats: list[str] | None = typer.Option( None, "--from", help="Input formats to accept. Use 'odf' for odt, ods, and odp. Defaults to all.", ), - to_formats: list[OutputFormat] = typer.Option( + to_formats: list[OutputFormat] | None = typer.Option( None, "--to", help="Specify output formats. Defaults to Markdown." ), chunker_type: ChunkerType = typer.Option( @@ -829,12 +843,12 @@ def convert( # noqa: C901 help="If enabled, the page images will show the bounding-boxes of the items.", ), ] = False, - headers: str = typer.Option( + headers: str | None = typer.Option( None, "--headers", help="Specify http request headers used when fetching url input sources in the form of a JSON string", ), - html_image_headers: str = typer.Option( + html_image_headers: str | None = typer.Option( None, "--html-image-headers", help="Specify http request headers used when fetching HTML and EPUB image resources in the form of a JSON string. They are only sent to the source document's origin, or to the origins given with --html-image-headers-origin.", @@ -1259,6 +1273,11 @@ def convert( # noqa: C901 PowerpointFormatOption, WordFormatOption, ) + from docling.models.factories import ( + get_layout_factory, + get_ocr_factory, + get_table_structure_factory, + ) from docling.pipeline.asr_pipeline import AsrPipeline from docling.pipeline.legacy_standard_pdf_pipeline import LegacyStandardPdfPipeline from docling.pipeline.standard_pdf_pipeline import StandardPdfPipeline @@ -1317,14 +1336,14 @@ def _resolve_pdf_backend() -> tuple[type[PdfDocumentBackend], PdfBackendOptions] settings.perf.page_batch_size = page_batch_size requested_from_formats = from_formats - from_formats = _expand_from_formats(from_formats) + allowed_formats = _expand_from_formats(from_formats) if pipeline == ProcessingPipeline.NATIVE: # The native pipeline reads the native content of a PDF; it has nothing to # offer for the other input formats, so it never silently handles them. if requested_from_formats is None: - from_formats = [InputFormat.PDF] - elif set(from_formats) != {InputFormat.PDF}: + allowed_formats = [InputFormat.PDF] + elif set(allowed_formats) != {InputFormat.PDF}: err_console.print( "[red]Error: --pipeline native is only available for --from pdf.[/red]" ) @@ -1369,7 +1388,7 @@ def _resolve_pdf_backend() -> tuple[type[PdfDocumentBackend], PdfBackendOptions] input_doc_paths: list[Path | str] = [] for src in source: try: - if _is_http_url(src) and _is_html_source(src, from_formats): + if _is_http_url(src) and _is_html_source(src, allowed_formats): input_doc_paths.append(src) continue @@ -1377,12 +1396,14 @@ def _resolve_pdf_backend() -> tuple[type[PdfDocumentBackend], PdfBackendOptions] if local_path.exists(): if local_path.is_dir(): input_doc_paths.extend( - _iter_input_paths_from_directory(local_path, from_formats) + _iter_input_paths_from_directory( + local_path, allowed_formats + ) ) elif _is_office_lock_file(local_path): _log.info(f"Ignoring temporary Office file: {local_path}") - elif _is_html_source(src, from_formats) or _is_latex_source( - local_path, from_formats + elif _is_html_source(src, allowed_formats) or _is_latex_source( + local_path, allowed_formats ): # Keep the file in place: these backends resolve images # and included files relative to the document. @@ -1410,7 +1431,9 @@ def _resolve_pdf_backend() -> tuple[type[PdfDocumentBackend], PdfBackendOptions] local_path = TypeAdapter(Path).validate_python(src) if local_path.exists() and local_path.is_dir(): input_doc_paths.extend( - _iter_input_paths_from_directory(local_path, from_formats) + _iter_input_paths_from_directory( + local_path, allowed_formats + ) ) elif local_path.exists(): if _is_office_lock_file(local_path): @@ -1476,7 +1499,7 @@ def _resolve_pdf_backend() -> tuple[type[PdfDocumentBackend], PdfBackendOptions] ): ocr_options.psm = psm accelerator_options = AcceleratorOptions(num_threads=num_threads, device=device) - pipeline_options: PipelineOptions + pipeline_options: PipelineOptions | None = None format_options: dict[InputFormat, FormatOption] = {} backend, pdf_backend_options = _resolve_pdf_backend() @@ -1523,9 +1546,14 @@ def _resolve_pdf_backend() -> tuple[type[PdfDocumentBackend], PdfBackendOptions] pipeline_options.table_structure_options.do_cell_matching = True pipeline_options.table_structure_options.mode = table_mode if picture_description_max_new_tokens is not None: - pipeline_options.picture_description_options.generation_config[ - "max_new_tokens" - ] = picture_description_max_new_tokens + picture_options = pipeline_options.picture_description_options + assert isinstance( + picture_options, + (PictureDescriptionVlmOptions, PictureDescriptionVlmEngineOptions), + ) + picture_options.generation_config["max_new_tokens"] = ( + picture_description_max_new_tokens + ) if _should_generate_export_images( image_export_mode, @@ -1556,9 +1584,14 @@ def _resolve_pdf_backend() -> tuple[type[PdfDocumentBackend], PdfBackendOptions] if artifacts_path is not None: simple_format_option.artifacts_path = artifacts_path if picture_description_max_new_tokens is not None: - simple_format_option.picture_description_options.generation_config[ - "max_new_tokens" - ] = picture_description_max_new_tokens + picture_options = simple_format_option.picture_description_options + assert isinstance( + picture_options, + (PictureDescriptionVlmOptions, PictureDescriptionVlmEngineOptions), + ) + picture_options.generation_config["max_new_tokens"] = ( + picture_description_max_new_tokens + ) html_backend_options: HTMLBackendOptions | None = None if ( @@ -1754,7 +1787,7 @@ def _resolve_pdf_backend() -> tuple[type[PdfDocumentBackend], PdfBackendOptions] # scipy, so we avoid paying that cost unless video input is used. # Check the expanded inputs, not the raw sources, so that videos found # in a directory or downloaded from a URL get these options too. - has_video_source = InputFormat.VIDEO in from_formats and any( + has_video_source = InputFormat.VIDEO in allowed_formats and any( _name_matches_format(str(path), InputFormat.VIDEO) for path in input_doc_paths ) @@ -1795,11 +1828,12 @@ def _resolve_pdf_backend() -> tuple[type[PdfDocumentBackend], PdfBackendOptions] # Common options for all pipelines if artifacts_path is not None: - pipeline_options.artifacts_path = artifacts_path + if pipeline_options is not None: + pipeline_options.artifacts_path = artifacts_path asr_pipeline_options.artifacts_path = artifacts_path doc_converter = DocumentConverter( - allowed_formats=from_formats, + allowed_formats=allowed_formats, format_options=format_options, ) diff --git a/docling/cli/remote.py b/docling/cli/remote.py index aef469221e..8177ef9a8d 100644 --- a/docling/cli/remote.py +++ b/docling/cli/remote.py @@ -360,7 +360,7 @@ def convert_remote( option_kwargs["image_export_mode"] = image_export_mode if parsed_page_range is not None: option_kwargs["page_range"] = parsed_page_range - options = ConvertDocumentsRequestOptions(**option_kwargs) + options = ConvertDocumentsRequestOptions.model_validate(option_kwargs) # Local export must match what the service produced: resolved value (the # user's override, or the options model default when unset). resolved_image_mode = options.image_export_mode diff --git a/docling/datamodel/asr_model_specs.py b/docling/datamodel/asr_model_specs.py index 5133ef7aa0..9e76086f3d 100644 --- a/docling/datamodel/asr_model_specs.py +++ b/docling/datamodel/asr_model_specs.py @@ -41,7 +41,7 @@ def _detect_hardware_and_libraries(): # Check if mlx-whisper is available try: - import mlx_whisper # type: ignore + import mlx_whisper # ty: ignore[unresolved-import] has_mlx_whisper = True except ImportError: diff --git a/docling/datamodel/backend_options.py b/docling/datamodel/backend_options.py index e11dd48d27..b764070682 100644 --- a/docling/datamodel/backend_options.py +++ b/docling/datamodel/backend_options.py @@ -15,7 +15,6 @@ PositiveInt, PrivateAttr, SecretStr, - conint, field_validator, model_validator, ) @@ -327,7 +326,7 @@ class ThreadedDoclingParseBackendOptions(PdfBackendOptions): "are consumed at, so pages are not rendered a second time on request." ), ) - release_native_memory_every_n_pages: conint(ge=0) = Field( + release_native_memory_every_n_pages: Annotated[int, Field(ge=0)] = Field( 128, description=( "Release native parser memory after every N decoded pages in the " diff --git a/docling/datamodel/base_models.py b/docling/datamodel/base_models.py index e337a8502e..82f649236f 100644 --- a/docling/datamodel/base_models.py +++ b/docling/datamodel/base_models.py @@ -38,6 +38,7 @@ field_validator, model_validator, ) +from typing_extensions import override if TYPE_CHECKING: from docling.backend.pdf_backend import PdfPageBackend @@ -543,7 +544,7 @@ def get_image( if max_size: assert self.size is not None - scale = min(scale, max_size / max(self.size.as_tuple())) + scale = min(scale, max_size / max(self.size.width, self.size.height)) if scale not in self._image_cache: if cropbox is None: @@ -664,17 +665,17 @@ def _score_to_grade(self, score: ScoreValue) -> QualityGrade: return QualityGrade.UNSPECIFIED - @computed_field # type: ignore + @computed_field @property def mean_grade(self) -> QualityGrade: return self._score_to_grade(self.mean_score) - @computed_field # type: ignore + @computed_field @property def low_grade(self) -> QualityGrade: return self._score_to_grade(self.low_score) - @computed_field # type: ignore + @computed_field @property def mean_score(self) -> ScoreValue: return ScoreValue( @@ -688,7 +689,7 @@ def mean_score(self) -> ScoreValue: ) ) - @computed_field # type: ignore + @computed_field @property def low_score(self) -> ScoreValue: return ScoreValue( @@ -745,8 +746,9 @@ def _accept_flat_confidence_scores(cls, value, handler): model._low_score_override = low_override return model - @computed_field # type: ignore + @computed_field @property + @override def mean_score(self) -> ScoreValue: if not np.isnan(self._mean_score_override): return self._mean_score_override @@ -758,8 +760,9 @@ def mean_score(self) -> ScoreValue: ) ) - @computed_field # type: ignore + @computed_field @property + @override def low_score(self) -> ScoreValue: if not np.isnan(self._low_score_override): return self._low_score_override diff --git a/docling/datamodel/chart_extraction_options.py b/docling/datamodel/chart_extraction_options.py index ac7658fcec..8f0eddb015 100644 --- a/docling/datamodel/chart_extraction_options.py +++ b/docling/datamodel/chart_extraction_options.py @@ -6,7 +6,7 @@ from typing import Any, ClassVar, Dict, Iterator, Literal, Optional from pydantic import Field, model_validator -from typing_extensions import Self +from typing_extensions import Self, override from docling.datamodel.stage_model_specs import StagePresetMixin, VlmModelSpec from docling.models.inference_engines.vlm.base import VlmEngineOptionsMixin @@ -164,10 +164,12 @@ class _ChartExtractionModelKindMeta(type): ``ChartExtractionModelKind["GRANITE_VISION"]`` does not raise TypeError. """ - def __iter__(cls) -> Iterator[Any]: + _members: ClassVar[dict[str, "ChartExtractionModelKind._Member"]] + + def __iter__(cls) -> Iterator["ChartExtractionModelKind._Member"]: return iter(cls._members.values()) - def __getitem__(cls, item: str) -> Any: + def __getitem__(cls, item: str) -> "ChartExtractionModelKind._Member": return cls._members[item] def __contains__(cls, item: object) -> bool: @@ -196,12 +198,15 @@ def __init__(self, name: str, value: str) -> None: self.name = name self.value = value + @override def __repr__(self) -> str: return f"" + @override def __str__(self) -> str: return self.value + @override def __eq__(self, other: object) -> bool: if isinstance(other, str): return self.value == other @@ -209,10 +214,11 @@ def __eq__(self, other: object) -> bool: return self.value == other.value return NotImplemented + @override def __hash__(self) -> int: return hash(self.value) - _members: ClassVar[Dict[str, "_ChartExtractionModelKindMeta"]] = {} + _members: ClassVar[dict[str, _Member]] = {} # Map old enum values to new preset IDs (V1 → V4 with deprecation) _PRESET_MAP: ClassVar[Dict[str, str]] = { @@ -227,7 +233,7 @@ def __hash__(self) -> int: ("GRANITE_VISION_V4", ChartExtractionModelKind.GRANITE_VISION_V4), ]: _member = ChartExtractionModelKind._Member(_name, _val) - ChartExtractionModelKind._members[_name] = _member # type: ignore[assignment] + ChartExtractionModelKind._members[_name] = _member setattr(ChartExtractionModelKind, _name, _member) diff --git a/docling/datamodel/document.py b/docling/datamodel/document.py index e6d114f5f5..db328e0df8 100644 --- a/docling/datamodel/document.py +++ b/docling/datamodel/document.py @@ -49,6 +49,7 @@ TypeAdapter, ValidationError, ) +from typing_extensions import override from docling.backend.abstract_backend import ( AbstractDocumentBackend, @@ -142,7 +143,7 @@ class InputDocument(BaseModel): limits: DocumentLimits = Field( DocumentLimits(), description="Limits in the input document for the conversion." ) - format: Annotated[InputFormat, Field(description="The document format.")] + format: Annotated[InputFormat | None, Field(description="The document format.")] filesize: Optional[int] = Field( None, description="Size of the input file, in bytes." @@ -156,7 +157,7 @@ class InputDocument(BaseModel): def __init__( self, path_or_stream: Union[BytesIO, Path], - format: InputFormat, + format: InputFormat | None, backend: Type[AbstractDocumentBackend], backend_options: Optional[BackendOptions] = None, filename: Optional[str] = None, @@ -601,17 +602,21 @@ class _DummyBackend(AbstractDocumentBackend): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) + @override def is_valid(self) -> bool: return False @classmethod + @override def supported_formats(cls) -> set[InputFormat]: return set() @classmethod + @override def supports_pagination(cls) -> bool: return False + @override def unload(self): return super().unload() @@ -643,6 +648,7 @@ def docs( self, format_options: Mapping[InputFormat, "BaseFormatOption"], ) -> Iterable[InputDocument]: + limits = self.limits or DocumentLimits() for item in self.path_or_stream_iterator: # `backend_input` is what backend_options_for_input() sees: the raw # URL string (not the HttpSource model) so HTML source_uri resolution @@ -662,7 +668,7 @@ def docs( obj = resolve_source_to_stream( source_uri, req_headers, - max_file_size=self.limits.max_file_size, + max_file_size=limits.max_file_size, ) except FileSizeLimitExceededError as exc: yield self._build_invalid_input_document( @@ -672,7 +678,7 @@ def docs( rejection=InputRejection( message=( f"File size {exc.size} exceeds the max_file_size " - f"limit of {self.limits.max_file_size} bytes." + f"limit of {limits.max_file_size} bytes." ), category=FailureCategory.POLICY, ), @@ -721,7 +727,7 @@ def docs( yield InputDocument( path_or_stream=path_or_stream, - format=format, # type: ignore[arg-type] + format=format, filename=obj.name, limits=self.limits, backend=backend, @@ -820,7 +826,7 @@ def _guess_format(self, obj: Union[Path, DocumentStream]) -> Optional[InputForma if office_mime is not None: mime = office_mime - elif isinstance(obj, DocumentStream): + else: if _DocumentConversionInput._has_doclang_extension(obj.name): return InputFormat.XML_DOCLANG if _DocumentConversionInput._has_dclx_extension(obj.name): diff --git a/docling/datamodel/pipeline_options.py b/docling/datamodel/pipeline_options.py index d3f52a7a6b..66579f1d0c 100644 --- a/docling/datamodel/pipeline_options.py +++ b/docling/datamodel/pipeline_options.py @@ -21,7 +21,7 @@ field_validator, model_validator, ) -from typing_extensions import deprecated +from typing_extensions import deprecated, override from docling.datamodel import ( asr_model_specs, @@ -1677,6 +1677,7 @@ class LayoutOptions(BaseLayoutOptions): ), ] = DOCLING_LAYOUT_HERON + @override def model_post_init(self, context: Any, /) -> None: super().model_post_init(context) warnings.warn( diff --git a/docling/datamodel/service/options.py b/docling/datamodel/service/options.py index 081c155cd8..200d97a9b9 100644 --- a/docling/datamodel/service/options.py +++ b/docling/datamodel/service/options.py @@ -661,7 +661,9 @@ class ConvertDocumentsOptions(BaseModel): ] = None chart_extraction_custom_config: Annotated[ - Optional[SerializeAsAny[Union[ChartExtractionVlmEngineOptions, dict]]], + Optional[ + SerializeAsAny[Union[ChartExtractionVlmEngineOptions, dict[str, Any]]] + ], Field( default=None, description=( @@ -806,7 +808,7 @@ class ConvertDocumentsOptions(BaseModel): # Option 2: Custom configuration (if allowed by config) vlm_pipeline_custom_config: Annotated[ - Optional[Union[VlmConvertOptions, dict]], + Optional[Union[VlmConvertOptions, dict[str, Any]]], Field( default=None, description="Custom VLM configuration including model spec and engine options. " @@ -815,7 +817,7 @@ class ConvertDocumentsOptions(BaseModel): ] = None picture_description_custom_config: Annotated[ - Optional[Union[PictureDescriptionVlmEngineOptions, dict]], + Optional[Union[PictureDescriptionVlmEngineOptions, dict[str, Any]]], Field( default=None, description="Custom picture description configuration including model spec and engine options.", @@ -823,7 +825,7 @@ class ConvertDocumentsOptions(BaseModel): ] = None code_formula_custom_config: Annotated[ - Optional[Union[CodeFormulaVlmOptions, dict]], + Optional[Union[CodeFormulaVlmOptions, dict[str, Any]]], Field( default=None, description="Custom code/formula extraction configuration including model spec and engine options.", diff --git a/docling/datamodel/service/requests.py b/docling/datamodel/service/requests.py index bb034bda45..ed88338c79 100644 --- a/docling/datamodel/service/requests.py +++ b/docling/datamodel/service/requests.py @@ -4,7 +4,7 @@ import enum from collections.abc import Mapping from functools import cache -from typing import Annotated, Any, Generic, Literal, TypeAlias, get_args +from typing import Annotated, Any, Generic, Literal, TypeAlias, cast, get_args from pydantic import ( AnyHttpUrl, @@ -261,9 +261,12 @@ def make_request_model( Dynamically create (and cache) a subclass of GenericChunkDocumentsRequest[opt_type] with chunking_options having a default factory. """ - return type( + parent = GenericChunkDocumentsRequest.__class_getitem__(opt_type) + assert isinstance(parent, type) + model = type( f"{opt_type.__name__}DocumentsRequest", - (GenericChunkDocumentsRequest[opt_type],), # type: ignore[valid-type] + # ty cannot resolve a runtime-parametrized Pydantic base class. + (parent,), # ty: ignore[unsupported-dynamic-base] { "__annotations__": {"chunking_options": opt_type}, "chunking_options": Field( @@ -271,3 +274,5 @@ def make_request_model( ), }, ) + assert issubclass(model, GenericChunkDocumentsRequest) + return cast(type[GenericChunkDocumentsRequest[ChunkingOptT]], model) diff --git a/docling/datamodel/service/responses.py b/docling/datamodel/service/responses.py index b144c3914d..53ee349d0a 100644 --- a/docling/datamodel/service/responses.py +++ b/docling/datamodel/service/responses.py @@ -5,7 +5,7 @@ import math import warnings from datetime import datetime -from typing import TYPE_CHECKING, Annotated, Literal, Optional +from typing import TYPE_CHECKING, Annotated, Any, Literal, Optional from docling_core.types.doc.document import DoclingDocument from pydantic import AliasChoices, AnyUrl, BaseModel, ConfigDict, Field @@ -23,7 +23,7 @@ from docling.datamodel.base_models import ConfidenceReport, PageConfidenceScores -def _nan_to_none(value: float) -> Optional[float]: +def _nan_to_none(value: float | None) -> Optional[float]: return None if value is None or math.isnan(value) else float(value) @@ -191,7 +191,8 @@ class ChunkedDocumentResultItem(BaseModel): Field(description="Page numbers where this chunk content appears"), ] = None metadata: Annotated[ - dict | None, Field(description="Additional metadata associated with this chunk") + dict[str, Any] | None, + Field(description="Additional metadata associated with this chunk"), ] = None @@ -199,7 +200,7 @@ class ChunkedDocumentResult(BaseModel): kind: Literal["ChunkedDocumentResponse"] = "ChunkedDocumentResponse" chunks: list[ChunkedDocumentResultItem] documents: list[ExportResult] - chunking_info: Optional[dict] = None + chunking_info: Optional[dict[str, Any]] = None class PresignedArtifactResult(BaseModel): diff --git a/docling/datamodel/stage_model_specs.py b/docling/datamodel/stage_model_specs.py index 8439bd1142..d5cb69d05c 100644 --- a/docling/datamodel/stage_model_specs.py +++ b/docling/datamodel/stage_model_specs.py @@ -16,6 +16,7 @@ from typing import TYPE_CHECKING, Any, ClassVar, Dict, List, Set from pydantic import BaseModel, Field +from typing_extensions import NotRequired, TypedDict from docling.datamodel.pipeline_options_vlm_model import ( ResponseFormat, @@ -693,7 +694,7 @@ def from_preset( # a single validated instance so that model_validators (e.g. "at least one # output") run against the final merged state rather than the pre-override # state. - preset_data: dict = dict( + preset_data: dict[str, Any] = dict( model_spec=preset.model_spec, engine_options=engine_options, scale=preset.scale, @@ -806,7 +807,9 @@ def from_preset( f"Unsupported engine type {preset.default_engine_type} for presets" ) - instance = cls( # type: ignore[call-arg] + if not issubclass(cls, BaseModel): + raise TypeError("Preset options must be Pydantic models") + instance = cls( model_spec=preset.model_spec, engine_options=engine_options, **preset.stage_options, @@ -921,7 +924,9 @@ def from_preset( f"Unsupported engine type {preset.default_engine_type} for presets" ) - instance = cls( # type: ignore[call-arg] + if not issubclass(cls, BaseModel): + raise TypeError("Preset options must be Pydantic models") + instance = cls( model_spec=preset.model_spec, engine_options=engine_options, **preset.stage_options, @@ -941,9 +946,20 @@ def from_preset( # SHARED MODEL SPECS (for reuse across multiple stages) # ----------------------------------------------------------------------------- + +class _SharedVlmModelSpec(TypedDict): + name: str + default_repo_id: str + stop_strings: NotRequired[list[str]] + max_new_tokens: NotRequired[int] + engine_overrides: NotRequired[dict[VlmEngineType, EngineModelConfig]] + api_overrides: NotRequired[dict[VlmEngineType, ApiModelConfig]] + supported_engines: NotRequired[set[VlmEngineType]] + + # Shared Granite Docling model spec used across VLM_CONVERT and CODE_FORMULA stages # Note: prompt and response_format are intentionally excluded here as they vary per stage -GRANITE_DOCLING_MODEL_SPEC_BASE = { +GRANITE_DOCLING_MODEL_SPEC_BASE: _SharedVlmModelSpec = { "name": "Granite-Docling-258M", "default_repo_id": "ibm-granite/granite-docling-258M", "stop_strings": ["", "<|end_of_text|>"], @@ -967,7 +983,7 @@ def from_preset( } # Shared Pixtral model spec used across VLM_CONVERT and PICTURE_DESCRIPTION stages -PIXTRAL_MODEL_SPEC_BASE = { +PIXTRAL_MODEL_SPEC_BASE: _SharedVlmModelSpec = { "name": "Pixtral-12B", "default_repo_id": "mistral-community/pixtral-12b", "engine_overrides": { @@ -981,7 +997,7 @@ def from_preset( } # Shared Granite Vision model spec used across VLM_CONVERT and PICTURE_DESCRIPTION stages -GRANITE_VISION_MODEL_SPEC_BASE = { +GRANITE_VISION_MODEL_SPEC_BASE: _SharedVlmModelSpec = { "name": "Granite-Vision-3.3-2B", "default_repo_id": "ibm-granite/granite-vision-3.3-2b", "supported_engines": { diff --git a/docling/document_converter.py b/docling/document_converter.py index 9f8bb6c6f6..1a77dc2ecc 100644 --- a/docling/document_converter.py +++ b/docling/document_converter.py @@ -16,7 +16,7 @@ from docling_core.types.doc.page import TextCellUnit from pydantic import ConfigDict, Field, model_validator, validate_call -from typing_extensions import Self +from typing_extensions import Self, override from docling.backend.abstract_backend import ( AbstractDocumentBackend, @@ -122,6 +122,7 @@ class FormatOption(BaseFormatOption): pipeline_cls: Type[BasePipeline] backend_options: Optional[BackendOptions] = None + @override def backend_options_for_input( self, source: Path | str | DocumentStream ) -> BackendOptions | None: @@ -188,6 +189,7 @@ class AsciiDocFormatOption(FormatOption): backend: Type[AbstractDocumentBackend] = AsciiDocBackend backend_options: AsciiDocBackendOptions | None = None + @override def backend_options_for_input( self, source: Path | str | DocumentStream ) -> AsciiDocBackendOptions | None: @@ -209,6 +211,7 @@ class HTMLFormatOption(FormatOption): backend: Type[AbstractDocumentBackend] = HTMLDocumentBackend backend_options: Optional[HTMLBackendOptions] = None + @override def backend_options_for_input( self, source: Path | str | DocumentStream ) -> HTMLBackendOptions | None: @@ -877,6 +880,7 @@ def _execute_pipeline( self, in_doc: InputDocument, raises_on_error: bool ) -> ConversionResult: if in_doc.valid: + assert in_doc.format is not None pipeline_started = False try: pipeline = self._get_pipeline(in_doc.format) diff --git a/docling/document_extractor.py b/docling/document_extractor.py index 5176dee8f1..cbb45d5d36 100644 --- a/docling/document_extractor.py +++ b/docling/document_extractor.py @@ -277,6 +277,7 @@ def _execute_extraction_pipeline( errors=build_invalid_input_errors(in_doc), ) + assert in_doc.format is not None pipeline = self._get_pipeline(in_doc.format) if pipeline is None: if raises_on_error: diff --git a/docling/experimental/models/table_crops_layout_model.py b/docling/experimental/models/table_crops_layout_model.py index 823cfba021..7f206105e0 100644 --- a/docling/experimental/models/table_crops_layout_model.py +++ b/docling/experimental/models/table_crops_layout_model.py @@ -10,6 +10,7 @@ from typing import Optional from docling_core.types.doc import BoundingBox, DocItemLabel +from typing_extensions import override from docling.datamodel.accelerator_options import AcceleratorOptions from docling.datamodel.base_models import Cluster, LayoutPrediction, Page @@ -44,9 +45,11 @@ def __init__( self.accelerator_options = accelerator_options @classmethod + @override def get_options_type(cls) -> type[TableCropsLayoutOptions]: return TableCropsLayoutOptions + @override def predict_layout( self, conv_res: ConversionResult, @@ -87,7 +90,7 @@ def _build_page_clusters(self, page: Page) -> list[Cluster]: cells=[], ) - clusters = [cluster] + clusters: list[Cluster] = [cluster] if not self.options.skip_cell_assignment: page_cells = list(page.cells) diff --git a/docling/experimental/pipeline/threaded_layout_vlm_pipeline.py b/docling/experimental/pipeline/threaded_layout_vlm_pipeline.py index 90c16d0d36..6a2e122203 100644 --- a/docling/experimental/pipeline/threaded_layout_vlm_pipeline.py +++ b/docling/experimental/pipeline/threaded_layout_vlm_pipeline.py @@ -18,6 +18,7 @@ from docling_core.types.doc import DoclingDocument, ImageRef, PictureItem from docling_core.types.doc.document import DocTagsDocument +from typing_extensions import override if TYPE_CHECKING: from docling_core.types.doc.page import SegmentedPage @@ -113,7 +114,8 @@ def _init_models(self) -> None: # Create layout-aware VLM options internally base_vlm_options = self.pipeline_options.vlm_options - class LayoutAwareVlmOptions(type(base_vlm_options)): # type: ignore[misc] + # Preserve the selected options subclass, including its Pydantic fields. + class LayoutAwareVlmOptions(type(base_vlm_options)): # ty: ignore[unsupported-base] def build_prompt( self, page: Optional[SegmentedPage], @@ -121,7 +123,8 @@ def build_prompt( _internal_page: Optional[Page] = None, ) -> str: base_prompt = self.prompt - augmented_prompt = base_prompt + assert isinstance(base_prompt, str) + augmented_prompt: str = base_prompt # Only augment convert to docling base prompts if base_prompt != DOCLING_BASE_PAGE_PROMPT: @@ -275,6 +278,7 @@ def _create_run_ctx(self) -> RunContext: stages=stages, first_stage=layout_stage, output_queue=output_q ) + @override def _build_document(self, conv_res: ConversionResult) -> ConversionResult: """Build document using threaded layout+VLM pipeline.""" assert isinstance(conv_res.input._backend, PdfDocumentBackend) @@ -484,20 +488,25 @@ def _integrate_results( else: conv_res.status = ConversionStatus.SUCCESS + @override def _assemble_document(self, conv_res: ConversionResult) -> ConversionResult: return conv_res @classmethod + @override def get_default_options(cls) -> ThreadedLayoutVlmPipelineOptions: return ThreadedLayoutVlmPipelineOptions() @classmethod + @override def is_backend_supported(cls, backend: AbstractDocumentBackend) -> bool: return isinstance(backend, PdfDocumentBackend) + @override def _determine_status(self, conv_res: ConversionResult) -> ConversionStatus: return conv_res.status + @override def _unload(self, conv_res: ConversionResult) -> None: for p in conv_res.pages: if p._backend is not None: diff --git a/docling/models/base_layout_model.py b/docling/models/base_layout_model.py index 1ccdcc687f..e6bd741bbf 100644 --- a/docling/models/base_layout_model.py +++ b/docling/models/base_layout_model.py @@ -8,6 +8,7 @@ from typing import Type from docling_core.types.doc import DocItemLabel +from typing_extensions import override from docling.datamodel.base_models import LayoutPrediction, Page from docling.datamodel.document import ConversionResult @@ -42,6 +43,7 @@ class BaseLayoutModel(BasePageModel, BaseModelWithOptions, ABC): @classmethod @abstractmethod + @override def get_options_type(cls) -> Type[BaseLayoutOptions]: """Return the options type supported by this layout model.""" @@ -53,6 +55,7 @@ def predict_layout( ) -> Sequence[LayoutPrediction]: """Produce layout predictions for the provided pages.""" + @override def __call__( self, conv_res: ConversionResult, diff --git a/docling/models/base_layout_postprocessing_model.py b/docling/models/base_layout_postprocessing_model.py index b54b31aaab..695bbe15e1 100644 --- a/docling/models/base_layout_postprocessing_model.py +++ b/docling/models/base_layout_postprocessing_model.py @@ -7,6 +7,8 @@ from collections.abc import Iterable, Sequence from typing import Type +from typing_extensions import override + from docling.datamodel.base_models import LayoutPrediction, Page from docling.datamodel.document import ConversionResult from docling.datamodel.pipeline_options import BaseLayoutPostprocessorOptions @@ -18,6 +20,7 @@ class BaseLayoutPostprocessingModel(BasePageModel, BaseModelWithOptions, ABC): @classmethod @abstractmethod + @override def get_options_type(cls) -> Type[BaseLayoutPostprocessorOptions]: """Return the options type supported by this post-processing model.""" @@ -29,6 +32,7 @@ def postprocess_layout( ) -> Sequence[LayoutPrediction]: """Finalize raw layout predictions for the provided pages.""" + @override def __call__( self, conv_res: ConversionResult, diff --git a/docling/models/base_model.py b/docling/models/base_model.py index 0f58773928..f2c61dc227 100644 --- a/docling/models/base_model.py +++ b/docling/models/base_model.py @@ -14,8 +14,9 @@ NodeItem, PictureItem, ) +from numpy.typing import NDArray from PIL.Image import Image -from typing_extensions import TypeVar +from typing_extensions import TypeVar, override from docling.datamodel.base_models import ( ItemAndImageEnrichmentElement, @@ -52,7 +53,7 @@ class BaseVlmModel(ABC): @abstractmethod def process_images( self, - image_batch: Iterable[Union[Image, np.ndarray]], + image_batch: Iterable[Union[Image, NDArray[np.generic]]], prompt: Union[str, list[str]], ) -> Iterable[VlmPrediction]: """Process raw images without page metadata. @@ -98,6 +99,7 @@ def _build_prompt_safe(self, page: Page) -> str: return self.vlm_options.build_prompt(page.parsed_page) @abstractmethod + @override def __call__( self, conv_res: ConversionResult, page_batch: Iterable[Page] ) -> Iterable[Page]: @@ -136,6 +138,7 @@ def formulate_prompt(self, user_prompt: str) -> str: prompt = self.processor.apply_chat_template( messages, add_generation_prompt=True ) + assert isinstance(prompt, str) return prompt raise RuntimeError( @@ -167,6 +170,7 @@ def __call__( class BaseEnrichmentModel(GenericEnrichmentModel[NodeItem]): + @override def prepare_element( self, conv_res: ConversionResult, element: NodeItem ) -> Optional[NodeItem]: @@ -181,6 +185,7 @@ class BaseItemAndImageEnrichmentModel( images_scale: float expansion_factor: float = 0.0 + @override def prepare_element( self, conv_res: ConversionResult, element: NodeItem ) -> Optional[ItemAndImageEnrichmentElement]: diff --git a/docling/models/base_ocr_model.py b/docling/models/base_ocr_model.py index 405eacfc63..6bdb91088d 100644 --- a/docling/models/base_ocr_model.py +++ b/docling/models/base_ocr_model.py @@ -21,6 +21,7 @@ ) from PIL import Image, ImageDraw from scipy.ndimage import binary_dilation, find_objects, label +from typing_extensions import override from docling.datamodel.accelerator_options import AcceleratorOptions from docling.datamodel.base_models import Page @@ -42,6 +43,7 @@ def _empty_segmented_page(page: Page) -> SegmentedPdfPage: """A minimal SegmentedPdfPage for pages whose native parse was skipped (PagePreprocessingOptions.skip_cell_extraction), sized from the page.""" + assert page.size is not None width, height = page.size.width, page.size.height # CoordOrigin.BOTTOMLEFT by convention: a real backend produces the # segmented page in bottom-left coordinates. @@ -183,6 +185,7 @@ def map_ocr_language(self, language: OcrLanguage) -> str | list[str]: supported=self.supported_ocr_languages(), detail="This engine needs a BCP-47 tag behind the `iso:` prefix.", ) + assert language.bcp47_language is not None return language.bcp47_language def resolve_ocr_languages(self) -> list[str]: @@ -229,7 +232,7 @@ def get_ocr_rects(self, page: Page) -> list[BoundingBox]: ocr_rects = self._find_pdf_aware_layout_ocr_rects(page) elif self.options.mode == OcrMode.LAYOUT_REGIONS: ocr_rects = self._find_layout_ocr_rects(page) - elif self.options.mode == OcrMode.FULL_PAGE: + else: # A big bbox covering the entire page ocr_rects = [ BoundingBox( @@ -249,6 +252,7 @@ def _find_layout_ocr_rects(self, page: Page) -> list[BoundingBox]: """ if page.predictions.layout is None: return [] + assert page.size is not None # Use every layout detection bbox as an initial ocr_rect ocr_rects = [c.bbox for c in page.predictions.layout.clusters] @@ -401,7 +405,7 @@ def _deduplicate_rects( ] # Compute area fraction on page covered by bitmaps - area_frac = np.sum(np_image > 0) / (size.width * size.height) + area_frac = float(np.sum(np_image > 0) / (size.width * size.height)) return (area_frac, bounding_boxes) # fraction covered # boxes def post_process_cells( @@ -490,7 +494,8 @@ def _merge_ocr_and_pdf_cells( if segregate_pdf_cells: visible_cells, invisible_cells = _segregate_by_visibility(pdf_cells) else: - visible_cells, invisible_cells = list(pdf_cells), [] + visible_cells = list(pdf_cells) + invisible_cells = [] # The prioritized cells are always kept # the secondary cells are added only where they don't overlap a prioritized cell. @@ -598,6 +603,7 @@ def draw_ocr_rects_and_cells(self, conv_res, page, ocr_rects, show: bool = False image.save(str(out_file), format="png") @abstractmethod + @override def __call__( self, conv_res: ConversionResult, page_batch: Iterable[Page] ) -> Iterable[Page]: @@ -605,5 +611,6 @@ def __call__( @classmethod @abstractmethod + @override def get_options_type(cls) -> type[OcrOptions]: pass diff --git a/docling/models/base_table_model.py b/docling/models/base_table_model.py index 0844d27a32..198cf709fb 100644 --- a/docling/models/base_table_model.py +++ b/docling/models/base_table_model.py @@ -7,6 +7,8 @@ from collections.abc import Iterable, Sequence from typing import Type +from typing_extensions import override + from docling.datamodel.base_models import Page, TableStructurePrediction from docling.datamodel.document import ConversionResult from docling.datamodel.pipeline_options import BaseTableStructureOptions @@ -20,6 +22,7 @@ class BaseTableStructureModel(BasePageModel, BaseModelWithOptions, ABC): @classmethod @abstractmethod + @override def get_options_type(cls) -> Type[BaseTableStructureOptions]: """Return the options type supported by this table model.""" @@ -31,6 +34,7 @@ def predict_tables( ) -> Sequence[TableStructurePrediction]: """Produce table structure predictions for the provided pages.""" + @override def __call__( self, conv_res: ConversionResult, diff --git a/docling/models/extraction/nuextract_transformers_model.py b/docling/models/extraction/nuextract_transformers_model.py index 6e4457467d..9de362e774 100644 --- a/docling/models/extraction/nuextract_transformers_model.py +++ b/docling/models/extraction/nuextract_transformers_model.py @@ -9,8 +9,16 @@ from typing import Any, Optional, Union import numpy as np +from numpy.typing import NDArray from PIL.Image import Image -from transformers import AutoModelForImageTextToText, AutoProcessor, GenerationConfig + +# Transformers exposes the model lazily when its torch dependency is installed. +from transformers import ( + AutoModelForImageTextToText, # ty: ignore[possibly-missing-import] + AutoProcessor, + GenerationConfig, +) +from typing_extensions import override from docling.datamodel.accelerator_options import ( AcceleratorOptions, @@ -87,7 +95,11 @@ def extract_example_images(example_item): ) # Ensure examples batch matches messages batch if provided - if examples and len(examples_batch) != len(messages_batch): + if ( + examples + and examples_batch is not None + and len(examples_batch) != len(messages_batch) + ): if not is_batch and len(examples_batch) == 1: # Single example set for a single input is fine pass @@ -98,7 +110,7 @@ def extract_example_images(example_item): all_images = [] for i, message_group in enumerate(messages_batch): # Get example images for this input - if examples and i < len(examples_batch): + if examples and examples_batch is not None and i < len(examples_batch): input_example_images = extract_example_images(examples_batch[i]) all_images.extend(input_example_images) @@ -162,16 +174,17 @@ def __init__( trust_remote_code=vlm_options.trust_remote_code, ) if sys.version_info < (3, 14): - self.vlm_model = torch.compile(self.vlm_model) # type: ignore + self.vlm_model = torch.compile(self.vlm_model) else: self.vlm_model.eval() # Load generation config self.generation_config = GenerationConfig.from_pretrained(artifacts_path) + @override def process_images( self, - image_batch: Iterable[Union[Image, np.ndarray]], + image_batch: Iterable[Union[Image, NDArray[np.generic]]], prompt: Union[str, list[str]], ) -> Iterable[VlmPrediction]: """ @@ -286,7 +299,7 @@ def process_images( trimmed_sequences = generated_ids[:, input_len:] # Decode with the processor/tokenizer - decoded_texts: list[str] = self.processor.batch_decode( + decoded_texts = self.processor.batch_decode( trimmed_sequences, skip_special_tokens=True, clean_up_tokenization_spaces=False, @@ -294,15 +307,16 @@ def process_images( # Optional logging num_tokens = None - if generated_ids.shape[0] > 0: # type: ignore + if generated_ids.shape[0] > 0: # Todo: confirm num tokens is actually from first item, code was already like this num_tokens = int(generated_ids[0].shape[0]) _log.debug( f"Generated {num_tokens} tokens in {generation_time:.2f}s " - f"for batch size {generated_ids.shape[0]}." # type: ignore + f"for batch size {generated_ids.shape[0]}." ) for text in decoded_texts: + assert isinstance(text, str) # Apply decode_response to the output text decoded_text = self.vlm_options.decode_response(text) yield VlmPrediction( diff --git a/docling/models/extraction/prompt_utils.py b/docling/models/extraction/prompt_utils.py index 7dedba097c..9695d3ecb3 100644 --- a/docling/models/extraction/prompt_utils.py +++ b/docling/models/extraction/prompt_utils.py @@ -7,7 +7,7 @@ tokenized inputs ready for model.generate(). """ -from typing import Any +from typing import Any, cast from PIL.Image import Image @@ -109,14 +109,17 @@ def _build_extraction_prompt(template: str) -> str: ) -def _process_all_vision_info(messages: list, examples: list | None = None) -> Any: +def _process_all_vision_info( + messages: list[dict[str, Any]] | list[list[dict[str, Any]]], + examples: list[dict[str, Any]] | list[list[dict[str, Any]] | None] | None = None, +) -> Any: """Process vision info from messages using qwen-vl-utils. Adapted from NuExtract source code. """ from qwen_vl_utils import fetch_image, process_vision_info - def extract_example_images(example_item: Any) -> list: + def extract_example_images(example_item: Any) -> list[Image]: if not example_item: return [] examples_to_process = ( @@ -131,8 +134,11 @@ def extract_example_images(example_item: Any) -> list: images.append(fetch_image(example["input"])) return images - is_batch = messages and isinstance(messages[0], list) - messages_batch = messages if is_batch else [messages] + is_batch = bool(messages) and isinstance(messages[0], list) + if is_batch: + messages_batch = cast(list[list[dict[str, Any]]], messages) + else: + messages_batch = [cast(list[dict[str, Any]], messages)] is_batch_examples = ( examples and isinstance(examples, list) diff --git a/docling/models/extraction/transformers_extraction_model.py b/docling/models/extraction/transformers_extraction_model.py index 2ef5a5817a..aae80fc3fa 100644 --- a/docling/models/extraction/transformers_extraction_model.py +++ b/docling/models/extraction/transformers_extraction_model.py @@ -12,8 +12,16 @@ import numpy as np import torch +from numpy.typing import NDArray from PIL.Image import Image -from transformers import AutoModelForImageTextToText, AutoProcessor, GenerationConfig + +# Transformers exposes the model lazily when its torch dependency is installed. +from transformers import ( + AutoModelForImageTextToText, # ty: ignore[possibly-missing-import] + AutoProcessor, + GenerationConfig, +) +from typing_extensions import override from docling.datamodel.accelerator_options import AcceleratorOptions from docling.datamodel.base_models import VlmPrediction, VlmStopReason @@ -122,7 +130,7 @@ def __init__( 3, 14, ): - self.vlm_model = torch.compile(self.vlm_model) # type: ignore + self.vlm_model = torch.compile(self.vlm_model) else: self.vlm_model.eval() @@ -133,9 +141,10 @@ def __init__( artifacts_path ) + @override def process_images( self, - image_batch: Iterable[Union[Image, np.ndarray]], + image_batch: Iterable[Union[Image, NDArray[np.generic]]], prompt: Union[str, list[str]], ) -> Iterable[VlmPrediction]: from PIL import Image as PILImage @@ -211,7 +220,7 @@ def process_images( input_len = processor_inputs["input_ids"].shape[1] trimmed_sequences = generated_ids[:, input_len:] - decoded_texts: list[str] = self.processor.batch_decode( + decoded_texts = self.processor.batch_decode( trimmed_sequences, skip_special_tokens=True, clean_up_tokenization_spaces=False, @@ -226,6 +235,7 @@ def process_images( ) for text in decoded_texts: + assert isinstance(text, str) decoded_text = self.vlm_options.decode_response(text) yield VlmPrediction( text=decoded_text, diff --git a/docling/models/factories/base_factory.py b/docling/models/factories/base_factory.py index 109d283044..c1d662222f 100644 --- a/docling/models/factories/base_factory.py +++ b/docling/models/factories/base_factory.py @@ -28,7 +28,7 @@ class FactoryMeta(BaseModel): class BaseFactory(Generic[A], metaclass=ABCMeta): default_plugin_name = "docling" - def __init__(self, plugin_attr_name: str, plugin_name=default_plugin_name): + def __init__(self, plugin_attr_name: str, plugin_name: str = default_plugin_name): self.plugin_name = plugin_name self.plugin_attr_name = plugin_attr_name @@ -39,7 +39,7 @@ def __init__(self, plugin_attr_name: str, plugin_name=default_plugin_name): def registered_kind(self) -> list[str]: return [opt.kind for opt in self._classes.keys()] - def get_enum(self) -> enum.Enum: + def get_enum(self) -> type[enum.Enum]: return enum.Enum( self.plugin_attr_name + "_enum", names={kind: kind for kind in self.registered_kind}, @@ -52,7 +52,7 @@ def classes(self): return self._classes @property - def registered_meta(self): + def registered_meta(self) -> dict[Type[BaseOptions], FactoryMeta]: return self._meta def create_instance(self, options: BaseOptions, **kwargs) -> A: diff --git a/docling/models/inference_engines/common/hf_vision_base.py b/docling/models/inference_engines/common/hf_vision_base.py index 9cca646b2c..ee11d496c9 100644 --- a/docling/models/inference_engines/common/hf_vision_base.py +++ b/docling/models/inference_engines/common/hf_vision_base.py @@ -84,10 +84,16 @@ def _load_preprocessor(self, model_folder: Path) -> BaseImageProcessor: ) try: - from transformers import AutoImageProcessor + from transformers.image_processing_utils import BaseImageProcessor + from transformers.models.auto.image_processing_auto import ( + AutoImageProcessor, + ) _log.debug("Loading image processor from %s", model_folder) - return AutoImageProcessor.from_pretrained(str(model_folder)) + processor = AutoImageProcessor.from_pretrained(str(model_folder)) + if not isinstance(processor, BaseImageProcessor): + raise TypeError("Expected a HuggingFace image processor") + return processor except Exception as exc: raise RuntimeError( f"Failed to load image processor from {model_folder}: {exc}" @@ -96,7 +102,7 @@ def _load_preprocessor(self, model_folder: Path) -> BaseImageProcessor: def _load_label_mapping(self, model_folder: Path) -> Dict[int, str]: """Load label mapping from HuggingFace model config.""" try: - from transformers import AutoConfig + from transformers.models.auto.configuration_auto import AutoConfig config = AutoConfig.from_pretrained(str(model_folder)) return { diff --git a/docling/models/inference_engines/common/kserve_v2_client_base.py b/docling/models/inference_engines/common/kserve_v2_client_base.py index 69ba3fdde4..d02159d9bb 100644 --- a/docling/models/inference_engines/common/kserve_v2_client_base.py +++ b/docling/models/inference_engines/common/kserve_v2_client_base.py @@ -8,6 +8,7 @@ from typing import Any, Dict, Mapping, Optional, Protocol import numpy as np +from numpy.typing import NDArray from docling.models.inference_engines.common.kserve_v2_types import ( KserveV2ModelMetadataResponse, @@ -23,10 +24,10 @@ def get_model_metadata(self) -> KserveV2ModelMetadataResponse: def infer( self, *, - inputs: Mapping[str, np.ndarray], + inputs: Mapping[str, NDArray[np.generic]], output_names: list[str], request_parameters: Optional[Mapping[str, Any]] = None, - ) -> Dict[str, np.ndarray]: + ) -> Dict[str, NDArray[np.generic]]: """Execute model inference and return outputs keyed by tensor name.""" def close(self) -> None: diff --git a/docling/models/inference_engines/common/kserve_v2_grpc.py b/docling/models/inference_engines/common/kserve_v2_grpc.py index fefe50cc87..54e8a6f0b7 100644 --- a/docling/models/inference_engines/common/kserve_v2_grpc.py +++ b/docling/models/inference_engines/common/kserve_v2_grpc.py @@ -12,6 +12,7 @@ from urllib.parse import SplitResult, urlsplit import numpy as np +from numpy.typing import NDArray try: import grpc # type: ignore[import-untyped] @@ -111,7 +112,7 @@ def _set_request_parameter( ) -def _encode_contents(tensor: np.ndarray, contents: Any) -> None: +def _encode_contents(tensor: NDArray[np.generic], contents: Any) -> None: """Populate an InferTensorContents message from a numpy array (non-binary path).""" flat = tensor.flatten() if tensor.dtype == np.float32: @@ -140,7 +141,7 @@ def _encode_contents(tensor: np.ndarray, contents: Any) -> None: def _decode_contents( contents: Any, np_dtype: np.dtype[Any], shape: tuple[int, ...] -) -> np.ndarray: +) -> NDArray[np.generic]: """Decode an InferTensorContents message to a numpy array (non-binary path).""" canonical_dtype = np.dtype(np_dtype) @@ -198,6 +199,9 @@ def __post_init__(self) -> None: "Install with: pip install 'docling[remote-serving]'" ) + self._grpc = grpc + self._service_pb2 = service_pb2 + endpoint = _resolve_grpc_endpoint( base_url=self.base_url, ) @@ -238,7 +242,8 @@ def close(self) -> None: self._channel.close() def get_model_metadata(self) -> KserveV2ModelMetadataResponse: - request = service_pb2.ModelMetadataRequest(name=self.model_name) + # Protobuf builds these message classes at import time from its descriptor. + request = self._service_pb2.ModelMetadataRequest(name=self.model_name) # ty: ignore[unresolved-attribute] if self.model_version: request.version = self.model_version @@ -248,7 +253,7 @@ def get_model_metadata(self) -> KserveV2ModelMetadataResponse: timeout=self.timeout, metadata=self._grpc_metadata, ) - except grpc.RpcError as exc: + except self._grpc.RpcError as exc: raise RuntimeError( f"gRPC metadata call failed for model {self.model_name}: {exc}" ) from exc @@ -281,17 +286,23 @@ def get_model_metadata(self) -> KserveV2ModelMetadataResponse: def infer( self, *, - inputs: Mapping[str, np.ndarray], + inputs: Mapping[str, NDArray[np.generic]], output_names: list[str], request_parameters: Mapping[str, Any] | None = None, - ) -> Dict[str, np.ndarray]: - _batch_size = next(iter(inputs.values())).shape[0] if inputs else 0 + ) -> Dict[str, NDArray[np.generic]]: + _first_tensor = next(iter(inputs.values()), None) + _batch_size = ( + _first_tensor.shape[0] + if _first_tensor is not None and _first_tensor.ndim > 0 + else 0 + ) - if _log.isEnabledFor(logging.DEBUG): - _t_ser_start = time.time() - _t_ser_mono = time.monotonic() + debug_logging = _log.isEnabledFor(logging.DEBUG) + _t_ser_start = time.time() if debug_logging else None + _t_ser_mono = time.monotonic() if debug_logging else None - request = service_pb2.ModelInferRequest(model_name=self.model_name) + # Protobuf builds these message classes at import time from its descriptor. + request = self._service_pb2.ModelInferRequest(model_name=self.model_name) # ty: ignore[unresolved-attribute] if self.model_version: request.model_version = self.model_version @@ -329,7 +340,7 @@ def infer( if self.use_binary_data: output_tensor.parameters["binary_data"].bool_param = True - if _log.isEnabledFor(logging.DEBUG): + if _t_ser_start is not None and _t_ser_mono is not None: _log.debug( "PIPELINE_PROFILING KServe gRPC infer serialization: batch_size=%d start=%.3f end=%.3f duration=%.3fs", _batch_size, @@ -337,8 +348,8 @@ def infer( time.time(), time.monotonic() - _t_ser_mono, ) - _t_grpc_start = time.time() - _t_grpc_mono = time.monotonic() + _t_grpc_start = time.time() if debug_logging else None + _t_grpc_mono = time.monotonic() if debug_logging else None try: response = self._stub.ModelInfer( @@ -346,12 +357,12 @@ def infer( timeout=self.timeout, metadata=self._grpc_metadata, ) - except grpc.RpcError as exc: + except self._grpc.RpcError as exc: raise RuntimeError( f"gRPC infer call failed for model {self.model_name}: {exc}" ) from exc - if _log.isEnabledFor(logging.DEBUG): + if _t_grpc_start is not None and _t_grpc_mono is not None: _log.debug( "PIPELINE_PROFILING KServe gRPC infer round-trip: batch_size=%d start=%.3f end=%.3f duration=%.3fs", _batch_size, @@ -359,10 +370,10 @@ def infer( time.time(), time.monotonic() - _t_grpc_mono, ) - _t_deser_start = time.time() - _t_deser_mono = time.monotonic() + _t_deser_start = time.time() if debug_logging else None + _t_deser_mono = time.monotonic() if debug_logging else None - decoded_outputs: Dict[str, np.ndarray] = {} + decoded_outputs: Dict[str, NDArray[np.generic]] = {} if self.use_binary_data: if len(response.raw_output_contents) != len(response.outputs): @@ -403,7 +414,7 @@ def infer( output_tensor.contents, np_dtype, shape ) - if _log.isEnabledFor(logging.DEBUG): + if _t_deser_start is not None and _t_deser_mono is not None: _log.debug( "PIPELINE_PROFILING KServe gRPC infer deserialization: batch_size=%d start=%.3f end=%.3f duration=%.3fs", _batch_size, diff --git a/docling/models/inference_engines/common/kserve_v2_http.py b/docling/models/inference_engines/common/kserve_v2_http.py index a91c8eed00..7f2d714632 100644 --- a/docling/models/inference_engines/common/kserve_v2_http.py +++ b/docling/models/inference_engines/common/kserve_v2_http.py @@ -18,6 +18,7 @@ import numpy as np import requests +from numpy.typing import NDArray from pydantic import BaseModel from docling.models.inference_engines.common.kserve_v2_types import ( @@ -34,7 +35,7 @@ _INFERENCE_HEADER_CONTENT_LENGTH = "Inference-Header-Content-Length" -def _tensor_kserve_dtype(tensor: np.ndarray) -> str: +def _tensor_kserve_dtype(tensor: NDArray[np.generic]) -> str: kserve_dtype = NUMPY_KSERVE_V2_DATATYPES.get(tensor.dtype) if kserve_dtype is None: raise ValueError( @@ -44,7 +45,7 @@ def _tensor_kserve_dtype(tensor: np.ndarray) -> str: return kserve_dtype -def _encode_input_tensor(name: str, tensor: np.ndarray) -> Dict[str, Any]: +def _encode_input_tensor(name: str, tensor: NDArray[np.generic]) -> Dict[str, Any]: kserve_dtype = _tensor_kserve_dtype(tensor) return { @@ -56,7 +57,7 @@ def _encode_input_tensor(name: str, tensor: np.ndarray) -> Dict[str, Any]: def _encode_binary_input_tensor( - name: str, tensor: np.ndarray + name: str, tensor: NDArray[np.generic] ) -> tuple[Dict[str, Any], bytes]: kserve_dtype = _tensor_kserve_dtype(tensor) if kserve_dtype == "BYTES": @@ -91,7 +92,7 @@ class KserveV2InferResponse(BaseModel): outputs: List[KserveV2OutputTensor] -def _decode_output_tensor(raw_output: KserveV2OutputTensor) -> np.ndarray: +def _decode_output_tensor(raw_output: KserveV2OutputTensor) -> NDArray[np.generic]: shape = tuple(int(dim) for dim in raw_output.shape) np_dtype = KSERVE_V2_NUMPY_DATATYPES.get(raw_output.datatype) if np_dtype is None: @@ -111,7 +112,7 @@ def _decode_output_tensor(raw_output: KserveV2OutputTensor) -> np.ndarray: def _decode_binary_output_tensor( raw_output: KserveV2OutputTensor, raw_payload: bytes -) -> np.ndarray: +) -> NDArray[np.generic]: np_dtype = KSERVE_V2_NUMPY_DATATYPES.get(raw_output.datatype) if np_dtype is None: raise RuntimeError( @@ -142,7 +143,7 @@ def _parse_binary_data_size(parameters: Mapping[str, Any] | None) -> int | None: def _build_binary_request( *, - inputs: Mapping[str, np.ndarray], + inputs: Mapping[str, NDArray[np.generic]], output_names: list[str], request_parameters: Optional[Mapping[str, Any]], ) -> tuple[Dict[str, str], bytes]: @@ -253,8 +254,6 @@ def _execute_http_request( response = requests.post( url, headers=request_headers, timeout=self.timeout, **kwargs ) - response.raise_for_status() - return response except requests.exceptions.Timeout as exc: raise requests.exceptions.Timeout( f"Timeout during {method} request to {url}" @@ -263,10 +262,13 @@ def _execute_http_request( raise requests.exceptions.ConnectionError( f"Failed to connect to {url}" ) from exc + try: + response.raise_for_status() except requests.exceptions.HTTPError as exc: raise requests.exceptions.HTTPError( f"HTTP error {response.status_code} from {url}: {response.text}" ) from exc + return response @property def model_metadata_url(self) -> str: @@ -318,10 +320,10 @@ def get_model_metadata(self) -> KserveV2ModelMetadataResponse: def infer( self, *, - inputs: Mapping[str, np.ndarray], + inputs: Mapping[str, NDArray[np.generic]], output_names: list[str], request_parameters: Optional[Mapping[str, Any]] = None, - ) -> Dict[str, np.ndarray]: + ) -> Dict[str, NDArray[np.generic]]: """Execute inference request against KServe v2 endpoint. Args: @@ -338,10 +340,15 @@ def infer( requests.exceptions.HTTPError: If server returns error status RuntimeError: If response format is invalid """ - if _log.isEnabledFor(logging.DEBUG): - _batch_size = next(iter(inputs.values())).shape[0] if inputs else 0 - _t_ser_start = time.time() - _t_ser_mono = time.monotonic() + debug_logging = _log.isEnabledFor(logging.DEBUG) + _first_tensor = next(iter(inputs.values()), None) + _batch_size = ( + _first_tensor.shape[0] + if _first_tensor is not None and _first_tensor.ndim > 0 + else 0 + ) + _t_ser_start = time.time() if debug_logging else None + _t_ser_mono = time.monotonic() if debug_logging else None request_kwargs: Dict[str, Any] if self.use_binary_data: binary_headers, request_body = _build_binary_request( @@ -371,7 +378,7 @@ def infer( request_kwargs = {"json": payload} - if _log.isEnabledFor(logging.DEBUG): + if _t_ser_start is not None and _t_ser_mono is not None: _log.debug( "PIPELINE_PROFILING KServe infer serialization: batch_size=%d start=%.3f end=%.3f duration=%.3fs", _batch_size, @@ -379,12 +386,12 @@ def infer( time.time(), time.monotonic() - _t_ser_mono, ) - _t_http_start = time.time() - _t_http_mono = time.monotonic() + _t_http_start = time.time() if debug_logging else None + _t_http_mono = time.monotonic() if debug_logging else None response = self._execute_http_request( self.infer_url, method="POST", **request_kwargs ) - if _log.isEnabledFor(logging.DEBUG): + if _t_http_start is not None and _t_http_mono is not None: _log.debug( "PIPELINE_PROFILING KServe infer http round-trip: batch_size=%d start=%.3f end=%.3f duration=%.3fs", _batch_size, @@ -392,8 +399,8 @@ def infer( time.time(), time.monotonic() - _t_http_mono, ) - _t_deser_start = time.time() - _t_deser_mono = time.monotonic() + _t_deser_start = time.time() if debug_logging else None + _t_deser_mono = time.monotonic() if debug_logging else None try: body = ( @@ -406,7 +413,7 @@ def infer( f"Invalid inference response from {self.infer_url}: {exc}" ) from exc - decoded_outputs: Dict[str, np.ndarray] = {} + decoded_outputs: Dict[str, NDArray[np.generic]] = {} header_len_text = response.headers.get(_INFERENCE_HEADER_CONTENT_LENGTH) raw_body = b"" if self.use_binary_data and header_len_text is not None: @@ -436,7 +443,7 @@ def infer( f"not consumed: {len(raw_body) - raw_offset} bytes" ) - if _log.isEnabledFor(logging.DEBUG): + if _t_deser_start is not None and _t_deser_mono is not None: _log.debug( "PIPELINE_PROFILING KServe infer deserialization: batch_size=%d start=%.3f end=%.3f duration=%.3fs", _batch_size, diff --git a/docling/models/inference_engines/common/kserve_v2_utils.py b/docling/models/inference_engines/common/kserve_v2_utils.py index b42df350ef..c2c280dfe7 100644 --- a/docling/models/inference_engines/common/kserve_v2_utils.py +++ b/docling/models/inference_engines/common/kserve_v2_utils.py @@ -8,6 +8,7 @@ from typing import Any import numpy as np +from numpy.typing import NDArray def encode_bytes_element(value: Any) -> bytes: @@ -21,7 +22,7 @@ def encode_bytes_element(value: Any) -> bytes: return str(value).encode("utf-8") -def encode_bytes_tensor(tensor: np.ndarray) -> bytes: +def encode_bytes_tensor(tensor: NDArray[np.generic]) -> bytes: """Encode a BYTES tensor as a length-prefixed byte stream.""" chunks: list[bytes] = [] for value in tensor.reshape(-1): @@ -31,7 +32,9 @@ def encode_bytes_tensor(tensor: np.ndarray) -> bytes: return b"".join(chunks) -def decode_bytes_tensor(raw_output: bytes, shape: tuple[int, ...]) -> np.ndarray: +def decode_bytes_tensor( + raw_output: bytes, shape: tuple[int, ...] +) -> NDArray[np.generic]: """Decode a length-prefixed BYTES payload to a numpy object array.""" strings, offset = [], 0 for _ in range(int(np.prod(shape))): diff --git a/docling/models/inference_engines/image_classification/api_kserve_v2_engine.py b/docling/models/inference_engines/image_classification/api_kserve_v2_engine.py index 26767b39d7..3c6a5e34be 100644 --- a/docling/models/inference_engines/image_classification/api_kserve_v2_engine.py +++ b/docling/models/inference_engines/image_classification/api_kserve_v2_engine.py @@ -10,6 +10,7 @@ from typing import TYPE_CHECKING, List, Optional, Union import numpy as np +from typing_extensions import override from docling.datamodel.accelerator_options import AcceleratorOptions from docling.datamodel.image_classification_engine_options import ( @@ -90,6 +91,7 @@ def _resolve_tensor_names(self) -> tuple[str, str]: output_name = metadata.outputs[0].name return input_name, output_name + @override def initialize(self) -> None: """Initialize preprocessor/labels and prepare remote client.""" _log.info("Initializing KServe v2 image-classification engine") @@ -143,6 +145,7 @@ def initialize(self) -> None: self._output_name, ) + @override def predict_batch( self, input_batch: List[ImageClassificationEngineInput] ) -> List[ImageClassificationEngineOutput]: diff --git a/docling/models/inference_engines/image_classification/base.py b/docling/models/inference_engines/image_classification/base.py index 929766acf2..0ff9cda64b 100644 --- a/docling/models/inference_engines/image_classification/base.py +++ b/docling/models/inference_engines/image_classification/base.py @@ -24,6 +24,7 @@ from PIL.Image import Image from pydantic import BaseModel, ConfigDict, Field, field_validator from pydantic_core import PydanticUndefined +from typing_extensions import override if TYPE_CHECKING: from docling.datamodel.stage_model_specs import EngineModelConfig @@ -64,6 +65,7 @@ class BaseImageClassificationEngineOptions(BaseModel): ] = {} @classmethod + @override def __pydantic_init_subclass__(cls, **kwargs): super().__pydantic_init_subclass__(**kwargs) diff --git a/docling/models/inference_engines/image_classification/hf_base.py b/docling/models/inference_engines/image_classification/hf_base.py index 50fee6b0d1..6d752855b1 100644 --- a/docling/models/inference_engines/image_classification/hf_base.py +++ b/docling/models/inference_engines/image_classification/hf_base.py @@ -9,6 +9,7 @@ from typing import TYPE_CHECKING, Any, Iterable, Optional, Union import numpy as np +from numpy.typing import NDArray from docling.datamodel.accelerator_options import AcceleratorOptions from docling.models.inference_engines.common import HfVisionModelMixin @@ -73,7 +74,7 @@ def _build_batch_outputs_from_probabilities( self, *, input_batch: list[ImageClassificationEngineInput], - probs_batch: np.ndarray, + probs_batch: NDArray[np.float16 | np.float32 | np.float64], ) -> list[ImageClassificationEngineOutput]: """Build outputs from a probability matrix of shape [batch, num_classes].""" batch_outputs: list[ImageClassificationEngineOutput] = [] @@ -100,8 +101,10 @@ def _build_batch_outputs_from_probabilities( return batch_outputs @staticmethod - def _softmax(logits: np.ndarray) -> np.ndarray: + def _softmax(logits: NDArray[np.float32]) -> NDArray[np.float32]: """Compute softmax probabilities over class dimension.""" shifted = logits - np.max(logits, axis=1, keepdims=True) exp_values = np.exp(shifted) - return exp_values / np.sum(exp_values, axis=1, keepdims=True) + return np.asarray( + exp_values / np.sum(exp_values, axis=1, keepdims=True), dtype=np.float32 + ) diff --git a/docling/models/inference_engines/image_classification/onnxruntime_engine.py b/docling/models/inference_engines/image_classification/onnxruntime_engine.py index c9cd9f7bd2..f21c2f86fe 100644 --- a/docling/models/inference_engines/image_classification/onnxruntime_engine.py +++ b/docling/models/inference_engines/image_classification/onnxruntime_engine.py @@ -10,6 +10,7 @@ from typing import TYPE_CHECKING, List, Optional, Union import numpy as np +from typing_extensions import override if TYPE_CHECKING: import onnxruntime as ort @@ -89,18 +90,29 @@ def _resolve_input_name(self, session: ort.InferenceSession) -> str: input_nodes = session.get_inputs() if not input_nodes: raise RuntimeError("ONNX model exposes no inputs") - return input_nodes[0].name + name = input_nodes[0].name + if not isinstance(name, str): + raise TypeError("ONNX input name must be a string") + return name def _resolve_output_name(self, session: ort.InferenceSession) -> str: """Resolve ONNX output name from the loaded model graph.""" output_nodes = session.get_outputs() if not output_nodes: raise RuntimeError("ONNX model exposes no outputs") - return output_nodes[0].name + name = output_nodes[0].name + if not isinstance(name, str): + raise TypeError("ONNX output name must be a string") + return name + @override def initialize(self) -> None: """Initialize ONNX session and preprocessor.""" import onnxruntime as ort + from onnxruntime.capi._pybind_state import ( + GraphOptimizationLevel, + SessionOptions, + ) _log.info("Initializing ONNX Runtime image-classification engine") @@ -110,9 +122,9 @@ def initialize(self) -> None: self._processor = self._load_preprocessor(model_folder) self._id_to_label = self._load_label_mapping(model_folder) - sess_options = ort.SessionOptions() + sess_options = SessionOptions() sess_options.intra_op_num_threads = self._accelerator_options.num_threads - sess_options.graph_optimization_level = ort.GraphOptimizationLevel( + sess_options.graph_optimization_level = GraphOptimizationLevel( self.options.graph_optimization_level ) providers = self._resolve_providers() @@ -154,6 +166,7 @@ def _resolve_providers(self) -> List[str]: ) return ["CPUExecutionProvider"] + @override def predict_batch( self, input_batch: List[ImageClassificationEngineInput] ) -> List[ImageClassificationEngineOutput]: diff --git a/docling/models/inference_engines/image_classification/transformers_engine.py b/docling/models/inference_engines/image_classification/transformers_engine.py index f473a52b31..3df5e29f9a 100644 --- a/docling/models/inference_engines/image_classification/transformers_engine.py +++ b/docling/models/inference_engines/image_classification/transformers_engine.py @@ -12,10 +12,11 @@ from typing import TYPE_CHECKING, List, Optional, Union from packaging import version +from typing_extensions import override if TYPE_CHECKING: import torch - from transformers import AutoModelForImageClassification + from transformers.modeling_utils import PreTrainedModel from docling.datamodel.accelerator_options import AcceleratorDevice, AcceleratorOptions from docling.datamodel.image_classification_engine_options import ( @@ -55,7 +56,7 @@ def __init__( artifacts_path=artifacts_path, ) self.options: TransformersImageClassificationEngineOptions = options - self._model: Optional[AutoModelForImageClassification] = None + self._model: Optional[PreTrainedModel] = None self._device: Optional[torch.device] = None def _resolve_device(self) -> torch.device: @@ -100,10 +101,13 @@ def _resolve_torch_dtype(self) -> Optional[torch.dtype]: ) return dtype + @override def initialize(self) -> None: """Initialize PyTorch model and preprocessor.""" import torch - from transformers import AutoModelForImageClassification + from transformers.models.auto.modeling_auto import ( + AutoModelForImageClassification, + ) _log.info("Initializing Transformers image-classification engine") @@ -171,6 +175,7 @@ def initialize(self) -> None: self._model.dtype, # type: ignore[union-attr] ) + @override def predict_batch( self, input_batch: List[ImageClassificationEngineInput] ) -> List[ImageClassificationEngineOutput]: diff --git a/docling/models/inference_engines/object_detection/api_kserve_v2_engine.py b/docling/models/inference_engines/object_detection/api_kserve_v2_engine.py index 4435a8d009..3d1e618f46 100644 --- a/docling/models/inference_engines/object_detection/api_kserve_v2_engine.py +++ b/docling/models/inference_engines/object_detection/api_kserve_v2_engine.py @@ -11,6 +11,7 @@ from typing import TYPE_CHECKING, List, Optional, Union import numpy as np +from typing_extensions import override from docling.datamodel.accelerator_options import AcceleratorOptions from docling.datamodel.kserve_transport_utils import resolve_kserve_transport_base_url @@ -104,6 +105,7 @@ def _resolve_tensor_names(self) -> tuple[str, str, str, str, str]: output_scores_name, ) + @override def initialize(self) -> None: """Initialize preprocessor/labels and prepare remote client.""" _log.info("Initializing KServe v2 object-detection engine") @@ -166,6 +168,7 @@ def initialize(self) -> None: self._output_scores_name, ) + @override def predict_batch( self, input_batch: List[ObjectDetectionEngineInput] ) -> List[ObjectDetectionEngineOutput]: @@ -184,9 +187,9 @@ def predict_batch( assert self._output_boxes_name is not None assert self._output_scores_name is not None - if _log.isEnabledFor(logging.DEBUG): - _t_preproc_start = time.time() - _t_preproc_mono = time.monotonic() + debug_logging = _log.isEnabledFor(logging.DEBUG) + _t_preproc_start = time.time() if debug_logging else None + _t_preproc_mono = time.monotonic() if debug_logging else None images = [item.image.convert("RGB") for item in input_batch] processed_inputs = self._processor(images=images, return_tensors="np") @@ -195,7 +198,7 @@ def predict_batch( [[image.width, image.height] for image in images], dtype=np.int64, ) - if _log.isEnabledFor(logging.DEBUG): + if _t_preproc_start is not None and _t_preproc_mono is not None: _log.debug( "PIPELINE_PROFILING KServe object-detection HF preprocessor: batch_size=%d start=%.3f end=%.3f duration=%.3fs", len(input_batch), diff --git a/docling/models/inference_engines/object_detection/base.py b/docling/models/inference_engines/object_detection/base.py index 9771f5881c..2fc68a6f05 100644 --- a/docling/models/inference_engines/object_detection/base.py +++ b/docling/models/inference_engines/object_detection/base.py @@ -24,6 +24,7 @@ from PIL.Image import Image from pydantic import BaseModel, ConfigDict, Field, field_validator from pydantic_core import PydanticUndefined +from typing_extensions import override if TYPE_CHECKING: from docling.datamodel.stage_model_specs import EngineModelConfig @@ -63,6 +64,7 @@ class BaseObjectDetectionEngineOptions(BaseModel): ] = {} @classmethod + @override def __pydantic_init_subclass__(cls, **kwargs): super().__pydantic_init_subclass__(**kwargs) diff --git a/docling/models/inference_engines/object_detection/onnxruntime_engine.py b/docling/models/inference_engines/object_detection/onnxruntime_engine.py index db1ae65df6..993e65e0aa 100644 --- a/docling/models/inference_engines/object_detection/onnxruntime_engine.py +++ b/docling/models/inference_engines/object_detection/onnxruntime_engine.py @@ -10,6 +10,7 @@ from typing import TYPE_CHECKING, List, Optional, Union import numpy as np +from typing_extensions import override if TYPE_CHECKING: import onnxruntime as ort @@ -98,9 +99,14 @@ def _resolve_model_filename(self) -> str: filename = extra_filename return filename + @override def initialize(self) -> None: """Initialize ONNX session and preprocessor.""" import onnxruntime as ort + from onnxruntime.capi._pybind_state import ( + GraphOptimizationLevel, + SessionOptions, + ) _log.info("Initializing ONNX Runtime object-detection engine") @@ -111,16 +117,16 @@ def initialize(self) -> None: # Load preprocessor (source of truth for preprocessing) self._processor = self._load_preprocessor(model_folder) - _log.debug(f"Loaded preprocessor with size: {self._processor.size}") # type: ignore[attr-defined] + _log.debug("Loaded preprocessor: %s", type(self._processor).__name__) # Load label mapping from config self._id_to_label = self._load_label_mapping(model_folder) _log.debug(f"Loaded label mapping with {len(self._id_to_label)} labels") # Create ONNX session - sess_options = ort.SessionOptions() + sess_options = SessionOptions() sess_options.intra_op_num_threads = self._accelerator_options.num_threads - sess_options.graph_optimization_level = ort.GraphOptimizationLevel( + sess_options.graph_optimization_level = GraphOptimizationLevel( self.options.graph_optimization_level ) providers = self._resolve_providers() @@ -157,6 +163,7 @@ def _resolve_providers(self) -> List[str]: ) return ["CPUExecutionProvider"] + @override def predict_batch( self, input_batch: List[ObjectDetectionEngineInput] ) -> List[ObjectDetectionEngineOutput]: @@ -197,7 +204,9 @@ def predict_batch( "[labels, boxes, scores]" ) - labels_batch, boxes_batch, scores_batch = output_tensors[:3] + labels_batch, boxes_batch, scores_batch = ( + np.asarray(tensor) for tensor in output_tensors[:3] + ) batch_outputs: List[ObjectDetectionEngineOutput] = [] for idx, input_item in enumerate(input_batch): diff --git a/docling/models/inference_engines/object_detection/transformers_engine.py b/docling/models/inference_engines/object_detection/transformers_engine.py index c5135b4e3f..83a67560fe 100644 --- a/docling/models/inference_engines/object_detection/transformers_engine.py +++ b/docling/models/inference_engines/object_detection/transformers_engine.py @@ -12,10 +12,11 @@ from typing import TYPE_CHECKING, List, Optional, Union from packaging import version +from typing_extensions import override if TYPE_CHECKING: import torch - from transformers import AutoModelForObjectDetection + from transformers.modeling_utils import PreTrainedModel from docling.datamodel.accelerator_options import AcceleratorDevice, AcceleratorOptions from docling.datamodel.object_detection_engine_options import ( @@ -66,7 +67,7 @@ def __init__( artifacts_path=artifacts_path, ) self.options: TransformersObjectDetectionEngineOptions = options - self._model: Optional[AutoModelForObjectDetection] = None + self._model: Optional[PreTrainedModel] = None self._device: Optional[torch.device] = None def _resolve_device(self) -> torch.device: @@ -116,10 +117,11 @@ def _resolve_torch_dtype(self) -> Optional[torch.dtype]: ) return dtype + @override def initialize(self) -> None: """Initialize PyTorch model and preprocessor.""" import torch - from transformers import AutoModelForObjectDetection + from transformers.models.auto.modeling_auto import AutoModelForObjectDetection _log.info("Initializing Transformers object-detection engine") @@ -141,7 +143,7 @@ def initialize(self) -> None: # Load preprocessor (source of truth for preprocessing) self._processor = self._load_preprocessor(model_folder) - _log.debug(f"Loaded preprocessor with size: {self._processor.size}") # type: ignore[attr-defined] + _log.debug("Loaded preprocessor: %s", type(self._processor).__name__) # Load label mapping from config self._id_to_label = self._load_label_mapping(model_folder) @@ -199,6 +201,7 @@ def initialize(self) -> None: f"Transformers engine ready (device={self._device}, dtype={self._model.dtype})" # type: ignore[union-attr] ) + @override def predict_batch( self, input_batch: List[ObjectDetectionEngineInput] ) -> List[ObjectDetectionEngineOutput]: @@ -231,7 +234,9 @@ def predict_batch( outputs = self._model(**inputs) # type: ignore[operator] # Post-process using HuggingFace processor - results = self._processor.post_process_object_detection( # type: ignore[attr-defined] + # Detection-specific processors supply this method; BaseImageProcessor + # does not declare the post-processing methods of its model subclasses. + results = self._processor.post_process_object_detection( # ty: ignore[unresolved-attribute] outputs, target_sizes=target_sizes, # type: ignore[arg-type] threshold=self.options.score_threshold, diff --git a/docling/models/inference_engines/vlm/_utils.py b/docling/models/inference_engines/vlm/_utils.py index 31eab2d292..798b9df5a9 100644 --- a/docling/models/inference_engines/vlm/_utils.py +++ b/docling/models/inference_engines/vlm/_utils.py @@ -13,6 +13,7 @@ from typing import Any, Callable, Dict, List, Optional, Union import numpy as np +from numpy.typing import NDArray from packaging import version from PIL import Image @@ -103,7 +104,9 @@ def engine_version_satisfied( return False -def normalize_image_to_pil(image: Union[Image.Image, np.ndarray]) -> Image.Image: +def normalize_image_to_pil( + image: Union[Image.Image, NDArray[np.generic]], +) -> Image.Image: """Convert any image format to RGB PIL Image. Args: @@ -134,7 +137,7 @@ def normalize_image_to_pil(image: Union[Image.Image, np.ndarray]) -> Image.Image def preprocess_image_batch( - images: List[Union[Image.Image, np.ndarray]], + images: List[Union[Image.Image, NDArray[np.generic]]], ) -> List[Image.Image]: """Preprocess a batch of images to RGB PIL Images. @@ -273,7 +276,12 @@ def format_prompt_for_vlm( ], } ] - return processor.apply_chat_template(messages, add_generation_prompt=True) + formatted_prompt = processor.apply_chat_template( + messages, add_generation_prompt=True, tokenize=False + ) + if not isinstance(formatted_prompt, str): + raise TypeError("Chat template must produce text when tokenize=False") + return formatted_prompt else: raise ValueError( f"Unknown prompt style: {prompt_style}. " diff --git a/docling/models/inference_engines/vlm/api_openai_compatible_engine.py b/docling/models/inference_engines/vlm/api_openai_compatible_engine.py index 29170930f0..7ecb8b1df9 100644 --- a/docling/models/inference_engines/vlm/api_openai_compatible_engine.py +++ b/docling/models/inference_engines/vlm/api_openai_compatible_engine.py @@ -10,6 +10,7 @@ from typing import TYPE_CHECKING, List, Optional from PIL.Image import Image +from typing_extensions import override from docling.datamodel.vlm_engine_options import ApiVlmEngineOptions from docling.exceptions import OperationNotAllowed @@ -73,14 +74,15 @@ def __init__( # correct priority order can be applied in predict_batch: # model_spec defaults < request-level generation settings < user params if model_config and "api_params" in model_config.extra_config: - self.model_api_params: dict = model_config.extra_config["api_params"].copy() + self.model_api_params = model_config.extra_config["api_params"].copy() else: self.model_api_params = {} # User-supplied params always win; if provided, model-spec defaults are # not mixed in (prevents conflicts for vendor-specific keys like model_id). - self.user_params: dict = self.options.params.copy() + self.user_params = self.options.params.copy() + @override def initialize(self) -> None: """Initialize the API engine. @@ -100,6 +102,7 @@ def initialize(self) -> None: self._initialized = True _log.info("API runtime initialized") + @override def predict_batch(self, input_batch: List[VlmEngineInput]) -> List[VlmEngineOutput]: """Run inference on a batch of inputs using concurrent API requests. @@ -154,6 +157,14 @@ def _process_single_input(input_data: VlmEngineInput) -> VlmEngineOutput: custom_stoppers = extract_generation_stoppers( input_data.extra_generation_config ) + usage_response_key = api_params.pop("usage_response_key", "usage") + token_extract_key = api_params.pop("token_extract_key", None) + if usage_response_key is not None and not isinstance( + usage_response_key, str + ): + raise ValueError("usage_response_key must be a string or None") + if token_extract_key is not None and not isinstance(token_extract_key, str): + raise ValueError("token_extract_key must be a string or None") request_start_time = time.time() stop_reason = "unspecified" error = None @@ -167,6 +178,8 @@ def _process_single_input(input_data: VlmEngineInput) -> VlmEngineOutput: headers=self.options.headers, generation_stoppers=custom_stoppers, timeout=self.options.timeout, + usage_response_key=usage_response_key, + token_extract_key=token_extract_key, **api_params, ) generated_text = api_response.text @@ -185,6 +198,8 @@ def _process_single_input(input_data: VlmEngineInput) -> VlmEngineOutput: prompt=input_data.prompt, headers=self.options.headers, timeout=self.options.timeout, + usage_response_key=usage_response_key, + token_extract_key=token_extract_key, **api_params, ) generated_text = api_response.text @@ -236,6 +251,7 @@ def _process_single_input(input_data: VlmEngineInput) -> VlmEngineOutput: return outputs + @override def cleanup(self) -> None: """Clean up API runtime resources. diff --git a/docling/models/inference_engines/vlm/auto_inline_engine.py b/docling/models/inference_engines/vlm/auto_inline_engine.py index caafdb6a26..423bc6b696 100644 --- a/docling/models/inference_engines/vlm/auto_inline_engine.py +++ b/docling/models/inference_engines/vlm/auto_inline_engine.py @@ -10,6 +10,8 @@ from pathlib import Path from typing import TYPE_CHECKING, List, Union +from typing_extensions import override + from docling.datamodel.accelerator_options import AcceleratorDevice, AcceleratorOptions from docling.datamodel.vlm_engine_options import ( AutoInlineVlmEngineOptions, @@ -172,6 +174,7 @@ def _engine_version_is_sufficient(self, engine_type: VlmEngineType) -> bool: ) return engine_version_satisfied(engine_type, required) + @override def initialize(self) -> None: """Initialize by selecting and creating the actual engine.""" if self._initialized: @@ -245,6 +248,7 @@ def initialize(self) -> None: f"Auto-inline engine initialized with {self.selected_engine_type.value}" ) + @override def predict_batch(self, input_batch: List[VlmEngineInput]) -> List[VlmEngineOutput]: """Run inference on a batch of inputs using the selected engine. @@ -262,6 +266,7 @@ def predict_batch(self, input_batch: List[VlmEngineInput]) -> List[VlmEngineOutp # Delegate to the actual engine's batch implementation return self.actual_engine.predict_batch(input_batch) + @override def cleanup(self) -> None: """Clean up the actual engine resources.""" if self.actual_engine is not None: diff --git a/docling/models/inference_engines/vlm/base.py b/docling/models/inference_engines/vlm/base.py index 68a66ee297..af1e923950 100644 --- a/docling/models/inference_engines/vlm/base.py +++ b/docling/models/inference_engines/vlm/base.py @@ -22,6 +22,7 @@ from PIL.Image import Image from pydantic import BaseModel, ConfigDict, Field, SerializeAsAny, field_validator from pydantic_core import PydanticUndefined +from typing_extensions import override if TYPE_CHECKING: from docling.datamodel.stage_model_specs import EngineModelConfig @@ -81,6 +82,7 @@ class BaseVlmEngineOptions(BaseModel): _registry: ClassVar[Dict[VlmEngineType, Type["BaseVlmEngineOptions"]]] = {} @classmethod + @override def __pydantic_init_subclass__(cls, **kwargs): super().__pydantic_init_subclass__(**kwargs) @@ -193,6 +195,8 @@ class BaseVlmEngine(ABC): These model specs are provided at construction time for eager initialization. """ + selected_engine_type: VlmEngineType | None = None + def __init__( self, options: BaseVlmEngineOptions, diff --git a/docling/models/inference_engines/vlm/factory.py b/docling/models/inference_engines/vlm/factory.py index ea0dc71be3..7b2d5a458a 100644 --- a/docling/models/inference_engines/vlm/factory.py +++ b/docling/models/inference_engines/vlm/factory.py @@ -52,7 +52,7 @@ def create_vlm_engine( # Generate model_config from model_spec if provided model_config: Optional[EngineModelConfig] = None - if model_spec is not None and engine_type != VlmEngineType.AUTO_INLINE: + if engine_type != VlmEngineType.AUTO_INLINE: # AUTO_INLINE handles model_spec internally model_config = model_spec.get_engine_config(engine_type) diff --git a/docling/models/inference_engines/vlm/mlx_engine.py b/docling/models/inference_engines/vlm/mlx_engine.py index 9c80eb2ae3..d9630dc6b0 100644 --- a/docling/models/inference_engines/vlm/mlx_engine.py +++ b/docling/models/inference_engines/vlm/mlx_engine.py @@ -13,6 +13,7 @@ from typing import TYPE_CHECKING, Any, Callable, List, Union from PIL.Image import Image +from typing_extensions import override from docling.datamodel.vlm_engine_options import MlxVlmEngineOptions from docling.models.inference_engines.vlm._utils import ( @@ -50,6 +51,7 @@ class _ProcessorKwargsWarningFilter(logging.Filter): + @override def filter(self, record: logging.LogRecord) -> bool: return _PROCESSOR_KWARGS_WARNING not in record.getMessage() @@ -105,6 +107,7 @@ def __init__( if self.model_config is not None: self.initialize() + @override def initialize(self) -> None: """Initialize the MLX model and processor.""" if self._initialized: @@ -188,7 +191,8 @@ def _load_model_for_repo(self, repo_id: str, revision: str = "main") -> None: # Load the model start_time = time.monotonic() self.vlm_model, self.processor = load( - artifacts_path, strict=not tied_word_embeddings + artifacts_path, + strict=not tied_word_embeddings, ) if tied_word_embeddings: language_model = self.vlm_model.language_model @@ -201,6 +205,7 @@ def _load_model_for_repo(self, repo_id: str, revision: str = "main") -> None: f"Loaded MLX model {repo_id} (revision: {revision}) in {load_time:.2f} sec." ) + @override def predict_batch(self, input_batch: List[VlmEngineInput]) -> List[VlmEngineOutput]: """Run inference on a batch of inputs. @@ -350,6 +355,7 @@ def predict_batch(self, input_batch: List[VlmEngineInput]) -> List[VlmEngineOutp return outputs + @override def cleanup(self) -> None: """Clean up model resources.""" if self.vlm_model is not None: diff --git a/docling/models/inference_engines/vlm/transformers_engine.py b/docling/models/inference_engines/vlm/transformers_engine.py index 0227e3e57a..5376f825ac 100644 --- a/docling/models/inference_engines/vlm/transformers_engine.py +++ b/docling/models/inference_engines/vlm/transformers_engine.py @@ -15,9 +15,6 @@ from packaging import version from PIL.Image import Image from transformers import ( - AutoModel, - AutoModelForCausalLM, - AutoModelForImageTextToText, AutoProcessor, BitsAndBytesConfig, GenerationConfig, @@ -25,6 +22,12 @@ StoppingCriteriaList, StopStringCriteria, ) +from transformers.models.auto.modeling_auto import ( + AutoModel, + AutoModelForCausalLM, + AutoModelForImageTextToText, +) +from typing_extensions import override from docling.datamodel.accelerator_options import AcceleratorDevice, AcceleratorOptions from docling.datamodel.pipeline_options_vlm_model import ( @@ -128,6 +131,7 @@ def __init__( if self.model_config is not None: self.initialize() + @override def initialize(self) -> None: """Initialize the Transformers model and processor.""" if self._initialized: @@ -263,6 +267,8 @@ def download_wrapper(repo_id: str, revision: str) -> Path: or self.model_config.extra_config.get("torch_dtype") ) + assert self.device is not None + # Load model attn_implementation: Optional[str] = ( "flash_attention_2" @@ -335,6 +341,7 @@ def _get_tokenizer(self) -> Any: return None return getattr(self.processor, "tokenizer", None) or self.processor + @override def predict_batch(self, input_batch: List[VlmEngineInput]) -> List[VlmEngineOutput]: """Run inference on a batch of inputs efficiently. @@ -507,8 +514,13 @@ def predict_batch(self, input_batch: List[VlmEngineInput]) -> List[VlmEngineOutp first_input.max_new_tokens, ) start_time = time.time() + generate = self.vlm_model.generate + if not callable(generate): + raise TypeError("VLM model must provide a generate method") with torch.inference_mode(): - generated_ids = self.vlm_model.generate(**gen_kwargs) # type: ignore[union-attr,operator] + generated_ids = generate(**gen_kwargs) + if not isinstance(generated_ids, torch.Tensor): + raise TypeError("VLM generation must produce a token tensor") generation_time = time.time() - start_time # Decode @@ -576,6 +588,7 @@ def predict_batch(self, input_batch: List[VlmEngineInput]) -> List[VlmEngineOutp return outputs + @override def cleanup(self) -> None: """Clean up model resources.""" if self.vlm_model is not None: diff --git a/docling/models/inference_engines/vlm/vllm_engine.py b/docling/models/inference_engines/vlm/vllm_engine.py index 016db6ffff..428363dc4e 100644 --- a/docling/models/inference_engines/vlm/vllm_engine.py +++ b/docling/models/inference_engines/vlm/vllm_engine.py @@ -9,6 +9,8 @@ from pathlib import Path from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union +from typing_extensions import override + from docling.datamodel.accelerator_options import AcceleratorDevice, AcceleratorOptions from docling.datamodel.pipeline_options_vlm_model import TransformersPromptStyle from docling.datamodel.vlm_engine_options import VllmVlmEngineOptions @@ -121,6 +123,7 @@ def __init__( if self.model_config is not None: self.initialize() + @override def initialize(self) -> None: """Initialize the vLLM engine.""" if self._initialized: @@ -261,6 +264,7 @@ def download_wrapper(repo_id: str, revision: str) -> Path: self._initialized = True _log.info("vLLM runtime initialized") + @override def predict_batch(self, input_batch: List[VlmEngineInput]) -> List[VlmEngineOutput]: """Run inference on a batch of inputs using vLLM. @@ -376,6 +380,7 @@ def predict_batch(self, input_batch: List[VlmEngineInput]) -> List[VlmEngineOutp return results + @override def cleanup(self) -> None: """Clean up vLLM resources.""" if self.llm is not None: diff --git a/docling/models/picture_description_base_model.py b/docling/models/picture_description_base_model.py index 033f39aa91..366f79a57e 100644 --- a/docling/models/picture_description_base_model.py +++ b/docling/models/picture_description_base_model.py @@ -16,6 +16,7 @@ ) from docling_core.types.doc.document import PictureDescriptionData from PIL import Image +from typing_extensions import override from docling.datamodel.accelerator_options import AcceleratorOptions from docling.datamodel.base_models import ApiImageRequestResult @@ -57,6 +58,7 @@ def __init__( self.elements_batch_size = options.batch_size self.images_scale = options.scale + @override def is_processable(self, doc: DoclingDocument, element: NodeItem) -> bool: return self.enabled and isinstance(element, PictureItem) @@ -65,6 +67,7 @@ def _annotate_images( ) -> Iterable[str | ApiImageRequestResult]: raise NotImplementedError + @override def __call__( self, doc: DoclingDocument, @@ -131,6 +134,7 @@ def __call__( @classmethod @abstractmethod + @override def get_options_type(cls) -> Type[PictureDescriptionBaseOptions]: pass diff --git a/docling/models/postprocessing/list_marker_processor.py b/docling/models/postprocessing/list_marker_processor.py index 14302005c3..58c14d20b4 100644 --- a/docling/models/postprocessing/list_marker_processor.py +++ b/docling/models/postprocessing/list_marker_processor.py @@ -128,7 +128,7 @@ def _create_list_item( return item def _find_match( - self, text: str, patterns: list[re.Pattern] + self, text: str, patterns: list[re.Pattern[str]] ) -> Optional[re.Match[str]]: for pattern in patterns: mtch = pattern.match(text) @@ -256,9 +256,7 @@ def process_text_item(self, item: TextItem) -> Union[TextItem, ListItem]: item.enumerated = is_enumerated return item - elif isinstance(item, TextItem) and ( - item.label not in [DocItemLabel.SECTION_HEADER, DocItemLabel.FOOTNOTE] - ): + elif item.label not in [DocItemLabel.SECTION_HEADER, DocItemLabel.FOOTNOTE]: # Create new ListItem return self._create_list_item( self_ref=item.get_ref().cref, @@ -303,6 +301,7 @@ def merge_markers_and_text_items_into_list_items( for ind, (self_ref, is_marker, is_enumerated) in self._matched_items.items(): if is_marker: marker_item = self_ref.resolve(doc=doc) + assert isinstance(marker_item, TextItem) if ind + 1 in self._other: next_item = self._other[ind + 1].resolve(doc=doc) diff --git a/docling/models/postprocessing/reading_order_rb.py b/docling/models/postprocessing/reading_order_rb.py index e45e1f8773..95a2c82ffa 100644 --- a/docling/models/postprocessing/reading_order_rb.py +++ b/docling/models/postprocessing/reading_order_rb.py @@ -13,6 +13,7 @@ from docling_core.types.doc.document import RefItem from docling_core.types.doc.labels import DocItemLabel from rtree import index as rtree_index +from typing_extensions import override _log = logging.getLogger(__name__) @@ -21,7 +22,7 @@ class PageElement(BoundingBox): eps: float = 1.0e-3 cid: int - ref: RefItem = RefItem(cref="#") # type: ignore + ref: RefItem = RefItem(cref="#") text: str = "" @@ -30,6 +31,7 @@ class PageElement(BoundingBox): label: DocItemLabel + @override def __str__(self): return f"{self.cid:6.2f}\t{self.label!s:<10}\t{self.l:6.2f}, {self.b:6.2f}, {self.r:6.2f}, {self.t:6.2f}" @@ -440,7 +442,11 @@ def predict_reading_order( # print("headers ....") for page_no, elems in page_to_headers.items(): - page_to_headers[page_no] = self._predict_page(elems) + page_to_headers[page_no] = [ + node + for node in self._predict_page(list[ReadingOrderNode](elems)) + if isinstance(node, PageElement) + ] # print("elems ....") for page_no, elems in page_to_elems.items(): @@ -465,7 +471,11 @@ def predict_reading_order( # print("footers ....") for page_no, elems in page_to_footers.items(): - page_to_footers[page_no] = self._predict_page(elems) + page_to_footers[page_no] = [ + node + for node in self._predict_page(list[ReadingOrderNode](elems)) + if isinstance(node, PageElement) + ] sorted_elements = [] for page_no in sorted(page_nos): diff --git a/docling/models/stages/chart_extraction/granite_vision.py b/docling/models/stages/chart_extraction/granite_vision.py index f6bcf49639..d10a64c5bd 100644 --- a/docling/models/stages/chart_extraction/granite_vision.py +++ b/docling/models/stages/chart_extraction/granite_vision.py @@ -23,6 +23,7 @@ TabularChartMetaField, ) from docling_core.types.doc.document import CodeMetaField +from typing_extensions import override from docling.datamodel.accelerator_options import AcceleratorOptions from docling.datamodel.base_models import ItemAndImageEnrichmentElement @@ -110,26 +111,28 @@ def __init__( # BaseItemAndImageEnrichmentModel protocol # ------------------------------------------------------------------ + @override def is_processable(self, doc: DoclingDocument, element: NodeItem) -> bool: if not self.enabled: return False if not isinstance(element, PictureItem): return False - if element.meta is None or not isinstance(element.meta, PictureMeta): + if element.meta is None: return False - if element.meta.classification is None or not isinstance( - element.meta.classification, PictureClassificationMetaField - ): + if element.meta.classification is None: return False main_pred = element.meta.classification.get_main_prediction() return main_pred.class_name in SUPPORTED_CHART_TYPES def _resolve_runtime_engine_type(self) -> VlmEngineType: - selected_engine_type = getattr(self.engine, "selected_engine_type", None) + selected_engine_type = ( + self.engine.selected_engine_type if self.engine is not None else None + ) if selected_engine_type is not None: return selected_engine_type return self.options.engine_options.engine_type + @override def __call__( self, doc: DoclingDocument, @@ -140,10 +143,10 @@ def __call__( yield el.item return - elements: List[PictureItem] = [] + elements: List[NodeItem] = [] images = [] for el in element_batch: - elements.append(el.item) # type: ignore[arg-type] + elements.append(el.item) images.append(el.image) # The prompts that are active for each image (one pass per prompt per image). @@ -193,7 +196,7 @@ def __call__( yield item continue - if item.meta is None or not isinstance(item.meta, PictureMeta): + if item.meta is None: item.meta = PictureMeta() handler = _OUTPUT_FORMAT_HANDLERS.get(self.options.output_format) @@ -277,18 +280,18 @@ def _handle_granite_vision_charts(prompt: str, result: str, item: PictureItem) - # --------------------------------------------------------------------------- -def _is_numeric(value: object) -> bool: - if pd.isna(value): # type: ignore[arg-type] +def _is_numeric(value: str | int | float | None) -> bool: + if pd.isna(value): return False try: - float(value) # type: ignore[arg-type] + float(value) return True except (ValueError, TypeError): return False -def _is_blank(value: object) -> bool: - return pd.isna(value) or not str(value).strip() # type: ignore[arg-type] +def _is_blank(value: str | int | float | None) -> bool: + return pd.isna(value) or not str(value).strip() def _is_monotone_integer_row(row: pd.Series) -> bool: @@ -389,4 +392,6 @@ def _extract_python_code(decoded_text: str) -> Optional[str]: python_match = re.search(r"```python\s*\n(.*?)\n```", decoded_text, re.DOTALL) if not python_match: return None - return python_match.group(1).strip() + code = python_match.group(1) + assert isinstance(code, str) + return code.strip() diff --git a/docling/models/stages/code_formula/code_formula_model.py b/docling/models/stages/code_formula/code_formula_model.py index b9d5779fc6..a5f9a65a01 100644 --- a/docling/models/stages/code_formula/code_formula_model.py +++ b/docling/models/stages/code_formula/code_formula_model.py @@ -15,8 +15,10 @@ TextItem, ) from docling_core.types.doc.labels import CodeLanguageLabel +from numpy.typing import NDArray from PIL import Image from pydantic import BaseModel +from typing_extensions import override from docling.datamodel.accelerator_options import AcceleratorDevice, AcceleratorOptions from docling.datamodel.base_models import ItemAndImageEnrichmentElement @@ -111,7 +113,11 @@ def __init__( else: artifacts_path = artifacts_path / self._model_repo_folder - from transformers import AutoModelForImageTextToText, AutoProcessor + # Transformers exposes this model lazily when torch is installed. + from transformers import ( + AutoModelForImageTextToText, # ty: ignore[possibly-missing-import] + AutoProcessor, + ) self._processor = AutoProcessor.from_pretrained( artifacts_path, @@ -136,6 +142,7 @@ def download_models( progress=progress, ) + @override def is_processable(self, doc: DoclingDocument, element: NodeItem) -> bool: """ Determines if a given element in a document can be processed by the model. @@ -249,7 +256,7 @@ def _get_prompt(self, label: str) -> str: prompt = self._processor.apply_chat_template( messages, add_generation_prompt=True ) - + assert isinstance(prompt, str) return prompt def _post_process(self, texts: list[str]) -> list[str]: @@ -282,6 +289,7 @@ def clean_text(text: str) -> str: return [clean_text(t) for t in texts] + @override def __call__( self, doc: DoclingDocument, @@ -308,11 +316,12 @@ def __call__( return labels: List[str] = [] - images: List[Union[Image.Image, np.ndarray]] = [] + images: List[Union[Image.Image, NDArray[np.generic]]] = [] elements: List[TextItem] = [] for el in element_batch: - elements.append(el.item) # type: ignore[arg-type] - labels.append(el.item.label) # type: ignore[attr-defined] + assert isinstance(el.item, TextItem) + elements.append(el.item) + labels.append(el.item.label) images.append(el.image) prompts = [self._get_prompt(label) for label in labels] diff --git a/docling/models/stages/code_formula/code_formula_vlm_model.py b/docling/models/stages/code_formula/code_formula_vlm_model.py index 176287b515..fe5f4dea17 100644 --- a/docling/models/stages/code_formula/code_formula_vlm_model.py +++ b/docling/models/stages/code_formula/code_formula_vlm_model.py @@ -22,7 +22,9 @@ TextItem, ) from docling_core.types.doc.labels import CodeLanguageLabel +from numpy.typing import NDArray from PIL import Image +from typing_extensions import override from docling.datamodel.accelerator_options import AcceleratorOptions from docling.datamodel.base_models import ItemAndImageEnrichmentElement @@ -122,6 +124,7 @@ def __init__( _log.info("CodeFormulaVlmModel initialized successfully") + @override def is_processable(self, doc: DoclingDocument, element: NodeItem) -> bool: """Determine if an element can be processed by this stage. @@ -227,6 +230,7 @@ def clean_text(text: str) -> str: return [clean_text(t) for t in texts] + @override def __call__( self, doc: DoclingDocument, @@ -250,7 +254,7 @@ def __call__( raise RuntimeError("Engine not initialized") labels: List[str] = [] - images: List[Union[Image.Image, np.ndarray]] = [] + images: List[Union[Image.Image, NDArray[np.generic]]] = [] elements: List[Union[CodeItem, TextItem]] = [] for el in element_batch: @@ -259,7 +263,7 @@ def __call__( labels.append(el.item.label) images.append(el.image) - extra_generation_config: dict = { + extra_generation_config: dict[str, object] = { "skip_special_tokens": False, # Keep special tokens for post-processing } # API engines switch to streaming when given a stopper, so they only get @@ -318,5 +322,5 @@ def __del__(self): self.engine.cleanup() except Exception as e: # _log may be None during interpreter shutdown - if _log is not None: + if _log is not None: # ty: ignore[redundant-condition-strict] _log.warning(f"Error cleaning up engine: {e}") diff --git a/docling/models/stages/heading_hierarchy/heading_hierarchy_model.py b/docling/models/stages/heading_hierarchy/heading_hierarchy_model.py index bfff966c2f..f6780c813c 100644 --- a/docling/models/stages/heading_hierarchy/heading_hierarchy_model.py +++ b/docling/models/stages/heading_hierarchy/heading_hierarchy_model.py @@ -427,7 +427,7 @@ def _item_page_and_top( return None, None prov = item.prov[0] page = document.pages.get(prov.page_no) - if page is not None and page.size is not None: + if page is not None: return prov.page_no, prov.bbox.to_top_left_origin(page.size.height).t return prov.page_no, None diff --git a/docling/models/stages/layout/layout_model.py b/docling/models/stages/layout/layout_model.py index ff6beb9183..3010aeca20 100644 --- a/docling/models/stages/layout/layout_model.py +++ b/docling/models/stages/layout/layout_model.py @@ -13,6 +13,8 @@ from pathlib import Path from typing import Optional +from typing_extensions import override + from docling.datamodel.accelerator_options import AcceleratorOptions from docling.datamodel.layout_model_specs import ( DOCLING_LAYOUT_HERON, @@ -80,6 +82,7 @@ def __init__( ) @classmethod + @override def get_options_type(cls) -> type[LayoutOptions]: # ty: ignore[invalid-method-override] # LayoutFactory dispatches on the exact options type, so the shim must # keep claiming LayoutOptions rather than inheriting the parent's. diff --git a/docling/models/stages/layout/layout_object_detection_model.py b/docling/models/stages/layout/layout_object_detection_model.py index 9a8a3b5e68..aeed302457 100644 --- a/docling/models/stages/layout/layout_object_detection_model.py +++ b/docling/models/stages/layout/layout_object_detection_model.py @@ -11,6 +11,7 @@ from docling_core.types.doc import BoundingBox, CoordOrigin, DocItemLabel from PIL import Image +from typing_extensions import override from docling.datamodel.accelerator_options import AcceleratorOptions from docling.datamodel.base_models import Cluster, LayoutPrediction, Page @@ -78,9 +79,11 @@ def _build_label_map(self) -> Dict[int, DocItemLabel]: return label_map @classmethod + @override def get_options_type(cls) -> type[LayoutObjectDetectionOptions]: return LayoutObjectDetectionOptions + @override def predict_layout( self, conv_res: ConversionResult, diff --git a/docling/models/stages/layout/layout_postprocessing_model.py b/docling/models/stages/layout/layout_postprocessing_model.py index 63ac2c7589..eb1e497447 100644 --- a/docling/models/stages/layout/layout_postprocessing_model.py +++ b/docling/models/stages/layout/layout_postprocessing_model.py @@ -14,6 +14,7 @@ from typing import List, Type import numpy as np +from typing_extensions import override from docling.datamodel.base_models import Cluster, LayoutPrediction, Page from docling.datamodel.document import ConversionResult @@ -34,9 +35,11 @@ def __init__(self, *, options: LayoutPostprocessorOptions) -> None: self.options = options @classmethod + @override def get_options_type(cls) -> Type[LayoutPostprocessorOptions]: return LayoutPostprocessorOptions + @override def postprocess_layout( self, conv_res: ConversionResult, diff --git a/docling/models/stages/ocr/_nemotron_ocr_model.py b/docling/models/stages/ocr/_nemotron_ocr_model.py index 28001bf15b..de03552c01 100644 --- a/docling/models/stages/ocr/_nemotron_ocr_model.py +++ b/docling/models/stages/ocr/_nemotron_ocr_model.py @@ -14,6 +14,8 @@ import numpy from docling_core.types.doc import BoundingBox, CoordOrigin from docling_core.types.doc.page import BoundingRectangle, TextCell +from numpy.typing import NDArray +from typing_extensions import override from docling.datamodel.accelerator_options import AcceleratorOptions from docling.datamodel.base_models import Page @@ -99,7 +101,7 @@ class _BufferedRect: state: _PageOcrState ocr_rect: BoundingBox - image: numpy.ndarray + image: NDArray[numpy.generic] image_size: tuple[int, int] @@ -153,6 +155,7 @@ def __init__( lang=code, ) + @override def supported_ocr_languages(self) -> OcrLanguageSupport: r"""Report the BCP74 and native languages without script whenever it is not needed""" return OcrLanguageSupport( @@ -163,6 +166,7 @@ def supported_ocr_languages(self) -> OcrLanguageSupport: native=sorted(_NEMOTRON_CODE_TO_ARTIFACT), ) + @override def map_ocr_language(self, language: OcrLanguage) -> str: if language.is_passthrough(): # `english`, `multilingual`: nemotron's own recognizer names. @@ -350,6 +354,7 @@ def _run_buffer( buffer.clear() + @override def __call__( self, conv_res: ConversionResult, page_batch: Iterable[Page] ) -> Iterable[Page]: @@ -436,5 +441,6 @@ def drain_completed() -> Iterable[Page]: yield from drain_completed() @classmethod + @override def get_options_type(cls) -> Type[OcrOptions]: return NemotronOcrOptions diff --git a/docling/models/stages/ocr/auto_ocr_model.py b/docling/models/stages/ocr/auto_ocr_model.py index 3ff1151f5c..fb7f4b7e87 100644 --- a/docling/models/stages/ocr/auto_ocr_model.py +++ b/docling/models/stages/ocr/auto_ocr_model.py @@ -7,6 +7,8 @@ from pathlib import Path from typing import NamedTuple, Optional, Type +from typing_extensions import override + from docling.datamodel.accelerator_options import AcceleratorOptions from docling.datamodel.base_models import Page from docling.datamodel.document import ConversionResult @@ -121,7 +123,7 @@ def __init__( if self._engine is None: try: import onnxruntime - from rapidocr import EngineType, RapidOCR # type: ignore + from rapidocr import EngineType, RapidOCR self._engine = RapidOcrModel( enabled=self.enabled, @@ -175,7 +177,7 @@ def __init__( if self._engine is None: try: import torch - from rapidocr import EngineType, RapidOCR # type: ignore + from rapidocr import EngineType, RapidOCR self._engine = RapidOcrModel( enabled=self.enabled, @@ -216,6 +218,7 @@ def __init__( ) _log.warning("No OCR engine found. Please review the install details.") + @override def __call__( self, conv_res: ConversionResult, page_batch: Iterable[Page] ) -> Iterable[Page]: @@ -225,5 +228,6 @@ def __call__( yield from self._engine(conv_res, page_batch) @classmethod + @override def get_options_type(cls) -> Type[OcrOptions]: return OcrAutoOptions diff --git a/docling/models/stages/ocr/easyocr_model.py b/docling/models/stages/ocr/easyocr_model.py index 82c5419b00..040a4bcbe0 100644 --- a/docling/models/stages/ocr/easyocr_model.py +++ b/docling/models/stages/ocr/easyocr_model.py @@ -13,6 +13,7 @@ import numpy from docling_core.types.doc import BoundingBox, CoordOrigin from docling_core.types.doc.page import BoundingRectangle, TextCell +from typing_extensions import override from docling.datamodel.accelerator_options import AcceleratorDevice, AcceleratorOptions from docling.datamodel.base_models import Page @@ -229,6 +230,7 @@ def __init__( verbose=False, ) + @override def supported_ocr_languages(self) -> OcrLanguageSupport: r"""Report the native and BCP74 languages without script whenever it is not needed""" tags: set[str] = set() @@ -246,6 +248,7 @@ def supported_ocr_languages(self) -> OcrLanguageSupport: native.add(code) return OcrLanguageSupport(bcp47=sorted(tags), native=sorted(native)) + @override def map_ocr_language(self, language: OcrLanguage) -> str | List[str]: code = _easyocr_code(language) if code is None: @@ -305,6 +308,7 @@ def download_models( return local_dir + @override def __call__( self, conv_res: ConversionResult, page_batch: Iterable[Page] ) -> Iterable[Page]: @@ -375,5 +379,6 @@ def __call__( yield page @classmethod + @override def get_options_type(cls) -> Type[OcrOptions]: return EasyOcrOptions diff --git a/docling/models/stages/ocr/kserve_v2_ocr_model.py b/docling/models/stages/ocr/kserve_v2_ocr_model.py index cb18e55717..e8f90b4b70 100644 --- a/docling/models/stages/ocr/kserve_v2_ocr_model.py +++ b/docling/models/stages/ocr/kserve_v2_ocr_model.py @@ -11,7 +11,9 @@ import numpy as np from docling_core.types.doc import BoundingBox, CoordOrigin from docling_core.types.doc.page import BoundingRectangle, TextCell +from numpy.typing import NDArray from PIL import Image +from typing_extensions import override from docling.datamodel.accelerator_options import AcceleratorOptions from docling.datamodel.base_models import ( @@ -141,7 +143,7 @@ def _initialize_client(self) -> None: self.options.transport, ) - def _preprocess_image(self, image: Image.Image) -> np.ndarray: + def _preprocess_image(self, image: Image.Image) -> NDArray[np.generic]: """Preprocess image for KServe v2 OCR inference. Converts PIL image to numpy array with shape (1, H, W, C) in UINT8 format, @@ -163,9 +165,9 @@ def _preprocess_image(self, image: Image.Image) -> np.ndarray: def _create_text_cells( self, - boxes: np.ndarray, - txts: np.ndarray, - scores: np.ndarray, + boxes: NDArray[np.generic], + txts: NDArray[np.generic], + scores: NDArray[np.generic], ocr_rect: BoundingBox, ) -> List[TextCell]: """Convert KServe v2 OCR outputs to TextCell objects. @@ -213,6 +215,7 @@ def _create_text_cells( return cells + @override def __call__( self, conv_res: ConversionResult, page_batch: Iterable[Page] ) -> Iterable[Page]: @@ -307,6 +310,7 @@ def __call__( yield page @classmethod + @override def get_options_type(cls) -> Type[OcrOptions]: """Get the options type for this OCR model. diff --git a/docling/models/stages/ocr/nemotron_ocr_model.py b/docling/models/stages/ocr/nemotron_ocr_model.py index 51f57721b2..8c81baab0d 100644 --- a/docling/models/stages/ocr/nemotron_ocr_model.py +++ b/docling/models/stages/ocr/nemotron_ocr_model.py @@ -14,6 +14,8 @@ import numpy from docling_core.types.doc import BoundingBox, CoordOrigin from docling_core.types.doc.page import BoundingRectangle, TextCell +from numpy.typing import NDArray +from typing_extensions import override from docling.datamodel.accelerator_options import AcceleratorOptions from docling.datamodel.base_models import Page @@ -118,7 +120,7 @@ class _BufferedRect: state: _PageOcrState ocr_rect: BoundingBox - image: numpy.ndarray + image: NDArray[numpy.generic] image_size: tuple[int, int] @@ -172,6 +174,7 @@ def __init__( lang=code, ) + @override def supported_ocr_languages(self) -> OcrLanguageSupport: r"""Report the BCP74 and native languages without script whenever it is not needed""" return OcrLanguageSupport( @@ -184,6 +187,7 @@ def supported_ocr_languages(self) -> OcrLanguageSupport: native=sorted(_NEMOTRON_CODE_TO_ARTIFACT), ) + @override def map_ocr_language(self, language: OcrLanguage) -> str: if language.is_passthrough(): # `english`, `multilingual`: nemotron's own recognizer names. @@ -381,6 +385,7 @@ def _run_buffer( buffer.clear() + @override def __call__( self, conv_res: ConversionResult, page_batch: Iterable[Page] ) -> Iterable[Page]: @@ -467,5 +472,6 @@ def drain_completed() -> Iterable[Page]: yield from drain_completed() @classmethod + @override def get_options_type(cls) -> Type[OcrOptions]: return NemotronOcrOptions diff --git a/docling/models/stages/ocr/ocr_mac_model.py b/docling/models/stages/ocr/ocr_mac_model.py index 36769e7ce9..f760671a53 100644 --- a/docling/models/stages/ocr/ocr_mac_model.py +++ b/docling/models/stages/ocr/ocr_mac_model.py @@ -10,6 +10,7 @@ from docling_core.types.doc import BoundingBox, CoordOrigin from docling_core.types.doc.page import BoundingRectangle, TextCell +from typing_extensions import override from docling.datamodel.accelerator_options import AcceleratorOptions from docling.datamodel.base_models import Page @@ -43,10 +44,11 @@ def _get_vision_languages(recognition: str) -> list[str]: "The OcrMac engine cannot be used on this system." ) try: - import Vision + import Vision # ty: ignore[unresolved-import] # pyobjc exposes the ObjC classes dynamically, so ty cannot see them. - request = Vision.VNRecognizeTextRequest.alloc().init() # ty: ignore[unresolved-attribute] + # pyobjc publishes this class dynamically at runtime. + request = Vision.VNRecognizeTextRequest.alloc().init() # Same mapping as ocrmac: 1 is fast, 0 (the default) is accurate. request.setRecognitionLevel_(1 if recognition == "fast" else 0) languages, error = request.supportedRecognitionLanguagesAndReturnError_(None) @@ -99,6 +101,7 @@ def __init__( self._vision_languages = _get_vision_languages(self.options.recognition) self._native_codes = self.resolve_ocr_languages() + @override def supported_ocr_languages(self) -> OcrLanguageSupport: # Map the Vision language tags to the canonical tags. tags: set[str] = set() @@ -128,6 +131,7 @@ def supported_ocr_languages(self) -> OcrLanguageSupport: return OcrLanguageSupport(bcp47=sorted(tags), native=sorted(native)) + @override def map_ocr_language(self, language: OcrLanguage) -> str | list[str]: if language.is_passthrough(): # One of Vision's own recognition languages @@ -148,6 +152,7 @@ def map_ocr_language(self, language: OcrLanguage) -> str | list[str]: ) return code + @override def __call__( self, conv_res: ConversionResult, page_batch: Iterable[Page] ) -> Iterable[Page]: @@ -178,7 +183,7 @@ def __call__( fname = image_file.name high_res_image.save(fname) - boxes = self.reader_RIL( + boxes = self.reader_RIL( # ty: ignore[unresolved-attribute] fname, recognition_level=self.options.recognition, framework=self.options.framework, @@ -235,5 +240,6 @@ def __call__( yield page @classmethod + @override def get_options_type(cls) -> Type[OcrOptions]: return OcrMacOptions diff --git a/docling/models/stages/ocr/rapid_ocr_model.py b/docling/models/stages/ocr/rapid_ocr_model.py index a8ada94930..7465e02041 100644 --- a/docling/models/stages/ocr/rapid_ocr_model.py +++ b/docling/models/stages/ocr/rapid_ocr_model.py @@ -11,6 +11,7 @@ import numpy from docling_core.types.doc import BoundingBox, CoordOrigin from docling_core.types.doc.page import BoundingRectangle, TextCell +from typing_extensions import override from docling.datamodel.accelerator_options import AcceleratorDevice, AcceleratorOptions from docling.datamodel.base_models import Page @@ -55,7 +56,7 @@ # Inference backends docling supports, derived from the RapidOcrOptions.backend # annotation so the two cannot drift. -_RAPIDOCR_BACKENDS: frozenset[str] = frozenset(get_args(RapidOcrBackend)) +_RAPIDOCR_BACKENDS: frozenset[str] = frozenset[str](get_args(RapidOcrBackend)) # Canonical BCP-47 to PP-OCR recognizer codes. The PP-OCRv6 set comes from the # installed `rapidocr`; the v4/v5 sets below have no upstream equivalent to read, @@ -483,7 +484,7 @@ def __init__( if self.enabled: try: - from rapidocr import ModelType, OCRVersion, RapidOCR # type: ignore + from rapidocr import ModelType, OCRVersion, RapidOCR except ImportError: raise ImportError( "RapidOCR is not installed. Please install it via `pip install rapidocr onnxruntime` to use this OCR engine. " @@ -667,9 +668,11 @@ def __init__( params=params, ) + @override def supported_ocr_languages(self) -> OcrLanguageSupport: return _ppocr_supported_languages(_rapidocr_vocabulary(self.options.backend)) + @override def resolve_ocr_languages(self) -> list[str]: # An empty `lang` list means "the engine's own default", which for PP-OCR # is the Simplified Chinese recognizer. Resolved the same way a request @@ -687,6 +690,7 @@ def resolve_ocr_languages(self) -> list[str]: return [code] return super().resolve_ocr_languages() + @override def map_ocr_language(self, language: OcrLanguage) -> str: code = _ppocr_code(language, _rapidocr_vocabulary(self.options.backend)) if code is None: @@ -741,6 +745,7 @@ def download_models( fw.write(buf.read()) return local_dir + @override def __call__( self, conv_res: ConversionResult, page_batch: Iterable[Page] ) -> Iterable[Page]: @@ -748,6 +753,8 @@ def __call__( yield from page_batch return + from rapidocr.utils.output import RapidOCROutput + for page in page_batch: assert page._backend is not None if not page._backend.is_valid(): @@ -771,43 +778,43 @@ def __call__( use_cls=self.options.use_cls, use_rec=self.options.use_rec, ) - if result is None or result.boxes is None: + if ( + not isinstance(result, RapidOCROutput) + or result.boxes is None + or result.txts is None + or result.scores is None + ): _log.warning("RapidOCR returned empty result!") continue - result = list( + rows = list( zip(result.boxes.tolist(), result.txts, result.scores) ) del high_res_image del im - if result is not None: - cells = [ - TextCell( - index=ix, - text=line[1], - orig=line[1], - confidence=line[2], - from_ocr=True, - rect=BoundingRectangle.from_bounding_box( - BoundingBox.from_tuple( - coord=( - (line[0][0][0] / self.scale) - + ocr_rect.l, - (line[0][0][1] / self.scale) - + ocr_rect.t, - (line[0][2][0] / self.scale) - + ocr_rect.l, - (line[0][2][1] / self.scale) - + ocr_rect.t, - ), - origin=CoordOrigin.TOPLEFT, - ) - ), - ) - for ix, line in enumerate(result) - ] - all_ocr_cells.extend(cells) + cells = [ + TextCell( + index=ix, + text=line[1], + orig=line[1], + confidence=line[2], + from_ocr=True, + rect=BoundingRectangle.from_bounding_box( + BoundingBox.from_tuple( + coord=( + (line[0][0][0] / self.scale) + ocr_rect.l, + (line[0][0][1] / self.scale) + ocr_rect.t, + (line[0][2][0] / self.scale) + ocr_rect.l, + (line[0][2][1] / self.scale) + ocr_rect.t, + ), + origin=CoordOrigin.TOPLEFT, + ) + ), + ) + for ix, line in enumerate(rows) + ] + all_ocr_cells.extend(cells) # Post-process the cells self.post_process_cells(all_ocr_cells, page, conv_res) @@ -819,5 +826,6 @@ def __call__( yield page @classmethod + @override def get_options_type(cls) -> Type[OcrOptions]: return RapidOcrOptions diff --git a/docling/models/stages/ocr/tesseract_ocr_cli_model.py b/docling/models/stages/ocr/tesseract_ocr_cli_model.py index 8b42884d87..db771778b8 100644 --- a/docling/models/stages/ocr/tesseract_ocr_cli_model.py +++ b/docling/models/stages/ocr/tesseract_ocr_cli_model.py @@ -16,6 +16,7 @@ import pandas as pd from docling_core.types.doc import BoundingBox, CoordOrigin from docling_core.types.doc.page import TextCell +from typing_extensions import override from docling.datamodel.accelerator_options import AcceleratorOptions from docling.datamodel.base_models import Page @@ -111,9 +112,11 @@ def __init__( self._sanitize_lang(lang) for lang in self.resolve_ocr_languages() ] + @override def supported_ocr_languages(self) -> OcrLanguageSupport: return installed_tesseract_languages(self._tesseract_vocabulary or []) + @override def map_ocr_language(self, language: OcrLanguage) -> str | List[str]: assert self._tesseract_vocabulary is not None name = language_to_tesseract_code(language) @@ -173,7 +176,7 @@ def _sanitize_filename(filename: str) -> str: def _get_name_and_version(self) -> Tuple[str, str]: if self._name is not None and self._version is not None: - return self._name, self._version # type: ignore + return self._name, self._version cmd = [self._safe_tesseract_cmd, "--version"] @@ -313,6 +316,7 @@ def _set_languages(self): df_list = pd.read_csv(io.StringIO(decoded_data), header=None) self._tesseract_vocabulary = tesseract_vocabulary(df_list[0].tolist()[1:]) + @override def __call__( self, conv_res: ConversionResult, page_batch: Iterable[Page] ) -> Iterable[Page]: @@ -427,6 +431,7 @@ def __call__( yield page @classmethod + @override def get_options_type(cls) -> Type[OcrOptions]: return TesseractCliOcrOptions diff --git a/docling/models/stages/ocr/tesseract_ocr_model.py b/docling/models/stages/ocr/tesseract_ocr_model.py index 800e5e2aec..a0b6575a44 100644 --- a/docling/models/stages/ocr/tesseract_ocr_model.py +++ b/docling/models/stages/ocr/tesseract_ocr_model.py @@ -5,10 +5,11 @@ import logging from pathlib import Path -from typing import Iterable, Optional, Type +from typing import TYPE_CHECKING, Iterable, Optional, Type, TypedDict from docling_core.types.doc import BoundingBox, CoordOrigin from docling_core.types.doc.page import TextCell +from typing_extensions import NotRequired, override from docling.datamodel.accelerator_options import AcceleratorOptions from docling.datamodel.base_models import Page @@ -36,6 +37,15 @@ _log = logging.getLogger(__name__) +if TYPE_CHECKING: + import tesserocr + + +class _TesserocrInitKwargs(TypedDict): + init: bool + oem: tesserocr.OEM + path: NotRequired[str] + class TesseractOcrModel(BaseOcrModel): multiple_languages = True @@ -112,7 +122,7 @@ def __init__( # Needs the installed language list and the prefix, so it runs here. self._native_codes = self.resolve_ocr_languages() - tesserocr_kwargs = { + tesserocr_kwargs: _TesserocrInitKwargs = { "init": True, "oem": tesserocr.OEM.DEFAULT, } @@ -128,11 +138,15 @@ def __init__( ) if self._auto_script: # No `lang`: the per-page OSD pass picks a script reader instead. - self._reader = tesserocr.PyTessBaseAPI(psm=main_psm, **tesserocr_kwargs) + # The native API accepts int PSM values; its stub requires PSM. + self._reader = tesserocr.PyTessBaseAPI( + psm=main_psm, # ty: ignore[invalid-argument-type] + **tesserocr_kwargs, + ) else: self._reader = tesserocr.PyTessBaseAPI( lang="+".join(self._native_codes), - psm=main_psm, + psm=main_psm, # ty: ignore[invalid-argument-type] **tesserocr_kwargs, ) # OSD reader must use PSM.OSD_ONLY for orientation detection @@ -141,9 +155,11 @@ def __init__( ) self._reader_RIL = tesserocr.RIL + @override def supported_ocr_languages(self) -> OcrLanguageSupport: return installed_tesseract_languages(self._tesseract_vocabulary) + @override def map_ocr_language(self, language: OcrLanguage) -> str | list[str]: name = language_to_tesseract_code(language) if name not in self._tesseract_vocabulary: @@ -162,6 +178,7 @@ def __del__(self): for reader in self._script_readers.values(): reader.End() + @override def __call__( self, conv_res: ConversionResult, page_batch: Iterable[Page] ) -> Iterable[Page]: @@ -197,7 +214,8 @@ def __call__( osd = self._osd_reader.DetectOrientationScript() # No text, or Orientation and Script detection failure - if osd is None: + # The method documents None on failure despite its dict annotation. + if osd is None: # ty: ignore[redundant-condition-strict] _log.error( "OSD failed for doc (doc %s, page: %s, " "OCR rectangle: %s)", @@ -234,9 +252,11 @@ def __call__( tesserocr.PyTessBaseAPI( path=self._reader.GetDatapath(), lang=lang, - psm=self.options.psm - if self.options.psm is not None - else tesserocr.PSM.AUTO, + psm=( + self.options.psm + if self.options.psm is not None + else tesserocr.PSM.AUTO + ), # ty: ignore[invalid-argument-type] init=True, oem=tesserocr.OEM.DEFAULT, ) @@ -299,5 +319,6 @@ def __call__( yield page @classmethod + @override def get_options_type(cls) -> Type[OcrOptions]: return TesseractOcrOptions diff --git a/docling/models/stages/ocr/tesseract_utils.py b/docling/models/stages/ocr/tesseract_utils.py index 447a0e2d61..fd382955c0 100644 --- a/docling/models/stages/ocr/tesseract_utils.py +++ b/docling/models/stages/ocr/tesseract_utils.py @@ -156,13 +156,12 @@ def tesseract_box_to_bounding_rectangle( if original_offset.coord_origin is not CoordOrigin.TOPLEFT: msg = f"expected coordinate origin to be {CoordOrigin.TOPLEFT.value}" raise ValueError(msg) - if original_offset is not None: - rect.r_x0 += original_offset.l - rect.r_x1 += original_offset.l - rect.r_x2 += original_offset.l - rect.r_x3 += original_offset.l - rect.r_y0 += original_offset.t - rect.r_y1 += original_offset.t - rect.r_y2 += original_offset.t - rect.r_y3 += original_offset.t + rect.r_x0 += original_offset.l + rect.r_x1 += original_offset.l + rect.r_x2 += original_offset.l + rect.r_x3 += original_offset.l + rect.r_y0 += original_offset.t + rect.r_y1 += original_offset.t + rect.r_y2 += original_offset.t + rect.r_y3 += original_offset.t return rect diff --git a/docling/models/stages/page_assemble/page_assemble_model.py b/docling/models/stages/page_assemble/page_assemble_model.py index a7bd223f63..d73e50416b 100644 --- a/docling/models/stages/page_assemble/page_assemble_model.py +++ b/docling/models/stages/page_assemble/page_assemble_model.py @@ -10,6 +10,7 @@ import numpy as np from docling_core.types.doc import BoundingBox from pydantic import AnyUrl, BaseModel, ValidationError +from typing_extensions import override from docling.datamodel.base_models import ( AssembledUnit, @@ -177,6 +178,7 @@ def sanitize_text(self, lines: List[str]) -> str: return sanitized_text.strip() # Strip any leading or trailing whitespace + @override def __call__( self, conv_res: ConversionResult, page_batch: Iterable[Page] ) -> Iterable[Page]: @@ -248,7 +250,6 @@ def __call__( label=cluster.label, id=cluster.id, text="", - data=None, cluster=cluster, page_no=page.page_no, ) diff --git a/docling/models/stages/page_preprocessing/page_preprocessing_model.py b/docling/models/stages/page_preprocessing/page_preprocessing_model.py index 26a188fb7e..db52e84fb9 100644 --- a/docling/models/stages/page_preprocessing/page_preprocessing_model.py +++ b/docling/models/stages/page_preprocessing/page_preprocessing_model.py @@ -13,6 +13,7 @@ from PIL import ImageDraw from PIL.Image import Image from pydantic import BaseModel +from typing_extensions import override from docling.datamodel.base_models import Page from docling.datamodel.document import ConversionResult @@ -56,6 +57,7 @@ def __init__(self, options: PagePreprocessingOptions): r"(?:/\w+\s*){2,}" ) # Two or more "/token " sequences + @override def __call__( self, conv_res: ConversionResult, page_batch: Iterable[Page] ) -> Iterable[Page]: diff --git a/docling/models/stages/picture_classifier/document_picture_classifier.py b/docling/models/stages/picture_classifier/document_picture_classifier.py index 9e17d3a9b5..fefa638194 100644 --- a/docling/models/stages/picture_classifier/document_picture_classifier.py +++ b/docling/models/stages/picture_classifier/document_picture_classifier.py @@ -16,7 +16,9 @@ PictureMeta, ) from docling_core.types.doc.document import PictureClassificationPrediction +from numpy.typing import NDArray from PIL import Image +from typing_extensions import override from docling.datamodel.accelerator_options import AcceleratorOptions from docling.datamodel.base_models import ItemAndImageEnrichmentElement @@ -105,6 +107,7 @@ def __init__( self.engine.initialize() self._classes = self.engine.get_label_mapping() + @override def is_processable(self, doc: DoclingDocument, element: NodeItem) -> bool: """ Determines if the given element can be processed by the classifier. @@ -123,6 +126,7 @@ def is_processable(self, doc: DoclingDocument, element: NodeItem) -> bool: """ return self.enabled and isinstance(element, PictureItem) + @override def __call__( self, doc: DoclingDocument, @@ -152,21 +156,13 @@ def __call__( if self.engine is None: raise RuntimeError("Picture classifier engine is not initialized.") - images: List[Union[Image.Image, np.ndarray]] = [] + images: List[Union[Image.Image, NDArray[np.generic]]] = [] elements: List[PictureItem] = [] for i, el in enumerate(element_batch): assert isinstance(el.item, PictureItem) elements.append(el.item) - raw_image = el.image - if isinstance(raw_image, Image.Image): - raw_image = raw_image.convert("RGB") - elif isinstance(raw_image, np.ndarray): - raw_image = Image.fromarray(raw_image).convert("RGB") - else: - raise TypeError( - "Supported input formats are PIL.Image.Image or numpy.ndarray." - ) + raw_image = el.image.convert("RGB") images.append(raw_image) engine_input_batch = [ @@ -205,7 +201,7 @@ def __call__( predictions=predictions, ) - if item.meta is not None and isinstance(item.meta, PictureMeta): + if item.meta is not None: item.meta.classification = classification_data else: item.meta = PictureMeta(classification=classification_data) diff --git a/docling/models/stages/picture_description/picture_description_api_model.py b/docling/models/stages/picture_description/picture_description_api_model.py index 6d9c736535..7d37456ab6 100644 --- a/docling/models/stages/picture_description/picture_description_api_model.py +++ b/docling/models/stages/picture_description/picture_description_api_model.py @@ -7,6 +7,7 @@ from typing import Optional, Type, Union from PIL import Image +from typing_extensions import override from docling.datamodel.accelerator_options import AcceleratorOptions from docling.datamodel.base_models import ApiImageRequestResult @@ -23,6 +24,7 @@ class PictureDescriptionApiModel(PictureDescriptionBaseModel): # elements_batch_size = 4 @classmethod + @override def get_options_type(cls) -> Type[PictureDescriptionBaseOptions]: return PictureDescriptionApiOptions @@ -51,6 +53,7 @@ def __init__( "pipeline_options.enable_remote_services=True." ) + @override def _annotate_images( self, images: Iterable[Image.Image] ) -> Iterable[ApiImageRequestResult]: diff --git a/docling/models/stages/picture_description/picture_description_vlm_engine_model.py b/docling/models/stages/picture_description/picture_description_vlm_engine_model.py index 2de16a721c..275de42d6f 100644 --- a/docling/models/stages/picture_description/picture_description_vlm_engine_model.py +++ b/docling/models/stages/picture_description/picture_description_vlm_engine_model.py @@ -13,6 +13,7 @@ from typing import Optional, Type, Union from PIL import Image +from typing_extensions import override from docling.datamodel.accelerator_options import AcceleratorOptions from docling.datamodel.base_models import ApiImageRequestResult, VlmStopReason @@ -62,6 +63,7 @@ class PictureDescriptionVlmEngineModel(PictureDescriptionBaseModel): """ @classmethod + @override def get_options_type(cls) -> Type[PictureDescriptionBaseOptions]: return PictureDescriptionVlmEngineOptions @@ -110,8 +112,12 @@ def __init__( self.provenance = f"{self.repo_id} ({engine_type.value})" def _resolve_runtime_engine_type(self) -> VlmEngineType: - selected_engine_type = getattr(self.engine, "selected_engine_type", None) - if selected_engine_type is not None: + selected_engine_type = ( + getattr(self.engine, "selected_engine_type", None) + if self.engine is not None + else None + ) + if isinstance(selected_engine_type, VlmEngineType): return selected_engine_type return self.options.engine_options.engine_type @@ -154,6 +160,7 @@ def _build_engine_inputs( for image in image_list ] + @override def _annotate_images( self, images: Iterable[Image.Image] ) -> Iterable[ApiImageRequestResult]: diff --git a/docling/models/stages/picture_description/picture_description_vlm_model.py b/docling/models/stages/picture_description/picture_description_vlm_model.py index 7c4c8947b0..026ff0d412 100644 --- a/docling/models/stages/picture_description/picture_description_vlm_model.py +++ b/docling/models/stages/picture_description/picture_description_vlm_model.py @@ -8,6 +8,7 @@ from typing import Optional, Type, Union from PIL import Image +from typing_extensions import override from docling.datamodel.accelerator_options import AcceleratorOptions from docling.datamodel.pipeline_options import ( @@ -29,6 +30,7 @@ class PictureDescriptionVlmModel( PictureDescriptionBaseModel, HuggingFaceModelDownloadMixin ): @classmethod + @override def get_options_type(cls) -> Type[PictureDescriptionBaseOptions]: return PictureDescriptionVlmOptions @@ -60,7 +62,8 @@ def __init__( try: import torch from transformers import ( - AutoModelForImageTextToText, + # Transformers exposes this model lazily when torch is installed. + AutoModelForImageTextToText, # ty: ignore[possibly-missing-import] AutoProcessor, ) except ImportError: @@ -86,12 +89,13 @@ def __init__( ), ) if sys.version_info < (3, 14): - self.model = torch.compile(self.model) # type: ignore + self.model = torch.compile(self.model) else: self.model.eval() self.provenance = f"{self.options.repo_id}" + @override def _annotate_images(self, images: Iterable[Image.Image]) -> Iterable[str]: from transformers import GenerationConfig diff --git a/docling/models/stages/reading_order/readingorder_model.py b/docling/models/stages/reading_order/readingorder_model.py index 28909aab70..2c13e67ac1 100644 --- a/docling/models/stages/reading_order/readingorder_model.py +++ b/docling/models/stages/reading_order/readingorder_model.py @@ -1,7 +1,7 @@ # SPDX-FileCopyrightText: The Docling Contributors # SPDX-License-Identifier: MIT -from collections.abc import Iterator +from collections.abc import Iterator, Sequence from pathlib import Path from statistics import median @@ -68,7 +68,9 @@ def _assembled_to_readingorder_elements( page_no_to_pages = {p.page_no: p for p in conv_res.pages} for element in conv_res.assembled.elements: - page_height = page_no_to_pages[element.page_no].size.height # type: ignore + page_size = page_no_to_pages[element.page_no].size + assert page_size is not None + page_height = page_size.height bbox = element.cluster.bbox.to_bottom_left_origin(page_height) text = element.text or "" @@ -78,7 +80,7 @@ def _assembled_to_readingorder_elements( ref=RefItem(cref=self._element_ref(element)), text=text, page_no=element.page_no, - page_size=page_no_to_pages[element.page_no].size, + page_size=page_size, label=element.label, l=bbox.l, r=bbox.r, @@ -151,6 +153,7 @@ def _add_table_element( parent=parent, ) + assert isinstance(table_item, TableItem) return table_item def _add_table_children( @@ -195,7 +198,9 @@ def _add_picture_element( charspan=(0, 0), bbox=element.cluster.bbox.to_bottom_left_origin(page_height), ) - return out_doc.add_picture(prov=prov, parent=parent) + picture_item = out_doc.add_picture(prov=prov, parent=parent) + assert isinstance(picture_item, PictureItem) + return picture_item def _add_code_element( self, @@ -209,7 +214,9 @@ def _add_code_element( charspan=(0, len(element.text)), bbox=element.cluster.bbox.to_bottom_left_origin(page_height), ) - return out_doc.add_code(text=element.text, prov=prov, parent=parent) + code_item = out_doc.add_code(text=element.text, prov=prov, parent=parent) + assert isinstance(code_item, CodeItem) + return code_item def _add_related_text_items( self, @@ -279,7 +286,7 @@ def _element_ref(element: BasePageElement) -> str: @classmethod def _match_table_pictures( cls, - elements: list[BasePageElement], + elements: Sequence[BasePageElement], excluded_picture_refs: set[str], ) -> dict[str, dict[int, list[FigureElement]]]: tables = [element for element in elements if isinstance(element, Table)] @@ -558,7 +565,7 @@ def materialize_siblings( ) self._add_child_elements(element, picture_item, out_doc) - elif isinstance(element, ContainerElement): + else: group_label = ( GroupLabel.FORM_AREA if element.label == DocItemLabel.FORM @@ -655,6 +662,7 @@ def _handle_text_element( content_layer=content_layer, hyperlink=element.hyperlink, ) + assert isinstance(new_item, TextItem) return new_item, current_list def _merge_elements(self, element, merged_elem, new_item, page_height): diff --git a/docling/models/stages/table_structure/table_structure_model.py b/docling/models/stages/table_structure/table_structure_model.py index 4957777f21..2b52330ce9 100644 --- a/docling/models/stages/table_structure/table_structure_model.py +++ b/docling/models/stages/table_structure/table_structure_model.py @@ -15,6 +15,7 @@ TextCellUnit, ) from PIL import ImageDraw +from typing_extensions import override from docling.datamodel.accelerator_options import AcceleratorDevice, AcceleratorOptions from docling.datamodel.base_models import Page, Table, TableStructurePrediction @@ -96,6 +97,7 @@ def __init__( self.scale = 2.0 # Scale up table input images to 144 dpi @classmethod + @override def get_options_type(cls) -> type[TableStructureOptions]: return TableStructureOptions @@ -178,6 +180,7 @@ def draw_table_and_cells( out_file = out_path / f"table_struct_page_{page.page_no:05}.png" image.save(str(out_file), format="png") + @override def predict_tables( self, conv_res: ConversionResult, diff --git a/docling/models/stages/table_structure/table_structure_model_granite_vision.py b/docling/models/stages/table_structure/table_structure_model_granite_vision.py index f49a39d892..59685ab300 100644 --- a/docling/models/stages/table_structure/table_structure_model_granite_vision.py +++ b/docling/models/stages/table_structure/table_structure_model_granite_vision.py @@ -10,7 +10,13 @@ import torch from docling_core.types.doc import DocItemLabel -from transformers import AutoModelForImageTextToText, AutoProcessor + +# Transformers exposes the model lazily when its torch dependency is installed. +from transformers import ( + AutoModelForImageTextToText, # ty: ignore[possibly-missing-import] + AutoProcessor, +) +from typing_extensions import override from docling.datamodel.accelerator_options import AcceleratorDevice, AcceleratorOptions from docling.datamodel.base_models import Page, Table, TableStructurePrediction @@ -64,6 +70,7 @@ def __init__( self._load_model(artifacts_path) @classmethod + @override def get_options_type(cls) -> type[GraniteVisionTableStructureOptions]: return GraniteVisionTableStructureOptions @@ -119,6 +126,7 @@ def _load_model(self, artifacts_path: Path) -> None: ) self._model.eval() + @override def predict_tables( self, conv_res: ConversionResult, diff --git a/docling/models/stages/table_structure/table_structure_model_v2.py b/docling/models/stages/table_structure/table_structure_model_v2.py index 9553a5f2ed..a2eff3f7ab 100644 --- a/docling/models/stages/table_structure/table_structure_model_v2.py +++ b/docling/models/stages/table_structure/table_structure_model_v2.py @@ -15,7 +15,10 @@ TextCell, ) from PIL import Image, ImageDraw -from transformers import AutoTokenizer + +# Transformers selects tokenizer exports from its optional dependency availability. +from transformers import AutoTokenizer # ty: ignore[possibly-missing-import] +from typing_extensions import override from docling.datamodel.accelerator_options import AcceleratorDevice, AcceleratorOptions from docling.datamodel.base_models import Cluster, Page, Table, TableStructurePrediction @@ -98,6 +101,7 @@ def __init__( self.scale = 2.0 # Scale up table input images to 144 dpi @classmethod + @override def get_options_type(cls) -> type[TableStructureV2Options]: return TableStructureV2Options @@ -171,7 +175,7 @@ def _do_prediction_on_image_to_table( element["text"] = text # Build TableCell objects, assigning text from text_cells by overlap - table_cells = [] + table_cells: list[TableCell] = [] for element in cell_data: tc = TableCell.model_validate(element) table_cells.append(tc) @@ -236,6 +240,7 @@ def _decode_otsl_sequence(self, token_ids: "torch.Tensor") -> list[str]: """ tags = [] skip_tokens = {"", "[UNK]", "", ""} + assert self.tokenizer is not None for tid in token_ids.tolist(): decoded_token = self.tokenizer.decode([tid]) @@ -254,7 +259,7 @@ def _build_table_cells( otsl_seq: list[str], bboxes: "torch.Tensor", table_bbox: list[float], - ) -> tuple[list[dict], int, int]: + ) -> tuple[list[dict[str, object]], int, int]: """ Build table cell structures from OTSL sequence and bboxes. @@ -289,7 +294,7 @@ def _build_table_cells( grid.append(row + [""] * (num_cols - len(row))) # Build cells - table_cells = [] + table_cells: list[dict[str, object]] = [] bbox_idx = 0 t_x1, t_y1, t_x2, t_y2 = table_bbox t_w = t_x2 - t_x1 @@ -397,6 +402,7 @@ def draw_table_and_cells( out_file = out_path / f"table_struct_v2_page_{page.page_no:05}.png" image.save(str(out_file), format="png") + @override def predict_tables( self, conv_res: ConversionResult, diff --git a/docling/models/stages/vlm_convert/vlm_convert_model.py b/docling/models/stages/vlm_convert/vlm_convert_model.py index 2547608e6b..aef9a56222 100644 --- a/docling/models/stages/vlm_convert/vlm_convert_model.py +++ b/docling/models/stages/vlm_convert/vlm_convert_model.py @@ -13,6 +13,7 @@ from pathlib import Path from PIL import Image as PILImage +from typing_extensions import override from docling.datamodel.accelerator_options import AcceleratorOptions from docling.datamodel.base_models import ( @@ -132,7 +133,7 @@ def __init__( def _resolve_runtime_engine_type(self) -> VlmEngineType: selected_engine_type = getattr(self.engine, "selected_engine_type", None) - if selected_engine_type is not None: + if isinstance(selected_engine_type, VlmEngineType): return selected_engine_type return self.options.engine_options.engine_type @@ -253,6 +254,7 @@ def _predict_mineru2(self, images: list[PILImage.Image]) -> list[VlmEngineOutput ) return combined_outputs + @override def __call__( self, conv_res: ConversionResult, page_batch: Iterable[Page] ) -> Iterable[Page]: diff --git a/docling/models/utils/generation_utils.py b/docling/models/utils/generation_utils.py index 0a30e33084..1b1f98d5b9 100644 --- a/docling/models/utils/generation_utils.py +++ b/docling/models/utils/generation_utils.py @@ -9,6 +9,8 @@ from copy import deepcopy from typing import TYPE_CHECKING, Any, List +from typing_extensions import override + if TYPE_CHECKING: from transformers import GenerationConfig @@ -145,6 +147,7 @@ def __init__( self.max_unit = max(1, int(max_unit)) self._lookback_tokens = max(1, int(lookback_tokens)) + @override def lookback_tokens(self) -> int: return self._lookback_tokens @@ -157,6 +160,7 @@ def repeated_unit(self, text: str) -> int | None: max_unit=self.max_unit, ) + @override def should_stop(self, s: str) -> bool: return self.repeated_unit(s) is not None @@ -207,6 +211,7 @@ def _regular(self, vals: List[int]) -> bool: tol = 0.2 * mean return all(abs(d - mean) <= tol for d in diffs) + @override def should_stop(self, s: str) -> bool: """ Trip only on **consecutive** runs (no other matched blocks between) of ≥3 items @@ -219,7 +224,7 @@ def should_stop(self, s: str) -> bool: prev_tag = prev_text = None run = [] # list of (x,y,w,h) - def run_repetitive(boxes: List[tuple]) -> bool: + def run_repetitive(boxes: List[tuple[int, int, int, int]]) -> bool: if len(boxes) < 3: return False # duplicates? diff --git a/docling/models/utils/hf_stopping_criteria.py b/docling/models/utils/hf_stopping_criteria.py index 18eefceb82..217ca1ad0a 100644 --- a/docling/models/utils/hf_stopping_criteria.py +++ b/docling/models/utils/hf_stopping_criteria.py @@ -5,6 +5,7 @@ import torch from transformers import StoppingCriteria +from typing_extensions import override from docling.models.utils.generation_utils import GenerationStopper @@ -28,6 +29,7 @@ def __init__( self.stopper = stopper self.skip_special_tokens = skip_special_tokens + @override def __call__(self, input_ids, scores, **kwargs) -> torch.BoolTensor: """Flag the rows whose decoded tail trips the stopper. @@ -59,4 +61,5 @@ def __call__(self, input_ids, scores, **kwargs) -> torch.BoolTensor: except Exception as e: _log.info(f"Error in TextStopper.should_stop: {e}") continue - return is_done # type: ignore[return-value] + # torch.zeros returns Tensor in the stubs even with dtype=torch.bool. + return is_done # ty: ignore[invalid-return-type] diff --git a/docling/models/vlm_pipeline_models/api_vlm_model.py b/docling/models/vlm_pipeline_models/api_vlm_model.py index a63bd57f9c..a6fa08fd3b 100644 --- a/docling/models/vlm_pipeline_models/api_vlm_model.py +++ b/docling/models/vlm_pipeline_models/api_vlm_model.py @@ -6,14 +6,15 @@ from typing import Union import numpy as np +from numpy.typing import NDArray from PIL.Image import Image +from typing_extensions import override from docling.datamodel.base_models import Page, VlmPrediction, VlmStopReason from docling.datamodel.document import ConversionResult from docling.datamodel.pipeline_options_vlm_model import ApiVlmOptions from docling.exceptions import OperationNotAllowed from docling.models.base_model import BaseVlmPageModel -from docling.models.utils.generation_utils import GenerationStopper from docling.utils.api_image_request import ( api_image_request, api_image_request_streaming, @@ -48,6 +49,7 @@ def __init__( "temperature": self.vlm_options.temperature, } + @override def __call__( self, conv_res: ConversionResult, page_batch: Iterable[Page] ) -> Iterable[Page]: @@ -97,9 +99,10 @@ def __call__( for page in original_order: yield page + @override def process_images( self, - image_batch: Iterable[Union[Image, np.ndarray]], + image_batch: Iterable[Union[Image, NDArray[np.generic]]], prompt: Union[str, list[str]], ) -> Iterable[VlmPrediction]: """Process raw images without page metadata.""" @@ -108,7 +111,7 @@ def process_images( # Handle prompt parameter if isinstance(prompt, str): prompts = [prompt] * len(images) - elif isinstance(prompt, list): + else: if len(prompt) != len(images): raise ValueError( f"Prompt list length ({len(prompt)}) must match image count ({len(images)})" @@ -138,18 +141,17 @@ def _process_single_image(image_prompt_pair): stop_reason = VlmStopReason.UNSPECIFIED error_message = None - if self.vlm_options.custom_stopping_criteria: - # Instantiate any GenerationStopper classes before passing to streaming - instantiated_stoppers = [] - for criteria in self.vlm_options.custom_stopping_criteria: - if isinstance(criteria, GenerationStopper): - instantiated_stoppers.append(criteria) - elif isinstance(criteria, type) and issubclass( - criteria, GenerationStopper - ): - instantiated_stoppers.append(criteria()) - # Skip non-GenerationStopper criteria (should have been caught in validation) + request_params = self.params.copy() + usage_response_key = request_params.pop("usage_response_key", "usage") + token_extract_key = request_params.pop("token_extract_key", None) + if usage_response_key is not None and not isinstance( + usage_response_key, str + ): + raise ValueError("usage_response_key must be a string or None") + if token_extract_key is not None and not isinstance(token_extract_key, str): + raise ValueError("token_extract_key must be a string or None") + if self.vlm_options.custom_stopping_criteria: # Streaming path with early abort support api_response = api_image_request_streaming( image=image, @@ -157,8 +159,10 @@ def _process_single_image(image_prompt_pair): url=self.vlm_options.url, timeout=self.timeout, headers=self.vlm_options.headers, - generation_stoppers=instantiated_stoppers, - **self.params, + generation_stoppers=self.vlm_options.custom_stopping_criteria, + usage_response_key=usage_response_key, + token_extract_key=token_extract_key, + **request_params, ) page_tags = api_response.text num_tokens = api_response.num_tokens @@ -170,7 +174,9 @@ def _process_single_image(image_prompt_pair): url=self.vlm_options.url, timeout=self.timeout, headers=self.vlm_options.headers, - **self.params, + usage_response_key=usage_response_key, + token_extract_key=token_extract_key, + **request_params, ) page_tags = api_response.text num_tokens = api_response.num_tokens diff --git a/docling/models/vlm_pipeline_models/hf_transformers_model.py b/docling/models/vlm_pipeline_models/hf_transformers_model.py index 94dbb358e7..4450f503bd 100644 --- a/docling/models/vlm_pipeline_models/hf_transformers_model.py +++ b/docling/models/vlm_pipeline_models/hf_transformers_model.py @@ -11,9 +11,11 @@ from typing import Any, Optional, Union import numpy as np +from numpy.typing import NDArray from packaging import version from PIL.Image import Image from transformers import StoppingCriteria, StoppingCriteriaList, StopStringCriteria +from typing_extensions import override from docling.datamodel.accelerator_options import ( AcceleratorOptions, @@ -70,13 +72,15 @@ def __init__( if self.enabled: import torch from transformers import ( - AutoModel, - AutoModelForCausalLM, - AutoModelForImageTextToText, AutoProcessor, BitsAndBytesConfig, GenerationConfig, ) + from transformers.models.auto.modeling_auto import ( + AutoModel, + AutoModelForCausalLM, + AutoModelForImageTextToText, + ) transformers_version = importlib.metadata.version("transformers") parsed_transformers_version = version.parse(transformers_version) @@ -172,7 +176,7 @@ def __init__( if is_dots_model: attn_implementation = "sdpa" - self.vlm_model = model_cls.from_pretrained( + model = model_cls.from_pretrained( artifacts_path, device_map=self.device, dtype=self.vlm_options.torch_dtype, @@ -180,8 +184,15 @@ def __init__( trust_remote_code=vlm_options.trust_remote_code, revision=vlm_options.revision, ) + if not isinstance(model, torch.nn.Module): + raise TypeError("HuggingFace VLM must be a PyTorch module") + self.vlm_model = model if sys.version_info < (3, 14): - self.vlm_model = torch.compile(self.vlm_model) # type: ignore + compiled_model = torch.compile(self.vlm_model) + # torch.compile preserves Module attributes on its optimized wrapper. + if not isinstance(compiled_model, torch.nn.Module): + raise TypeError("Compiled VLM must remain a PyTorch module") + self.vlm_model = compiled_model else: self.vlm_model.eval() @@ -190,6 +201,7 @@ def __init__( artifacts_path, revision=vlm_options.revision ) + @override def __call__( self, conv_res: ConversionResult, page_batch: Iterable[Page] ) -> Iterable[Page]: @@ -246,9 +258,10 @@ def __call__( for page in valid_pages: yield page + @override def process_images( self, - image_batch: Iterable[Union[Image, np.ndarray]], + image_batch: Iterable[Union[Image, NDArray[np.generic]]], prompt: Union[str, list[str]], ) -> Iterable[VlmPrediction]: """ @@ -290,6 +303,8 @@ def process_images( ) user_prompts = prompt + prompts: list[str] = [] + # Use your prompt formatter verbatim if self.vlm_options.transformers_prompt_style == TransformersPromptStyle.NONE: inputs = self.processor( @@ -299,7 +314,7 @@ def process_images( **self.vlm_options.extra_processor_kwargs, ) else: - prompts: list[str] = [self.formulate_prompt(p) for p in user_prompts] + prompts = [self.formulate_prompt(p) for p in user_prompts] # -- Processor performs BOTH text+image preprocessing + batch padding (recommended) inputs = self.processor( @@ -338,7 +353,9 @@ def process_images( stopping_criteria_list.append(wrapped_criteria) elif issubclass(criteria, StoppingCriteria): # It's a StoppingCriteria class, instantiate with tokenizer - criteria_instance = criteria(self.processor.tokenizer) + # Custom HF stopper subclasses can accept a tokenizer; + # the base StoppingCriteria constructor declares no arguments. + criteria_instance = criteria(self.processor.tokenizer) # ty: ignore[too-many-positional-arguments] stopping_criteria_list.append(criteria_instance) elif isinstance(criteria, GenerationStopper): # Wrap GenerationStopper instances in HFStoppingCriteriaWrapper @@ -393,8 +410,13 @@ def process_images( gen_kwargs["stopping_criteria"] = stopping_criteria start_time = time.time() + generate = self.vlm_model.generate + if not callable(generate): + raise TypeError("VLM model must provide a generate method") with torch.inference_mode(): - generated_ids = self.vlm_model.generate(**gen_kwargs) + generated_ids = generate(**gen_kwargs) + if not isinstance(generated_ids, torch.Tensor): + raise TypeError("VLM generation must produce a token tensor") generation_time = time.time() - start_time input_len = inputs["input_ids"].shape[1] # common right-aligned prompt length @@ -409,10 +431,11 @@ def process_images( "Neither processor.batch_decode nor tokenizer.batch_decode is available." ) - decoded_texts: list[str] = decode_fn( - trimmed_sequences, - **decoder_config, - ) + decoded_texts: list[str] = [] + for text in decode_fn(trimmed_sequences, **decoder_config): + if not isinstance(text, str): + raise TypeError("VLM decoder must produce strings") + decoded_texts.append(text) # -- Clip off pad tokens from decoded texts pad_token = self.processor.tokenizer.pad_token diff --git a/docling/models/vlm_pipeline_models/mlx_model.py b/docling/models/vlm_pipeline_models/mlx_model.py index 9ec3f9b0d9..edffa9bb62 100644 --- a/docling/models/vlm_pipeline_models/mlx_model.py +++ b/docling/models/vlm_pipeline_models/mlx_model.py @@ -10,8 +10,10 @@ from typing import Union import numpy as np +from numpy.typing import NDArray from PIL.Image import Image from transformers import StoppingCriteria +from typing_extensions import override from docling.datamodel.accelerator_options import ( AcceleratorOptions, @@ -54,9 +56,9 @@ def __init__( if self.enabled: try: - from mlx_vlm import generate, load, stream_generate # type: ignore - from mlx_vlm.prompt_utils import apply_chat_template # type: ignore - from mlx_vlm.utils import load_config # type: ignore + from mlx_vlm import generate, load, stream_generate + from mlx_vlm.prompt_utils import apply_chat_template + from mlx_vlm.utils import load_config except ImportError: raise ImportError( "mlx-vlm is not installed. Please install it via `pip install mlx-vlm` to use MLX VLM models." @@ -95,7 +97,7 @@ def __init__( ) ## Load the model - self.vlm_model, self.processor = load(artifacts_path) + self.vlm_model, self.processor = load(str(artifacts_path)) self.config = load_config(artifacts_path) # Validate custom stopping criteria - MLX doesn't support HF StoppingCriteria @@ -114,6 +116,7 @@ def __init__( f"Found {criteria.__name__}. Use GenerationStopper instead." ) + @override def __call__( self, conv_res: ConversionResult, page_batch: Iterable[Page] ) -> Iterable[Page]: @@ -170,9 +173,10 @@ def __call__( for page in valid_pages: yield page + @override def process_images( self, - image_batch: Iterable[Union[Image, np.ndarray]], + image_batch: Iterable[Union[Image, NDArray[np.generic]]], prompt: Union[str, list[str]], ) -> Iterable[VlmPrediction]: """Process raw images without page metadata. @@ -236,6 +240,9 @@ def process_images( self.processor, self.config, user_prompt, num_images=1 ) + if not isinstance(formatted_prompt, str): + raise TypeError("MLX chat template must produce text") + # Stream generate with stop strings and custom stopping criteria support start_time = time.time() _log.debug("start generating ...") @@ -248,34 +255,47 @@ def process_images( self.vlm_model, self.processor, formatted_prompt, - [image], # MLX stream_generate expects list of images + # mlx-vlm prepare_inputs supports PIL images, but the stream + # signature currently lists only image path strings. + [image], max_tokens=self.max_tokens, verbose=False, temp=self.temperature, ): - # Collect token information - if len(token.logprobs.shape) == 1: - tokens.append( - VlmPredictionToken( - text=token.text, - token=token.token, - logprob=float(token.logprobs[token.token]), + # Newer mlx-vlm can omit token log probabilities and returns + # Python lists; older versions return MLX arrays. + if token.token is not None and token.logprobs is not None: + raw_logprobs: object = token.logprobs + # mlx-vlm annotates lists; supported older releases return + # MLX arrays, including bfloat16 (covered by the stream test). + if isinstance(raw_logprobs, (list, np.ndarray)): + logprobs = np.asarray(raw_logprobs) + else: + # MLX is an optional Apple Silicon dependency. + import mlx.core as mx # ty: ignore[unresolved-import] + + if not isinstance(raw_logprobs, mx.array): + raise TypeError( + "Expected a list or array of token log probabilities" + ) + # Keep MLX scalars in MLX until float() converts them; + # NumPy cannot read the bfloat16 buffer format. + logprobs = raw_logprobs + if logprobs.ndim == 1: + logprob = float(logprobs[token.token]) + elif logprobs.ndim == 2 and logprobs.shape[0] == 1: + logprob = float(logprobs[0, token.token]) + else: + _log.warning( + "incompatible shape for logprobs: %s", logprobs.shape ) - ) - elif ( - len(token.logprobs.shape) == 2 and token.logprobs.shape[0] == 1 - ): - tokens.append( - VlmPredictionToken( - text=token.text, - token=token.token, - logprob=float(token.logprobs[0, token.token]), + logprob = None + if logprob is not None: + tokens.append( + VlmPredictionToken( + text=token.text, token=token.token, logprob=logprob + ) ) - ) - else: - _log.warning( - f"incompatible shape for logprobs: {token.logprobs.shape}" - ) output += token.text @@ -299,6 +319,9 @@ def process_images( ): stopper = criteria() + else: + continue + # Determine the text window to check based on lookback_tokens lookback_tokens = stopper.lookback_tokens() # Check only the last N characters worth of text diff --git a/docling/models/vlm_pipeline_models/vllm_model.py b/docling/models/vlm_pipeline_models/vllm_model.py index c9d6aedc28..19d19cdd58 100644 --- a/docling/models/vlm_pipeline_models/vllm_model.py +++ b/docling/models/vlm_pipeline_models/vllm_model.py @@ -9,7 +9,9 @@ from typing import Any, Dict, Optional, Union import numpy as np +from numpy.typing import NDArray from PIL.Image import Image +from typing_extensions import override from docling.datamodel.accelerator_options import AcceleratorOptions from docling.datamodel.base_models import ( @@ -210,6 +212,7 @@ def __init__( **gen_cfg, ) + @override def __call__( self, conv_res: ConversionResult, page_batch: Iterable[Page] ) -> Iterable[Page]: @@ -269,9 +272,10 @@ def __call__( for page in original_order: yield page + @override def process_images( self, - image_batch: Iterable[Union[Image, np.ndarray]], + image_batch: Iterable[Union[Image, NDArray[np.generic]]], prompt: Union[str, list[str]], ) -> Iterable[VlmPrediction]: """Process images in a single batched vLLM inference call.""" @@ -321,7 +325,7 @@ def process_images( # Generate assert self.llm is not None and self.sampling_params is not None start_time = time.time() - outputs = self.llm.generate(llm_inputs, sampling_params=self.sampling_params) # type: ignore + outputs = self.llm.generate(llm_inputs, sampling_params=self.sampling_params) generation_time = time.time() - start_time # Optional debug diff --git a/docling/pipeline/asr_pipeline.py b/docling/pipeline/asr_pipeline.py index 737421cdcb..67f1212073 100644 --- a/docling/pipeline/asr_pipeline.py +++ b/docling/pipeline/asr_pipeline.py @@ -4,6 +4,7 @@ import logging from docling_core.types.doc import DoclingDocument +from typing_extensions import override from docling.backend.abstract_backend import AbstractDocumentBackend from docling.backend.noop_backend import NoOpBackend @@ -50,6 +51,7 @@ def _has_text(self, document: "DoclingDocument") -> bool: return True return False + @override def _determine_status(self, conv_res: ConversionResult) -> ConversionStatus: if conv_res.status == ConversionStatus.FAILURE or conv_res.errors: return ConversionStatus.FAILURE @@ -62,9 +64,11 @@ def _determine_status(self, conv_res: ConversionResult) -> ConversionStatus: return ConversionStatus.SUCCESS @classmethod + @override def get_default_options(cls) -> AsrPipelineOptions: return AsrPipelineOptions() + @override def _build_document(self, conv_res: ConversionResult) -> ConversionResult: _log.info("Transcribing audio document %s.", conv_res.input.file.name) with TimeRecorder(conv_res, "doc_build", scope=ProfilingScope.DOCUMENT): @@ -72,5 +76,6 @@ def _build_document(self, conv_res: ConversionResult) -> ConversionResult: return conv_res @classmethod + @override def is_backend_supported(cls, backend: AbstractDocumentBackend): return isinstance(backend, NoOpBackend) diff --git a/docling/pipeline/asr_transcriber.py b/docling/pipeline/asr_transcriber.py index d4d7f160da..12529b5e77 100644 --- a/docling/pipeline/asr_transcriber.py +++ b/docling/pipeline/asr_transcriber.py @@ -17,6 +17,7 @@ TrackSource, ) from pydantic import BaseModel, Field +from typing_extensions import override from docling.backend.abstract_backend import AbstractDocumentBackend from docling.backend.noop_backend import NoOpBackend @@ -162,6 +163,7 @@ def __lt__(self, other): return NotImplemented return self.start_time < other.start_time + @override def __eq__(self, other): if not isinstance(other, _ConversationItem): return NotImplemented @@ -207,7 +209,7 @@ def __init__( if self.enabled: try: - import whisper # type: ignore + import whisper except ImportError: if sys.version_info < (3, 14): raise ImportError( @@ -347,7 +349,11 @@ def transcribe(self, fpath: Path) -> list[_ConversationItem]: language=self.language, word_timestamps=self.word_timestamps, beam_size=self.beam_size, - condition_on_previous_text=self.condition_on_previous_text, + condition_on_previous_text=( + True + if self.condition_on_previous_text is None + else self.condition_on_previous_text + ), temperature=self.temperature, ) @@ -387,7 +393,7 @@ def __init__( if self.enabled: try: - import mlx_whisper # type: ignore + import mlx_whisper # ty: ignore[unresolved-import] except ImportError: raise ImportError( "mlx-whisper is not installed. Please install it via " @@ -533,7 +539,7 @@ def __init__( if self.enabled: try: - import whisper_s2t # type: ignore + import whisper_s2t except ImportError: raise ImportError( "whisper_s2t is not installed. Please install it via " @@ -604,7 +610,7 @@ def __init__( self.initial_prompt = asr_options.initial_prompt self.word_timestamps = asr_options.word_timestamps - def _parse_device(self, device_str: str) -> tuple: + def _parse_device(self, device_str: str) -> tuple[str, int]: """Parse device string like 'cuda:0' into ('cuda', 0).""" if ":" in device_str: parts = device_str.split(":") diff --git a/docling/pipeline/base_pipeline.py b/docling/pipeline/base_pipeline.py index 186ead3c7f..c69f80d1ee 100644 --- a/docling/pipeline/base_pipeline.py +++ b/docling/pipeline/base_pipeline.py @@ -11,6 +11,7 @@ from typing import Any, Callable, List, Optional from docling_core.types.doc import ContentLayer, DocItem, DoclingDocument, NodeItem +from typing_extensions import override from docling.backend.abstract_backend import ( AbstractDocumentBackend, @@ -58,7 +59,9 @@ class BasePipeline(ABC): def __init__(self, pipeline_options: PipelineOptions): self.pipeline_options = pipeline_options self.keep_images = False - self.build_pipe: List[Callable] = [] + self.build_pipe: List[ + Callable[[ConversionResult, Iterable[Page]], Iterable[Page]] + ] = [] self.enrichment_pipe: List[GenericEnrichmentModel[Any]] = [] self.artifacts_path: Optional[Path] = None @@ -272,6 +275,7 @@ def _get_picture_description_model( @classmethod @abstractmethod + @override def get_default_options(cls) -> ConvertPipelineOptions: pass @@ -289,6 +293,7 @@ def _apply_on_pages( yield from page_batch + @override def _build_document(self, conv_res: ConversionResult) -> ConversionResult: if not isinstance(conv_res.input._backend, PaginatedDocumentBackend): raise RuntimeError( @@ -388,6 +393,7 @@ def _build_document(self, conv_res: ConversionResult) -> ConversionResult: return conv_res + @override def _unload(self, conv_res: ConversionResult) -> ConversionResult: for page in conv_res.pages: if page._backend is not None: @@ -398,6 +404,7 @@ def _unload(self, conv_res: ConversionResult) -> ConversionResult: return conv_res + @override def _determine_status(self, conv_res: ConversionResult) -> ConversionStatus: status = conv_res.status if status in [ diff --git a/docling/pipeline/extraction_vlm_pipeline.py b/docling/pipeline/extraction_vlm_pipeline.py index 7d204cca6b..417ae7ec81 100644 --- a/docling/pipeline/extraction_vlm_pipeline.py +++ b/docling/pipeline/extraction_vlm_pipeline.py @@ -10,6 +10,7 @@ from PIL.Image import Image from pydantic import BaseModel +from typing_extensions import override from docling.backend.pdf_backend import PdfDocumentBackend, iter_pdf_page_backends from docling.datamodel.base_models import ( @@ -54,6 +55,7 @@ def __init__(self, pipeline_options: VlmExtractionPipelineOptions): prompt_style=pipeline_options.extraction_prompt_style, ) + @override def _extract_data( self, ext_res: ExtractionResult, @@ -158,6 +160,7 @@ def _extract_data( return ext_res + @override def _determine_status(self, ext_res: ExtractionResult) -> ConversionStatus: """Determine the status based on extraction results.""" if ext_res.pages and not any(page.errors for page in ext_res.pages): @@ -232,5 +235,6 @@ class ExtractionTemplateFactory(ModelFactory[template]): # type: ignore raise ValueError(f"Unsupported template type: {type(template)}") @classmethod + @override def get_default_options(cls) -> PipelineOptions: return VlmExtractionPipelineOptions() diff --git a/docling/pipeline/native_pdf_pipeline.py b/docling/pipeline/native_pdf_pipeline.py index 3059c4b804..db174ca88b 100644 --- a/docling/pipeline/native_pdf_pipeline.py +++ b/docling/pipeline/native_pdf_pipeline.py @@ -19,6 +19,7 @@ Size, ) from docling_core.types.doc.page import SegmentedPdfPage +from typing_extensions import override from docling.backend.abstract_backend import AbstractDocumentBackend from docling.backend.pdf_backend import ( @@ -93,6 +94,7 @@ def __init__(self, pipeline_options: NativePdfPipelineOptions) -> None: self.pipeline_options: NativePdfPipelineOptions = pipeline_options self.keep_images = pipeline_options.generate_page_images + @override def _build_document(self, conv_res: ConversionResult) -> ConversionResult: backend = conv_res.input._backend if not isinstance(backend, PdfDocumentBackend): @@ -280,6 +282,7 @@ def _record_page_failure( ) conv_res.status = ConversionStatus.PARTIAL_SUCCESS + @override def _assemble_document(self, conv_res: ConversionResult) -> ConversionResult: timings = _NativeAssemblyTimings() with TimeRecorder(conv_res, "doc_assemble", scope=ProfilingScope.DOCUMENT): @@ -413,11 +416,13 @@ def _to_prov_bbox(bbox: BoundingBox, size: Size) -> BoundingBox: return bbox return bbox.to_bottom_left_origin(page_height=size.height) + @override def _determine_status(self, conv_res: ConversionResult) -> ConversionStatus: if conv_res.status in (ConversionStatus.PENDING, ConversionStatus.STARTED): return ConversionStatus.SUCCESS return conv_res.status + @override def _unload(self, conv_res: ConversionResult) -> None: for page in conv_res.pages: if page._backend is not None: @@ -426,9 +431,11 @@ def _unload(self, conv_res: ConversionResult) -> None: conv_res.input._backend.unload() @classmethod + @override def get_default_options(cls) -> NativePdfPipelineOptions: return NativePdfPipelineOptions() @classmethod + @override def is_backend_supported(cls, backend: AbstractDocumentBackend) -> bool: return isinstance(backend, PdfDocumentBackend) diff --git a/docling/pipeline/simple_pipeline.py b/docling/pipeline/simple_pipeline.py index 0e0c8f7a07..2558b6d554 100644 --- a/docling/pipeline/simple_pipeline.py +++ b/docling/pipeline/simple_pipeline.py @@ -3,6 +3,8 @@ import logging +from typing_extensions import override + from docling.backend.abstract_backend import ( AbstractDocumentBackend, DeclarativeDocumentBackend, @@ -26,6 +28,7 @@ class SimplePipeline(ConvertPipeline): def __init__(self, pipeline_options: ConvertPipelineOptions): super().__init__(pipeline_options) + @override def _build_document(self, conv_res: ConversionResult) -> ConversionResult: if not isinstance(conv_res.input._backend, DeclarativeDocumentBackend): raise RuntimeError( @@ -43,6 +46,7 @@ def _build_document(self, conv_res: ConversionResult) -> ConversionResult: conv_res.document = conv_res.input._backend.convert() return conv_res + @override def _determine_status(self, conv_res: ConversionResult) -> ConversionStatus: # This is called only if the previous steps didn't raise. # Since we don't have anything else to evaluate, we can @@ -50,9 +54,11 @@ def _determine_status(self, conv_res: ConversionResult) -> ConversionStatus: return ConversionStatus.SUCCESS @classmethod + @override def get_default_options(cls) -> ConvertPipelineOptions: return ConvertPipelineOptions() @classmethod + @override def is_backend_supported(cls, backend: AbstractDocumentBackend): return isinstance(backend, DeclarativeDocumentBackend) diff --git a/docling/pipeline/standard_pdf_pipeline.py b/docling/pipeline/standard_pdf_pipeline.py index 664caf1aa5..c1151b121d 100644 --- a/docling/pipeline/standard_pdf_pipeline.py +++ b/docling/pipeline/standard_pdf_pipeline.py @@ -37,6 +37,7 @@ Size, TableItem, ) +from typing_extensions import override from docling.backend.abstract_backend import AbstractDocumentBackend from docling.backend.pdf_backend import PdfDocumentBackend, iter_pdf_page_backends @@ -52,6 +53,7 @@ from docling.datamodel.pipeline_options import ( KserveV2OcrOptions, LayoutPostprocessorOptions, + PdfPipelineOptions, ThreadedPdfPipelineOptions, ) from docling.datamodel.settings import settings @@ -180,7 +182,8 @@ def put(self, item: ThreadedItem, timeout: float | None = None) -> bool: if self._closed: return False start = time.monotonic() - while len(self._items) >= self._max and not self._closed: + # Condition.wait() releases the lock; another thread can close the queue. + while len(self._items) >= self._max and not self._closed: # ty: ignore[redundant-condition-strict] if timeout is not None: remaining = timeout - (time.monotonic() - start) if remaining <= 0: @@ -188,7 +191,8 @@ def put(self, item: ThreadedItem, timeout: float | None = None) -> bool: self._not_full.wait(remaining) else: self._not_full.wait() - if self._closed: + # Recheck after wait(), which permits another thread to close the queue. + if self._closed: # ty: ignore[redundant-condition-strict] return False self._items.append(item) self._not_empty.notify() @@ -355,19 +359,21 @@ def _process_batch(self, batch: Sequence[ThreadedItem]) -> list[ThreadedItem]: continue pages: list[Page] = [payload for _, payload in pages_with_payloads] - if _log.isEnabledFor(logging.DEBUG): - _t_start = time.time() - _t_mono = time.monotonic() + timing = ( + (time.time(), time.monotonic()) + if _log.isEnabledFor(logging.DEBUG) + else None + ) processed_pages = list(self.model(good[0].conv_res, pages)) # type: ignore[arg-type] - if _log.isEnabledFor(logging.DEBUG): + if timing is not None: _log.debug( "PIPELINE_PROFILING Stage %s: run_id=%d pages=%s start=%.3f end=%.3f duration=%.3fs", self.name, rid, [it.page_no for it in good], - _t_start, + timing[0], time.time(), - time.monotonic() - _t_mono, + time.monotonic() - timing[1], ) if len(processed_pages) != len(pages): # strict mismatch guard raise RuntimeError( @@ -433,6 +439,7 @@ def __init__( timed_out_run_ids=timed_out_run_ids, ) + @override def _process_batch(self, batch: Sequence[ThreadedItem]) -> list[ThreadedItem]: groups: dict[int, list[ThreadedItem]] = defaultdict(list) for itm in batch: @@ -517,21 +524,23 @@ def _process_batch(self, batch: Sequence[ThreadedItem]) -> list[ThreadedItem]: continue try: - if _log.isEnabledFor(logging.DEBUG): - _t_start = time.time() - _t_mono = time.monotonic() + timing = ( + (time.time(), time.monotonic()) + if _log.isEnabledFor(logging.DEBUG) + else None + ) pages = [page for _, page in valid] processed_pages = list( self.model(valid[0][0].conv_res, pages) # type: ignore[arg-type] ) - if _log.isEnabledFor(logging.DEBUG): + if timing is not None: _log.debug( "PIPELINE_PROFILING Stage preprocess: run_id=%d pages=%s start=%.3f end=%.3f duration=%.3fs", rid, [it.page_no for it, _ in valid], - _t_start, + timing[0], time.time(), - time.monotonic() - _t_mono, + time.monotonic() - timing[1], ) if len(processed_pages) != len(pages): raise RuntimeError( @@ -586,9 +595,9 @@ class RunContext: class StandardPdfPipeline(ConvertPipeline): """High-performance PDF pipeline with multi-threaded stages.""" - def __init__(self, pipeline_options: ThreadedPdfPipelineOptions) -> None: + def __init__(self, pipeline_options: PdfPipelineOptions) -> None: super().__init__(pipeline_options) - self.pipeline_options: ThreadedPdfPipelineOptions = pipeline_options + self.pipeline_options: PdfPipelineOptions = pipeline_options self._run_seq = itertools.count(1) # deterministic, monotonic run ids self._page_sizes_by_no: dict[int, Size] = {} @@ -712,7 +721,7 @@ def _make_ocr_model(self, art_path: Path | None) -> Any: **extra, ) - def _release_page_resources(self, item: ThreadedItem) -> None: + def _release_threaded_item_resources(self, item: ThreadedItem) -> None: page = item.payload if page is None: return @@ -781,7 +790,7 @@ def _create_run_ctx(self) -> RunContext: batch_timeout=opts.batch_polling_interval_seconds, queue_max_size=opts.queue_max_size, shutdown_timeout=opts.stage_shutdown_timeout_seconds, - postprocess=self._release_page_resources, + postprocess=self._release_threaded_item_resources, timed_out_run_ids=timed_out_run_ids, ) @@ -806,6 +815,7 @@ def _create_run_ctx(self) -> RunContext: def _get_expected_page_nos(self, conv_res: ConversionResult) -> list[int]: return get_expected_page_nos(conv_res) + @override def _build_document(self, conv_res: ConversionResult) -> ConversionResult: """Stream-build the document with a dedicated producer thread. @@ -895,10 +905,7 @@ def _produce_pages() -> None: try: while proc.success_count + proc.failure_count < total_pages: # Check timeout - if ( - self.pipeline_options.document_timeout is not None - and not timeout_exceeded - ): + if self.pipeline_options.document_timeout is not None: elapsed_time = time.monotonic() - start_time if elapsed_time > self.pipeline_options.document_timeout: _log.warning( @@ -1045,6 +1052,7 @@ def _integrate_results( p.parsed_page = None # ---------------------------------------------------------------- assemble + @override def _assemble_document(self, conv_res: ConversionResult) -> ConversionResult: elements, headers, body = [], [], [] with TimeRecorder(conv_res, "doc_assemble", scope=ProfilingScope.DOCUMENT): @@ -1089,7 +1097,7 @@ def _assemble_document(self, conv_res: ConversionResult) -> ConversionResult: page_ix = element.prov[0].page_no page = next( (p for p in conv_res.pages if p.page_no == page_ix), - cast("Page", None), + None, ) assert page is not None assert page.size is not None @@ -1158,9 +1166,6 @@ def _add_failed_pages_to_document( The failed pages are added with their size information (if available from the backend) but without any content. """ - if conv_res.document is None: - return - # Find pages that are missing from the document existing_page_nos = set(conv_res.document.pages.keys()) missing_page_nos = set(expected_page_nos) - existing_page_nos @@ -1186,16 +1191,20 @@ def _add_failed_pages_to_document( # ---------------------------------------------------------------- misc @classmethod + @override def get_default_options(cls) -> ThreadedPdfPipelineOptions: return ThreadedPdfPipelineOptions() @classmethod + @override def is_backend_supported(cls, backend: AbstractDocumentBackend) -> bool: return isinstance(backend, PdfDocumentBackend) + @override def _determine_status(self, conv_res: ConversionResult) -> ConversionStatus: return conv_res.status + @override def _unload(self, conv_res: ConversionResult) -> None: self._page_sizes_by_no = {} for p in conv_res.pages: diff --git a/docling/pipeline/video_pipeline.py b/docling/pipeline/video_pipeline.py index 69a3d99725..a93404bb8d 100644 --- a/docling/pipeline/video_pipeline.py +++ b/docling/pipeline/video_pipeline.py @@ -26,6 +26,7 @@ ImageRef, TrackSource, ) +from typing_extensions import override from docling.backend.abstract_backend import AbstractDocumentBackend from docling.backend.noop_backend import NoOpBackend @@ -39,6 +40,7 @@ from docling.datamodel.pipeline_options import VideoPipelineOptions from docling.pipeline.asr_transcriber import ( _AsrModelFactory, + _ConversationItem, _merge_into_sentences, ) from docling.pipeline.base_pipeline import BasePipeline @@ -118,13 +120,16 @@ def __init__(self, pipeline_options: VideoPipelineOptions): ) @classmethod + @override def get_default_options(cls) -> VideoPipelineOptions: return VideoPipelineOptions() @classmethod + @override def is_backend_supported(cls, backend: AbstractDocumentBackend) -> bool: return isinstance(backend, NoOpBackend) + @override def _determine_status(self, conv_res: ConversionResult) -> ConversionStatus: if conv_res.status == ConversionStatus.FAILURE: return ConversionStatus.FAILURE @@ -146,6 +151,7 @@ def _determine_status(self, conv_res: ConversionResult) -> ConversionStatus: return ConversionStatus.PARTIAL_SUCCESS return ConversionStatus.SUCCESS + @override def _build_document(self, conv_res: ConversionResult) -> ConversionResult: _log.info("Processing video document %s.", conv_res.input.file.name) with TimeRecorder(conv_res, "doc_build", scope=ProfilingScope.DOCUMENT): @@ -292,7 +298,7 @@ def _process_video(self, conv_res: ConversionResult) -> None: ) # 6. Merge transcript + frames by timestamp - events: list[tuple[float, int, object]] = [] + events: list[tuple[float, int, _ConversationItem | VideoFrame]] = [] for item in transcript_items: if ( item.start_time is not None @@ -305,9 +311,9 @@ def _process_video(self, conv_res: ConversionResult) -> None: events.sort(key=lambda e: (e[0], e[1])) - for _, event_type, payload in events: - if event_type == 0: - frame = payload # type: ignore[assignment] + for _, _, payload in events: + if isinstance(payload, VideoFrame): + frame = payload try: conv_res.document.add_picture( image=ImageRef.from_pil(frame.image, dpi=72), @@ -329,7 +335,8 @@ def _process_video(self, conv_res: ConversionResult) -> None: exc, ) else: - item = payload # type: ignore[assignment] + item = payload + assert item.start_time is not None and item.end_time is not None try: conv_res.document.add_text( label=DocItemLabel.TEXT, diff --git a/docling/pipeline/vlm_pipeline.py b/docling/pipeline/vlm_pipeline.py index ff67a579fa..b9c8abfe14 100644 --- a/docling/pipeline/vlm_pipeline.py +++ b/docling/pipeline/vlm_pipeline.py @@ -84,9 +84,7 @@ def __init__(self, pipeline_options: VlmPipelineOptions): else: self._initialize_legacy_vlm_models(pipeline_options) - self.enrichment_pipe: list = [ - # Other models working on `NodeItem` elements in the DoclingDocument - ] + self.enrichment_pipe = [] def _initialize_new_runtime_system( self, pipeline_options: VlmPipelineOptions @@ -134,6 +132,8 @@ def _initialize_legacy_vlm_models( Note: This method is deprecated and will be removed in a future version. """ + vlm_options = pipeline_options.vlm_options + assert isinstance(vlm_options, (InlineVlmOptions, ApiVlmOptions)) # Legacy path - using old InlineVlmOptions or ApiVlmOptions warnings.warn( "Using legacy VLM options (InlineVlmOptions/ApiVlmOptions) is deprecated. " @@ -147,21 +147,20 @@ def _initialize_legacy_vlm_models( # force_backend_text = True - get text from backend using bounding boxes predicted by SmolDocling doctags self.force_backend_text = ( pipeline_options.force_backend_text - and pipeline_options.vlm_options.response_format == ResponseFormat.DOCTAGS # type: ignore[union-attr] + and vlm_options.response_format == ResponseFormat.DOCTAGS ) self.keep_images = self.pipeline_options.generate_page_images - if isinstance(pipeline_options.vlm_options, ApiVlmOptions): + if isinstance(vlm_options, ApiVlmOptions): self.build_pipe = [ ApiVlmModel( enabled=True, enable_remote_services=self.pipeline_options.enable_remote_services, - vlm_options=cast(ApiVlmOptions, self.pipeline_options.vlm_options), + vlm_options=vlm_options, ), ] - elif isinstance(self.pipeline_options.vlm_options, InlineVlmOptions): - vlm_options = cast(InlineVlmOptions, self.pipeline_options.vlm_options) + else: if vlm_options.inference_framework == InferenceFramework.MLX: self.build_pipe = [ HuggingFaceMlxModel( @@ -196,14 +195,13 @@ def _initialize_legacy_vlm_models( f"Could not instantiate the right type of VLM pipeline: {vlm_options.inference_framework}" ) + @override def initialize_page(self, conv_res: ConversionResult, page: Page) -> Page: raise NotImplementedError("VlmPipeline initializes pages in _build_document()") def _initialize_page(self, conv_res: ConversionResult, page: Page) -> Page: with TimeRecorder(conv_res, "page_init"): - images_scale = self.pipeline_options.images_scale - if images_scale is not None: - page._default_image_scale = images_scale + page._default_image_scale = self.pipeline_options.images_scale if page._backend is not None and page._backend.is_valid(): page.size = page._backend.get_size() @@ -212,6 +210,7 @@ def _initialize_page(self, conv_res: ConversionResult, page: Page) -> Page: return page + @override def _build_document(self, conv_res: ConversionResult) -> ConversionResult: if not isinstance(conv_res.input._backend, PdfDocumentBackend): raise RuntimeError( @@ -441,6 +440,7 @@ def _determine_status(self, conv_res: ConversionResult) -> ConversionStatus: status = ConversionStatus.PARTIAL_SUCCESS return status + @override def _assemble_document(self, conv_res: ConversionResult) -> ConversionResult: return conv_res @@ -705,7 +705,9 @@ def _extract_code_block(self, text: str) -> str: if mtch: # Return only the content of the first capturing group - return mtch.group(1) + code = mtch.group(1) + assert isinstance(code, str) + return code else: # No code blocks found, return original text return text @@ -804,9 +806,11 @@ def _finalize_page_output(self, document: DoclingDocument, page: Page) -> None: ) @classmethod + @override def get_default_options(cls) -> VlmPipelineOptions: return VlmPipelineOptions() @classmethod + @override def is_backend_supported(cls, backend: AbstractDocumentBackend): return isinstance(backend, PdfDocumentBackend) diff --git a/docling/service_client/_async_client.py b/docling/service_client/_async_client.py index 0c9a3eab61..77098a9f0c 100644 --- a/docling/service_client/_async_client.py +++ b/docling/service_client/_async_client.py @@ -215,7 +215,10 @@ async def submit( headers: dict[str, str] | None = None, *, target: PresignedUrlTarget | None = None, - ) -> AsyncConversionJob[PresignedUrlConvertResponse | ConversionResult]: ... + ) -> ( + AsyncConversionJob[PresignedUrlConvertResponse] + | AsyncConversionJob[ConversionResult] + ): ... async def submit( self, @@ -317,7 +320,8 @@ async def submit_batch( AsyncConversionJob[PresignedUrlConvertDocumentResponse] | AsyncConversionJob[PresignedUrlConvertResponse] ): - assert self._async_client is not None, "client not open — use async with" + async_client = self._async_client + assert async_client is not None, "client not open — use async with" if target is None and targets is None: raise ValueError("submit_batch() requires either 'target' or 'targets'.") if target is not None and targets is not None: @@ -348,7 +352,7 @@ async def submit_batch( options=submit_options, target=request.target, targets=request.targets, - async_client=self._async_client, + async_client=async_client, request_headers=headers, ) @@ -366,7 +370,7 @@ async def fetch_result( return await self._fetch_presigned_document_result( task_id=task_id, last_status=last_status, - async_client=self._async_client, + async_client=async_client, ) else: @@ -378,7 +382,7 @@ async def fetch_result( return await self._fetch_presigned_result( task_id=task_id, last_status=last_status, - async_client=self._async_client, + async_client=async_client, ) handlers = _AsyncJobHandlers[Any]( diff --git a/docling/service_client/client.py b/docling/service_client/client.py index 37f3ebdd1d..d2517adf35 100644 --- a/docling/service_client/client.py +++ b/docling/service_client/client.py @@ -24,7 +24,7 @@ from enum import Enum from io import BytesIO from pathlib import Path, PurePath -from typing import IO, TYPE_CHECKING, Any, TypeAlias, TypeVar +from typing import IO, TYPE_CHECKING, Any, TypeAlias, TypeVar, overload from urllib.parse import urlencode, urlparse import httpx @@ -32,7 +32,14 @@ from docling_core.types.doc.common.constants import CURRENT_VERSION from docling_core.types.io import DocumentStream from PIL import Image as PILImage -from pydantic import AnyHttpUrl, SecretBytes, SecretStr, TypeAdapter, ValidationError +from pydantic import ( + AnyHttpUrl, + BaseModel, + SecretBytes, + SecretStr, + TypeAdapter, + ValidationError, +) from docling.backend.noop_backend import NoOpBackend from docling.datamodel.base_models import ( @@ -118,6 +125,7 @@ BatchSubmitTarget: TypeAlias = BatchTargetRequestInput logger = logging.getLogger(__name__) _T = TypeVar("_T") +_ResultModelT = TypeVar("_ResultModelT", bound=BaseModel) SUCCESS_CONVERSION_STATUSES: set[ConversionStatus] = { @@ -276,8 +284,8 @@ def __init__( def _parse_result_model_response( self, response: httpx.Response, - model_cls: type[_T], - ) -> _T: + model_cls: type[_ResultModelT], + ) -> _ResultModelT: try: return model_cls.model_validate_json(response.text) except (ValidationError, ValueError) as exc: @@ -381,7 +389,7 @@ def _options_for_output_formats( self, options: ConvertDocumentsRequestOptions, output_formats: list[OutputFormat] | None, - target: SubmitTarget | GenericTargetRequest, + target: SubmitTarget | GenericTargetRequest | None, ) -> ConvertDocumentsRequestOptions: effective = options if output_formats is not None: @@ -485,7 +493,9 @@ def _filename_from_headers(self, headers: httpx.Headers) -> str | None: match = re.search(r'filename="?(?P[^";]+)"?', disposition) if match is None: return None - return match.group("name") + filename = match.group("name") + assert isinstance(filename, str) + return filename def _describe_source(self, source: SourceType) -> _SourceDescriptor: source = self._normalize_source(source) @@ -619,7 +629,7 @@ def _retry_with_retry_after_header( ) def _exponential_backoff_delay(self, attempt: int) -> float: - return HTTP_RETRY_BACKOFF_BASE_SECONDS * (2**attempt) + return HTTP_RETRY_BACKOFF_BASE_SECONDS * (2.0**attempt) def _transport_retry_delay( self, @@ -1188,6 +1198,49 @@ def _convert_single( raise TypeError("Conversion submission returned an unexpected result type.") return result + @overload + def _submit_conversion_job( + self, + source: SourceType, + options: ConvertDocumentsRequestOptions, + limits: DocumentLimits, + target: InBodyTarget, + descriptor: _SourceDescriptor | None = None, + request_headers: dict[str, str] | None = None, + materialize_presigned: bool = False, + ) -> ConversionJob[ConversionResult]: ... + + @overload + def _submit_conversion_job( + self, + source: SourceType, + options: ConvertDocumentsRequestOptions, + limits: DocumentLimits, + target: PresignedUrlTarget, + descriptor: _SourceDescriptor | None = None, + request_headers: dict[str, str] | None = None, + materialize_presigned: bool = False, + ) -> ( + ConversionJob[ConversionResult] | ConversionJob[PresignedUrlConvertResponse] + ): ... + + @overload + def _submit_conversion_job( + self, + source: SourceType, + options: ConvertDocumentsRequestOptions, + limits: DocumentLimits, + target: SubmitTarget, + descriptor: _SourceDescriptor | None = None, + request_headers: dict[str, str] | None = None, + materialize_presigned: bool = False, + ) -> ( + ConversionJob[ConversionResult] + | ConversionJob[RawServiceResult] + | ConversionJob[PresignedUrlConvertDocumentResponse] + | ConversionJob[PresignedUrlConvertResponse] + ): ... + def _submit_conversion_job( self, source: SourceType, @@ -2241,6 +2294,7 @@ def make_items() -> Iterator[ConversionItem]: limits=resolved.limits, ) else: + assert isinstance(outcome, ConvertDocumentResponse) result = self._build_conversion_result( payload=outcome, descriptor=metadata.descriptor, diff --git a/docling/service_client/exceptions.py b/docling/service_client/exceptions.py index 4c672cbf88..6353cbb45b 100644 --- a/docling/service_client/exceptions.py +++ b/docling/service_client/exceptions.py @@ -7,6 +7,8 @@ from dataclasses import dataclass +from typing_extensions import override + from docling.datamodel.service.responses import PublicFailureInfo @@ -22,6 +24,7 @@ class ServiceError(DoclingServiceClientError): status_code: int | None = None detail: str | None = None + @override def __str__(self) -> str: if self.status_code is None: return self.message @@ -54,6 +57,7 @@ def __init__( self.current_usage = current_usage self.limit = limit + @override def __str__(self) -> str: base = super().__str__() if self.current_usage is None or self.limit is None: @@ -110,5 +114,6 @@ class BatchConversionError(DoclingServiceClientError): message: str failures: list[Exception] + @override def __str__(self) -> str: return f"{self.message} ({len(self.failures)} failure(s))" diff --git a/docling/utils/chandra_utils.py b/docling/utils/chandra_utils.py index 9381ee7d39..1e51684156 100644 --- a/docling/utils/chandra_utils.py +++ b/docling/utils/chandra_utils.py @@ -38,6 +38,7 @@ ) from PIL import Image as PILImage from pydantic import AnyUrl, ValidationError +from typing_extensions import override from docling.utils.code_language import detect_code_language @@ -144,6 +145,7 @@ def __init__(self, content: str): ) self.close() + @override def handle_starttag(self, tag: str, attrs: list[tuple[str, str | None]]) -> None: # Accept omitted closing tags for adjacent cells, rows, and list items. closes = { @@ -166,17 +168,20 @@ def handle_starttag(self, tag: str, attrs: list[tuple[str, str | None]]) -> None if tag not in _VOID_TAGS: self.stack.append(node) + @override def handle_startendtag(self, tag: str, attrs: list[tuple[str, str | None]]) -> None: self.handle_starttag(tag, attrs) if tag not in _VOID_TAGS: self.handle_endtag(tag) + @override def handle_endtag(self, tag: str) -> None: for index in range(len(self.stack) - 1, 0, -1): if self.stack[index].tag == tag: del self.stack[index:] break + @override def handle_data(self, data: str) -> None: self.stack[-1].children.append(data) @@ -428,13 +433,15 @@ def _emit_runs( label = DocItemLabel.TEXT if len(merged) > 1: if label != DocItemLabel.TEXT: - parent = self.doc.add_text( + text_parent = self.doc.add_text( label=label, text="", parent=parent, prov=self._text_prov(prov, ""), content_layer=layer, ) + assert isinstance(text_parent, NodeItem) + parent = text_parent parent = self.doc.add_inline_group(parent=parent, content_layer=layer) for run in merged: run_label = ( diff --git a/docling/utils/layout_postprocessor.py b/docling/utils/layout_postprocessor.py index 1735885a40..7c5fe64b9b 100644 --- a/docling/utils/layout_postprocessor.py +++ b/docling/utils/layout_postprocessor.py @@ -46,7 +46,7 @@ def union(self, x, y): def get_groups(self) -> dict[int, list[int]]: """Returns groups as {root: [elements]}.""" - groups = defaultdict(list) + groups = defaultdict[int, list[int]](list) for elem in self.parent: groups[self.find(elem)].append(elem) return groups @@ -480,7 +480,7 @@ def _handle_cross_type_overlaps(self, special_clusters) -> list[Cluster]: return [c for c in special_clusters if c.id not in clusters_to_remove] def _should_prefer_cluster( - self, candidate: Cluster, other: Cluster, params: dict + self, candidate: Cluster, other: Cluster, params: dict[str, float] ) -> bool: """Determine if candidate cluster should be preferred over other cluster based on rules. Returns True if candidate should be preferred, False if not.""" @@ -518,7 +518,7 @@ def _should_prefer_cluster( def _select_best_cluster_from_group( self, group_clusters: list[Cluster], - params: dict, + params: dict[str, float], ) -> Cluster: """Select best cluster from a group of overlapping clusters based on all rules.""" current_best = None diff --git a/docling/utils/mineru_utils.py b/docling/utils/mineru_utils.py index 5b007543a7..b7f0bcf292 100644 --- a/docling/utils/mineru_utils.py +++ b/docling/utils/mineru_utils.py @@ -439,6 +439,7 @@ def parse_mineru2( prov=provenance, ) if region.type == "text": + assert isinstance(text_item, TextItem) previous_text_item = text_item link_adjacent_captions(document, block_starts, caption_owners) return document diff --git a/docling/utils/model_downloader.py b/docling/utils/model_downloader.py index aeb250e794..3f6f15372c 100644 --- a/docling/utils/model_downloader.py +++ b/docling/utils/model_downloader.py @@ -77,7 +77,7 @@ def download_models( with_easyocr: bool = False, easyocr_languages: Optional[list[str]] = None, # BCP-47 tags with_nemotron_ocr: bool = False, -): +) -> Path: if easyocr_languages is not None and not with_easyocr: raise ValueError("easyocr_languages requires with_easyocr=True") if rapidocr_models is not None and not with_rapidocr: diff --git a/docling/utils/nemotron_parse_utils.py b/docling/utils/nemotron_parse_utils.py index 134f794201..7a4dd032b2 100644 --- a/docling/utils/nemotron_parse_utils.py +++ b/docling/utils/nemotron_parse_utils.py @@ -26,6 +26,7 @@ TabularChartMetaField, ) from PIL import Image as PILImage +from typing_extensions import override from docling.backend.latex_backend import LatexDocumentBackend from docling.backend.md_backend import MarkdownDocumentBackend @@ -72,14 +73,17 @@ def __init__(self) -> None: super().__init__(convert_charrefs=True) self._parts: list[str] = [] + @override def handle_starttag(self, tag: str, attrs: list[tuple[str, str | None]]) -> None: if tag.lower() == "br": self._append_break() + @override def handle_endtag(self, tag: str) -> None: if tag.lower() in self._BLOCK_TAGS: self._append_break() + @override def handle_data(self, data: str) -> None: self._parts.append(re.sub(r"\s+", " ", data)) diff --git a/docling/utils/ocr_language.py b/docling/utils/ocr_language.py index 003dd9b1e0..dc61094055 100644 --- a/docling/utils/ocr_language.py +++ b/docling/utils/ocr_language.py @@ -22,6 +22,7 @@ import langcodes from pydantic import BaseModel, ConfigDict +from typing_extensions import override _log = logging.getLogger(__name__) @@ -124,6 +125,7 @@ def has_default_script(self) -> bool: default_script = None return self.bcp47_script == default_script + @override def __str__(self) -> str: return self.tag() diff --git a/docling/utils/otsl.py b/docling/utils/otsl.py index e012b141fb..9833f14588 100644 --- a/docling/utils/otsl.py +++ b/docling/utils/otsl.py @@ -27,7 +27,9 @@ def parse_otsl_output(text: str) -> tuple[list[str], list[TableCell], int, int]: otsl_match = re.search(r"(.*)", text, re.DOTALL) if otsl_match: - text = otsl_match.group(1) + content = otsl_match.group(1) + assert isinstance(content, str) + text = content token_pairs: list[tuple[str, str]] = [] for match in _TAG_PATTERN.finditer(text): diff --git a/docling/utils/paddleocr_vl_utils.py b/docling/utils/paddleocr_vl_utils.py index 6b3d05241e..24d19986f5 100644 --- a/docling/utils/paddleocr_vl_utils.py +++ b/docling/utils/paddleocr_vl_utils.py @@ -142,7 +142,11 @@ class _PaddleOCRVLPage(BaseModel): def _strip_formatted_centering(text: str) -> str: """Undo Paddle's optional centered ``div`` around textual captions.""" match = _CENTERED_DIV_PATTERN.match(text) - return match.group(1) if match is not None else text + if match is None: + return text + content = match.group(1) + assert isinstance(content, str) + return content def _load_page_payload( diff --git a/docling/utils/pdf_outline.py b/docling/utils/pdf_outline.py index 24d45bcf18..62f64e01eb 100644 --- a/docling/utils/pdf_outline.py +++ b/docling/utils/pdf_outline.py @@ -89,6 +89,8 @@ def _dest_top_pdf(dest: pdfium.PdfDest) -> tuple[int | None, float | None]: mode, pos = dest.get_view() idx = _view_top_index().get(mode) y_pdf = pos[idx] if idx is not None and idx < len(pos) else None + assert page_index is None or isinstance(page_index, int) + assert y_pdf is None or isinstance(y_pdf, (int, float)) return page_index, y_pdf diff --git a/docling/utils/profiling.py b/docling/utils/profiling.py index b2c7b0f047..f5bb967db7 100644 --- a/docling/utils/profiling.py +++ b/docling/utils/profiling.py @@ -27,19 +27,19 @@ class ProfilingItem(BaseModel): start_timestamps: List[datetime] = [] def total(self) -> float: - return np.sum(self.times) # type: ignore + return float(np.sum(self.times)) def avg(self) -> float: - return np.average(self.times) # type: ignore + return float(np.average(self.times)) def std(self) -> float: - return np.std(self.times) # type: ignore + return float(np.std(self.times)) def mean(self) -> float: - return np.mean(self.times) # type: ignore + return float(np.mean(self.times)) def percentile(self, perc: float) -> float: - return np.percentile(self.times, perc) # type: ignore + return float(np.percentile(self.times, perc)) class TimeRecorder: diff --git a/docling/utils/speaker_diarization.py b/docling/utils/speaker_diarization.py index 5844a50b8e..b0e36cd10d 100644 --- a/docling/utils/speaker_diarization.py +++ b/docling/utils/speaker_diarization.py @@ -16,6 +16,7 @@ from typing import Any import numpy as np +from numpy.typing import NDArray _log = logging.getLogger(__name__) @@ -45,7 +46,7 @@ class DiarizationResult: speaker_ids: list[str] = field(default_factory=list) -def _estimate_num_speakers(embeddings: np.ndarray) -> int: +def _estimate_num_speakers(embeddings: NDArray[np.float32]) -> int: """Estimate optimal speaker count via silhouette score. Args: @@ -137,7 +138,7 @@ def diarize( step_samples = int(_WINDOW_STEP * sr) timestamps: list[float] = [] - wav_splits: list[np.ndarray] = [] + wav_splits: list[NDArray[np.float32]] = [] i = 0 while i + window_samples <= len(wav): diff --git a/docling/utils/utils.py b/docling/utils/utils.py index 5053dcf4e8..ae114f82d0 100644 --- a/docling/utils/utils.py +++ b/docling/utils/utils.py @@ -34,7 +34,7 @@ def _hash_buf(binary_stream): if isinstance(path_or_stream, Path): with path_or_stream.open("rb") as afile: _hash_buf(afile) - elif isinstance(path_or_stream, BytesIO): + else: _hash_buf(path_or_stream) return hasher.hexdigest() diff --git a/docling/utils/video_frame_sampling.py b/docling/utils/video_frame_sampling.py index bf6d638921..0b87ec1ebe 100644 --- a/docling/utils/video_frame_sampling.py +++ b/docling/utils/video_frame_sampling.py @@ -28,6 +28,7 @@ from typing import Final import numpy as np +from numpy.typing import NDArray from PIL import Image from pydantic import BaseModel, ConfigDict, Field @@ -282,7 +283,7 @@ def _require_ffmpeg() -> None: Higher = stricter (fewer scenes on busy video); lower = more sensitive.""" -def _auto_prominence(diffs: np.ndarray) -> float: +def _auto_prominence(diffs: NDArray[np.float64]) -> float: """Adapt the scene-cut threshold to how busy the video is. Uses the median frame difference as the ambient-motion floor and the @@ -527,7 +528,7 @@ def __init__( self.sharpness_candidates = sharpness_candidates self.runner = runner if runner is not None else FfmpegRunner() - def _probe_diffs(self, video_path: Path) -> tuple[list[float], np.ndarray]: + def _probe_diffs(self, video_path: Path) -> tuple[list[float], NDArray[np.float64]]: """Decode probe thumbnails at ``probe_fps`` and diff consecutive frames. Returns: @@ -549,7 +550,7 @@ def _on_frame(image: Image.Image) -> None: video_path, self.probe_fps, self.probe_size, self.runner, _on_frame ) timestamps = [i / self.probe_fps for i in range(count)] - return timestamps, np.array(diffs) + return timestamps, np.array(diffs, dtype=np.float64) @staticmethod def _mean_abs_diff(a: Image.Image, b: Image.Image) -> float: diff --git a/docling/utils/vlm_utils.py b/docling/utils/vlm_utils.py index d3ed9c5d11..3e2e9ce1a4 100644 --- a/docling/utils/vlm_utils.py +++ b/docling/utils/vlm_utils.py @@ -68,7 +68,9 @@ def parse_markdown_heading(text: str) -> tuple[str, int | None]: match = _MARKDOWN_HEADING_PATTERN.match(text) if match is None: return text, None - return match.group(2), len(match.group(1)) + heading = match.group(2) + assert isinstance(heading, str) + return heading, len(match.group(1)) def strip_stop_strings(texts: list[str], stop_strings: list[str]) -> list[str]: diff --git a/pyproject.toml b/pyproject.toml index c1f7868ed6..b27ab84db6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -321,7 +321,7 @@ all = [ [dependency-groups] typecheck = [ - "ty==0.0.33", + "ty==0.0.84", "types-setuptools~=70.3", "pandas-stubs~=2.1", "types-openpyxl~=3.1", diff --git a/tests/test_asr_pipeline.py b/tests/test_asr_pipeline.py index 374334dde1..4039db0aac 100644 --- a/tests/test_asr_pipeline.py +++ b/tests/test_asr_pipeline.py @@ -208,7 +208,12 @@ def test_native_and_mlx_transcribe_language_handling(monkeypatch, tmp_path): mm.mlx_whisper.transcribe.assert_called() -def test_native_whisper_passes_decode_options_to_transcribe(tmp_path): +@pytest.mark.parametrize( + "condition, expected", [(False, False), (True, True), (None, True)] +) +def test_native_whisper_passes_decode_options_to_transcribe( + tmp_path, condition, expected +): """Native Whisper forwards decoding options to whisper.transcribe.""" from docling.pipeline.asr_pipeline import _NativeWhisperModel @@ -216,7 +221,7 @@ def test_native_whisper_passes_decode_options_to_transcribe(tmp_path): repo_id="tiny", word_timestamps=False, beam_size=3, - condition_on_previous_text=False, + condition_on_previous_text=condition, ) with patch.dict("sys.modules", {"whisper": Mock()}): @@ -236,7 +241,7 @@ def test_native_whisper_passes_decode_options_to_transcribe(tmp_path): call_kwargs = model.model.transcribe.call_args.kwargs assert call_kwargs["beam_size"] == 3 - assert call_kwargs["condition_on_previous_text"] is False + assert call_kwargs["condition_on_previous_text"] is expected def test_native_init_with_artifacts_path_and_device_logging(tmp_path): diff --git a/tests/test_backend_opendocument.py b/tests/test_backend_opendocument.py index f8512b4a5a..358a651e0d 100644 --- a/tests/test_backend_opendocument.py +++ b/tests/test_backend_opendocument.py @@ -1040,6 +1040,38 @@ def test_odt_table_cell_image_creates_rich_cell_picture(tmp_path: Path): assert pictures[0].image is not None +@pytest.mark.parametrize("in_cell", [False, True]) +def test_ods_embedded_image_is_emitted_once(tmp_path: Path, in_cell: bool): + image_path = tmp_path / "picture.png" + Image.new("RGB", (2, 3), "red").save(image_path) + path = tmp_path / "image.ods" + source = OdfDocument("spreadsheet") + source.body.clear() + table = Table("Images", width=1, height=1) + href = source.add_file(str(image_path)) + frame = Frame.image_frame(href, size=("1cm", "1cm")) + if in_cell: + cell = table.get_cell("A1") + cell.set_value("Image cell") + cell.append(frame) + table.set_cell("A1", cell) + else: + table.append(frame) + source.body.append(table) + source.save(str(path)) + + result = DocumentConverter(allowed_formats=[InputFormat.ODS]).convert(path) + pictures = result.document.pictures + assert len(pictures) == 1 + assert pictures[0].image is not None + pixels = pictures[0].image.pil_image + assert pixels is not None + assert pixels.size == (2, 3) + assert pixels.getpixel((0, 0)) == (255, 0, 0) + html = result.document.export_to_html(image_mode=ImageRefMode.EMBEDDED) + assert html.count("data:image/png;base64,") == 1 + + def _rich_cell_descendants(doc: DoclingDocument, cell: RichTableCell): group = cell.ref.resolve(doc) return [ diff --git a/tests/test_check_skill_routes.py b/tests/test_check_skill_routes.py new file mode 100644 index 0000000000..b15ad00201 --- /dev/null +++ b/tests/test_check_skill_routes.py @@ -0,0 +1,68 @@ +# SPDX-FileCopyrightText: The Docling Contributors +# SPDX-License-Identifier: MIT + +from __future__ import annotations + +import importlib.util +from pathlib import Path + +import pytest + +MODULE_PATH = ( + Path(__file__).resolve().parents[1] / ".github/scripts/check_skill_routes.py" +) +SPEC = importlib.util.spec_from_file_location("check_skill_routes", MODULE_PATH) +assert SPEC is not None and SPEC.loader is not None +skill_routes = importlib.util.module_from_spec(SPEC) +SPEC.loader.exec_module(skill_routes) + + +def make_skill(repo_root: Path, name: str) -> Path: + skill = repo_root / ".agents/skills" / name + skill.mkdir(parents=True) + (skill / "SKILL.md").write_text("# Review\n", encoding="utf-8") + for agent in ("codex", "claude"): + link = repo_root / f".{agent}/skills" / name + link.parent.mkdir(parents=True, exist_ok=True) + link.symlink_to(f"../../.agents/skills/{name}", target_is_directory=True) + return skill + + +def test_routing_rejects_missing_skills_and_broken_discovery(tmp_path: Path) -> None: + skill = make_skill(tmp_path, "review") + (skill.parent / "skill-router.json").write_text( + '{"review": "Review a PR"}', encoding="utf-8" + ) + assert skill_routes.check_routes(tmp_path) == [] + + (skill / "SKILL.md").unlink() + assert "Route has no SKILL.md: review" in skill_routes.check_routes(tmp_path) + (skill / "SKILL.md").write_text("# Review\n", encoding="utf-8") + link = tmp_path / ".codex/skills/review" + link.unlink() + link.symlink_to("../../.agents/skills/missing", target_is_directory=True) + assert any( + ".codex/skills/review must link" in error + for error in skill_routes.check_routes(tmp_path) + ) + + +def test_routing_rejects_unrouted_contributor_and_agent_skills(tmp_path: Path) -> None: + skill = make_skill(tmp_path, "review") + (skill.parent / "skill-router.json").write_text( + '{"review": "Review a PR"}', encoding="utf-8" + ) + make_skill(tmp_path, "new-skill") + errors = skill_routes.check_routes(tmp_path) + assert "Contributor skill has no task route: new-skill" in errors + assert "Agent skill has no task route: .claude/skills/new-skill" in errors + + +@pytest.mark.parametrize( + "content", ["{", "[]", "{}", '{"../escape": "Task"}', '{"review": null}'] +) +def test_routing_rejects_invalid_router(tmp_path: Path, content: str) -> None: + skills = tmp_path / ".agents/skills" + skills.mkdir(parents=True) + (skills / "skill-router.json").write_text(content, encoding="utf-8") + assert skill_routes.check_routes(tmp_path) diff --git a/tests/test_cli.py b/tests/test_cli.py index e973777d29..f9bd2311ea 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -1807,3 +1807,12 @@ def convert_all( assert captured["pipeline_options"].parser_threads == 3 assert captured["backend_options"].parser_threads == 3 assert captured["pipeline_options"].accelerator_options.num_threads == 7 + + +def test_cli_asr_artifacts_path_reaches_audio_options(tmp_path, monkeypatch): + result, options = _capture_cli_asr_pipeline_options( + monkeypatch, ["--pipeline", "asr", "--artifacts-path", str(tmp_path)], tmp_path + ) + assert result.exit_code == 0, result.output + assert options is not None + assert options.artifacts_path == tmp_path diff --git a/tests/test_inference_engine_type_contracts.py b/tests/test_inference_engine_type_contracts.py new file mode 100644 index 0000000000..4b9618f951 --- /dev/null +++ b/tests/test_inference_engine_type_contracts.py @@ -0,0 +1,120 @@ +# SPDX-FileCopyrightText: The Docling Contributors +# SPDX-License-Identifier: MIT + +import json +import logging +from types import SimpleNamespace + +import numpy as np +import pytest +import requests +from PIL import Image + +from docling.datamodel.accelerator_options import AcceleratorOptions +from docling.datamodel.pipeline_options_vlm_model import ( + ApiVlmOptions, + InferenceFramework, + InlineVlmOptions, + ResponseFormat, +) +from docling.datamodel.vlm_engine_options import ApiVlmEngineOptions +from docling.models.inference_engines.common.kserve_v2_http import KserveV2HttpClient +from docling.models.inference_engines.vlm.api_openai_compatible_engine import ( + ApiVlmEngine, +) +from docling.models.inference_engines.vlm.base import VlmEngineInput +from docling.models.vlm_pipeline_models.api_vlm_model import ApiVlmModel +from docling.models.vlm_pipeline_models.mlx_model import HuggingFaceMlxModel + +pytestmark = pytest.mark.cross_platform + + +def test_kserve_profiling_survives_a_log_level_change(monkeypatch, caplog): + logger_name = "docling.models.inference_engines.common.kserve_v2_http" + caplog.set_level(logging.INFO, logger=logger_name) + + def post(*args, **kwargs): + logging.getLogger(logger_name).setLevel(logging.DEBUG) + response = requests.Response() + response.status_code = 200 + response._content = json.dumps( + { + "outputs": [ + {"name": "out", "datatype": "FP32", "shape": [1], "data": [1.0]} + ] + } + ).encode() + return response + + monkeypatch.setattr(requests, "post", post) + client = KserveV2HttpClient( + base_url="http://test", + model_name="model", + model_version=None, + timeout=1, + headers={}, + use_binary_data=False, + ) + result = client.infer( + inputs={"input": np.zeros((1,), dtype=np.float32)}, output_names=["out"] + ) + np.testing.assert_array_equal(result["out"], np.array([1.0], dtype=np.float32)) + + +@pytest.mark.parametrize("key", ["usage_response_key", "token_extract_key"]) +@pytest.mark.parametrize("engine", [False, True]) +def test_api_response_keys_require_text_before_a_request(key, engine): + image = Image.new("RGB", (2, 2)) + with pytest.raises(ValueError, match=key + " must be a string or None"): + if engine: + model = ApiVlmEngine( + True, ApiVlmEngineOptions(url="http://test", params={key: 42}) + ) + model.predict_batch([VlmEngineInput(image=image, prompt="Read")]) + else: + model = ApiVlmModel( + True, + True, + ApiVlmOptions( + url="http://test", + params={key: 42}, + prompt="Read", + response_format=ResponseFormat.PLAINTEXT, + ), + ) + list(model.process_images([image], "Read")) + + +@pytest.mark.parametrize( + "logprobs", [None, [-0.25], np.array([[-0.25]], dtype=np.float32), "mlx-bfloat16"] +) +def test_mlx_stream_accepts_optional_and_list_logprobs(logprobs): + if isinstance(logprobs, str): + mx = pytest.importorskip("mlx.core") + logprobs = mx.array([-0.25], dtype=mx.bfloat16) + model = HuggingFaceMlxModel( + False, + None, + AcceleratorOptions(), + InlineVlmOptions( + repo_id="test/model", + prompt="Read", + inference_framework=InferenceFramework.MLX, + response_format=ResponseFormat.PLAINTEXT, + ), + ) + model.vlm_model = object() + model.processor = object() + model.config = {} + model.apply_chat_template = lambda *args, **kwargs: "Read" + model.stream_generate = lambda *args, **kwargs: iter( + [ + SimpleNamespace(text="a", token=0, logprobs=logprobs), + ] + ) + (prediction,) = model.process_images([Image.new("RGB", (2, 2))], "Read") + assert prediction.text == "a" + if logprobs is None: + assert prediction.generated_tokens == [] + else: + assert prediction.generated_tokens[0].logprob == pytest.approx(-0.25) diff --git a/tests/test_pdf_streaming_foundation.py b/tests/test_pdf_streaming_foundation.py index 9453ace829..6814e8328a 100644 --- a/tests/test_pdf_streaming_foundation.py +++ b/tests/test_pdf_streaming_foundation.py @@ -162,7 +162,7 @@ def test_standard_pipeline_bounds_live_streaming_backends() -> None: generate_parsed_pages=False, ) pipeline._create_run_ctx = MethodType( - lambda self: _make_run_context(self._release_page_resources), pipeline + lambda self: _make_run_context(self._release_threaded_item_resources), pipeline ) conv_res = _make_conversion_result(backend, page_count) diff --git a/tests/test_picture_description_vlm_engine_model.py b/tests/test_picture_description_vlm_engine_model.py index 98d1f6ef18..e8eb801da5 100644 --- a/tests/test_picture_description_vlm_engine_model.py +++ b/tests/test_picture_description_vlm_engine_model.py @@ -20,6 +20,8 @@ class _DummyEngine: + selected_engine_type = None + def __init__(self): self.received_inputs: list[VlmEngineInput] = [] diff --git a/tests/test_pipeline_type_contracts.py b/tests/test_pipeline_type_contracts.py new file mode 100644 index 0000000000..64e369f446 --- /dev/null +++ b/tests/test_pipeline_type_contracts.py @@ -0,0 +1,91 @@ +# SPDX-FileCopyrightText: The Docling Contributors +# SPDX-License-Identifier: MIT + +import logging +import threading +from io import BytesIO + +import pytest + +from docling.datamodel.base_models import DocumentStream, InputFormat, Page +from docling.datamodel.document import ConversionResult, _DocumentConversionInput +from docling.document_converter import HTMLFormatOption +from docling.pipeline.standard_pdf_pipeline import ( + ThreadedItem, + ThreadedPipelineStage, + ThreadedQueue, +) + + +def test_url_input_with_none_limits_uses_default_limits(monkeypatch): + calls = [] + + def resolve(source, headers, *, max_file_size): + calls.append(max_file_size) + return DocumentStream(name="page.html", stream=BytesIO(b"

    Hello

    ")) + + monkeypatch.setattr("docling.datamodel.document.resolve_source_to_stream", resolve) + inputs = _DocumentConversionInput( + path_or_stream_iterator=["https://example.com/page.html"], limits=None + ) + documents = list(inputs.docs({InputFormat.HTML: HTMLFormatOption()})) + assert len(documents) == 1 and documents[0].valid + assert documents[0]._backend.convert().export_to_markdown() == "Hello" + assert len(calls) == 1 + + +@pytest.mark.parametrize("timeout", [None, 5.0]) +def test_queue_close_wakes_blocked_producer(monkeypatch, timeout): + queue = ThreadedQueue(max_size=0) + waiting = threading.Event() + original_wait = queue._not_full.wait + + def wait(timeout=None): + waiting.set() + return original_wait(timeout) + + monkeypatch.setattr(queue._not_full, "wait", wait) + accepted = [] + # A full zero-capacity queue never inspects the item payload. + producer = threading.Thread( + target=lambda: accepted.append(queue.put(None, timeout)) + ) + producer.start() + reached_wait = waiting.wait(2.0) + queue.close() + producer.join(2.0) + assert reached_wait + assert not producer.is_alive() + assert accepted == [False] + + +def test_log_level_change_during_stage_does_not_fail_page(): + inputs = _DocumentConversionInput( + path_or_stream_iterator=[ + DocumentStream(name="page.html", stream=BytesIO(b"

    Hello

    ")) + ] + ) + document = next(iter(inputs.docs({InputFormat.HTML: HTMLFormatOption()}))) + result = ConversionResult(input=document) + item = ThreadedItem(payload=Page(page_no=1), run_id=1, page_no=1, conv_res=result) + logger = logging.getLogger("docling.pipeline.standard_pdf_pipeline") + original_level = logger.level + + def model(result, pages): + logger.setLevel(logging.DEBUG) + return pages + + stage = ThreadedPipelineStage( + name="test", + model=model, + batch_size=1, + batch_timeout=0.1, + queue_max_size=1, + ) + try: + logger.setLevel(logging.INFO) + assert stage._process_batch([item]) == [item] + assert not item.is_failed + assert item.error is None + finally: + logger.setLevel(original_level) diff --git a/tests/test_rapid_ocr_model.py b/tests/test_rapid_ocr_model.py index fdd98401ba..9ed20396a4 100644 --- a/tests/test_rapid_ocr_model.py +++ b/tests/test_rapid_ocr_model.py @@ -4,10 +4,15 @@ from io import BytesIO from pathlib import Path +import numpy as np import pytest +from PIL import Image +from docling.backend.image_backend import ImageDocumentBackend from docling.datamodel.accelerator_options import AcceleratorOptions -from docling.datamodel.pipeline_options import RapidOcrOptions +from docling.datamodel.base_models import InputFormat, Page +from docling.datamodel.document import ConversionResult, InputDocument +from docling.datamodel.pipeline_options import OcrMode, RapidOcrOptions from docling.models.stages.ocr.rapid_ocr_model import RapidOcrModel pytestmark = pytest.mark.ml_ocr @@ -106,3 +111,52 @@ def test_rapidocr_pins_explicit_model_paths( assert params["Rec.model_path"] is not None assert "Det.lang_type" not in params assert "Rec.lang_type" not in params + + +@pytest.mark.parametrize( + "result_kind", ["complete", "detection", "no_text", "no_scores"] +) +def test_rapidocr_uses_only_complete_text_results(result_kind: str): + from rapidocr.ch_ppocr_det.utils import TextDetOutput + from rapidocr.utils.output import RapidOCROutput + + boxes = np.array([[[0, 0], [12, 0], [12, 6], [0, 6]]]) + if result_kind == "detection": + result = TextDetOutput(boxes=boxes, scores=[0.9]) + else: + result = RapidOCROutput( + boxes=boxes, + txts=None if result_kind == "no_text" else ("word",), + scores=None if result_kind == "no_scores" else (0.9,), + ) + + source = BytesIO() + Image.new("RGB", (20, 20), "white").save(source, format="PNG") + source.seek(0) + in_doc = InputDocument( + path_or_stream=source, + format=InputFormat.IMAGE, + backend=ImageDocumentBackend, + filename="ocr.png", + ) + conv_res = ConversionResult(input=in_doc) + page = Page(page_no=1) + page._backend = in_doc._backend.load_page(0) + page.size = page._backend.get_size() + + model = RapidOcrModel( + enabled=False, + artifacts_path=None, + options=RapidOcrOptions(mode=OcrMode.FULL_PAGE, scale=2), + accelerator_options=AcceleratorOptions(), + ) + model.enabled = True + model.reader = lambda *args, **kwargs: result + + assert list(model(conv_res, [page])) == [page] + if result_kind == "complete": + assert [cell.text for cell in page.cells] == ["word"] + assert page.cells[0].confidence == 0.9 + assert page.cells[0].rect.to_bounding_box().as_tuple() == (0, 0, 6, 3) + else: + assert page.cells == [] diff --git a/tests/test_service_client_sdk_unit.py b/tests/test_service_client_sdk_unit.py index d4286ffc03..46dad8dc66 100644 --- a/tests/test_service_client_sdk_unit.py +++ b/tests/test_service_client_sdk_unit.py @@ -48,6 +48,8 @@ S3SourceRequest, ) from docling.datamodel.service.responses import ( + ConvertDocumentResponse, + ExportDocumentResponse, FailureCategory, FailurePhase, MessageKind, @@ -110,14 +112,15 @@ def _status_response(task_id: str, status: str) -> TaskStatusResponse: ) -def _convert_payload(source_name: str) -> SimpleNamespace: - return SimpleNamespace( +def _convert_payload(source_name: str) -> ConvertDocumentResponse: + return ConvertDocumentResponse( + processing_time=0.0, status=ConversionStatus.SUCCESS, errors=[], timings={}, confidence=None, - document=SimpleNamespace( - filename=None, + document=ExportDocumentResponse( + filename="", json_content=DoclingDocument(name=PurePath(source_name).stem), ), ) diff --git a/tests/test_vlm_convert_model.py b/tests/test_vlm_convert_model.py index 239b693eae..001803c3b3 100644 --- a/tests/test_vlm_convert_model.py +++ b/tests/test_vlm_convert_model.py @@ -21,6 +21,8 @@ class _CapturingEngine: + selected_engine_type = None + def __init__(self) -> None: self.batches = [] diff --git a/uv.lock b/uv.lock index 42881ee979..cd06dfc666 100644 --- a/uv.lock +++ b/uv.lock @@ -2432,7 +2432,7 @@ dev = [ { name = "pytest-xdist", specifier = "~=3.3" }, { name = "python-semantic-release", specifier = "~=7.32" }, { name = "tach", specifier = ">=0.34.0,<1" }, - { name = "ty", specifier = "==0.0.33" }, + { name = "ty", specifier = "==0.0.84" }, { name = "types-defusedxml", specifier = ">=0.7.0.20250822,<0.8.0" }, { name = "types-openpyxl", specifier = "~=3.1" }, { name = "types-requests", specifier = "~=2.31" }, @@ -2466,7 +2466,7 @@ pr-fast-checks = [ { name = "pydantic", specifier = ">=2.0.0,<3.0.0" }, { name = "pydantic-settings", specifier = ">=2.3.0,<3.0.0" }, { name = "ruff", specifier = "==0.15.12" }, - { name = "ty", specifier = "==0.0.33" }, + { name = "ty", specifier = "==0.0.84" }, { name = "types-defusedxml", specifier = ">=0.7.0.20250822,<0.8.0" }, { name = "types-openpyxl", specifier = "~=3.1" }, { name = "types-requests", specifier = "~=2.31" }, @@ -2477,7 +2477,7 @@ pr-fast-checks = [ typecheck = [ { name = "boto3-stubs", specifier = "~=1.37" }, { name = "pandas-stubs", specifier = "~=2.1" }, - { name = "ty", specifier = "==0.0.33" }, + { name = "ty", specifier = "==0.0.84" }, { name = "types-defusedxml", specifier = ">=0.7.0.20250822,<0.8.0" }, { name = "types-openpyxl", specifier = "~=3.1" }, { name = "types-requests", specifier = "~=2.31" }, @@ -10549,26 +10549,27 @@ wheels = [ [[package]] name = "ty" -version = "0.0.33" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/84/44/9478c50c266826c1bf30d1692e589755bffa8f1c0a3eb7af8a346c255991/ty-0.0.33.tar.gz", hash = "sha256:46d63bda07403322cb6c28ccfdd5536be916e13df725c29f7ccd0a21f06bd9e8", size = 5559373, upload-time = "2026-04-28T10:45:13.18Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/e9/24/e287388c63a19191be26b32ff4dbd06029834068150ebe2532939bc4c851/ty-0.0.33-py3-none-linux_armv6l.whl", hash = "sha256:94d0a9d2234261a8911396d59e506b5923fe0971dbda43b9dcea287936887fcc", size = 11021308, upload-time = "2026-04-28T10:45:43.34Z" }, - { url = "https://files.pythonhosted.org/packages/00/ca/ba1eed819895bd239fba8ee35dfcd5fcb266c203b0914a17a59579096bb5/ty-0.0.33-py3-none-macosx_10_12_x86_64.whl", hash = "sha256:e4a2b5ba078f90de342f56b5f7979bb77c9b9b1d8625a041352ffc6ee93c4073", size = 10777272, upload-time = "2026-04-28T10:45:32.905Z" }, - { url = "https://files.pythonhosted.org/packages/25/a8/c3131d37b44b3fea1d6654a1c929a0cd0873822f77a90482b8ec28f6fbbd/ty-0.0.33-py3-none-macosx_11_0_arm64.whl", hash = "sha256:84ff5707825e9af9668d2bcf66975f93e520a63b524ab494e3a8265735be2563", size = 10201078, upload-time = "2026-04-28T10:45:23.374Z" }, - { url = "https://files.pythonhosted.org/packages/7b/db/d8e37ff0045810cc65e1ff36aa0da0a2253c05659787ac987df8a16c7897/ty-0.0.33-py3-none-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:e375285736f57886868e7af0b11c7b0ec5b6543fa15e7ad2a714fed9f077d4e0", size = 10732347, upload-time = "2026-04-28T10:45:21.444Z" }, - { url = "https://files.pythonhosted.org/packages/e0/1a/20e83a412506a918e4684fc67b567cf7cc13b105470b3428cb23c3d5aa13/ty-0.0.33-py3-none-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:5680f6350c3b4e46b8bff6d7bb132366ea239463d6cad4892725d06046e65464", size = 10808238, upload-time = "2026-04-28T10:45:38.565Z" }, - { url = "https://files.pythonhosted.org/packages/5d/4b/d0a39f4464dc6cb4cc2c159473ce216bd1846bfb684c0323a3cb36dce5c6/ty-0.0.33-py3-none-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:c5535538bad8d0f7e62bcdff02197cdb30e41451d80b35d27e17d128f2e1dc5d", size = 11288348, upload-time = "2026-04-28T10:45:08.419Z" }, - { url = "https://files.pythonhosted.org/packages/35/7e/f1745e0f9583363d7a83d9a4990fc244f76ecc30840ddad83dc16a33c52d/ty-0.0.33-py3-none-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:da196c42bbbc069e1e21e3e52107c061aa9660352dae57a41930690b56e2c02d", size = 11789907, upload-time = "2026-04-28T10:45:19.064Z" }, - { url = "https://files.pythonhosted.org/packages/a5/71/25f39f46a12d662859d45bc648555d0661044eb43db6b5648c9947487da9/ty-0.0.33-py3-none-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:9281672921ef6d4460e03146b5e6c18cb1a3e3a3b8a1a88f6f33226d05a469b7", size = 11500774, upload-time = "2026-04-28T10:45:48.012Z" }, - { url = "https://files.pythonhosted.org/packages/94/ec/136959ecbb7c71cb90537f5aea441c73f4ab24612868a6ecdc9d7444d32d/ty-0.0.33-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:82c1b8f303f82da64e878108e764be3ecbcd7c9903ac0a7f7031614ed00b97ab", size = 11360314, upload-time = "2026-04-28T10:45:05.402Z" }, - { url = "https://files.pythonhosted.org/packages/cf/95/32809575c222f00beed498cb728e9290a0f5009f930025381bb7253b2206/ty-0.0.33-py3-none-musllinux_1_2_aarch64.whl", hash = "sha256:efe3af412c9ff67bce5fa37d0a2b0d8555c24072b145a5bac6c79637f1c83abe", size = 10707785, upload-time = "2026-04-28T10:45:10.836Z" }, - { url = "https://files.pythonhosted.org/packages/13/89/c8e9531f7aa4a093359e15fa32c8e1277fbbe90d16894d7c6032d29f4b34/ty-0.0.33-py3-none-musllinux_1_2_armv7l.whl", hash = "sha256:aeec29c91ea768601747da546c3efc20b72c2fb1bd52bcc786a5c6eeff51d27b", size = 10834987, upload-time = "2026-04-28T10:45:40.738Z" }, - { url = "https://files.pythonhosted.org/packages/31/16/9835fbcf5338af1a1917bd28fdb8a7193c210b83f243aa286fa9f79cb3ad/ty-0.0.33-py3-none-musllinux_1_2_i686.whl", hash = "sha256:a535977c52bbb5f7e96b8b70a6ad375ad077f4a9ff2492508ea3816a2b403819", size = 10968968, upload-time = "2026-04-28T10:45:30.26Z" }, - { url = "https://files.pythonhosted.org/packages/36/69/64c76aabc1bc70c7f24b686cd93c3407f8ea430905e395f59bf9603ef571/ty-0.0.33-py3-none-musllinux_1_2_x86_64.whl", hash = "sha256:1d732facf39fcb221ba279d469c5040d37883e964f123b1563888efd34818180", size = 11458077, upload-time = "2026-04-28T10:45:45.971Z" }, - { url = "https://files.pythonhosted.org/packages/91/84/fae27b0c4718776a298690d31ca4cc1995f2e3e1c63a7b59e84c41498e9a/ty-0.0.33-py3-none-win32.whl", hash = "sha256:d90960b574428dc252f85e8598ec5fcb7f619794196b2fc95a90da075ed4681c", size = 10345364, upload-time = "2026-04-28T10:45:16.836Z" }, - { url = "https://files.pythonhosted.org/packages/3c/a0/a2938b23ae3e1a09a2d7c189e2ac5f7113676bae4e0e23948b568e18e5f8/ty-0.0.33-py3-none-win_amd64.whl", hash = "sha256:c1c3aec62c44de610c6e95f0a4e97ac3dbc07934bfdbf1fd90d758c9ff72f48e", size = 11342470, upload-time = "2026-04-28T10:45:26.455Z" }, - { url = "https://files.pythonhosted.org/packages/ab/62/7fb948aace38d2f6329261bb33c035a8484549c74f1db28649c7a4c6fed9/ty-0.0.33-py3-none-win_arm64.whl", hash = "sha256:0d44f99ba1b441e55e2aa301b2ac0a21112784931b46a5f66f4ea9efe5620d97", size = 10742673, upload-time = "2026-04-28T10:45:35.555Z" }, +version = "0.0.84" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/e1/f6/34f8869c45ea87fe6a931ffa824b1fe4428167173543d92988d631add355/ty-0.0.84.tar.gz", hash = "sha256:0cefdb0cd5d399418dbe83c349ba605047b7dff815ae6f460c80e8522b40be67", size = 7409537, upload-time = "2026-09-24T13:16:24.367Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/aa/0c/600648778e1c79133ce6b6a5c4c2531ae16b90c9b30f5041b38b574f61b8/ty-0.0.84-py3-none-linux_armv6l.whl", hash = "sha256:868438e2b9abfc835dd5b5b07ed693db95d8348ac40e3228252511d63ae87695", size = 13935333, upload-time = "2026-09-24T13:15:31.167Z" }, + { url = "https://files.pythonhosted.org/packages/c5/3b/27f06f49945c36170765fad82076898039ae3873338e9865f876dd62f12c/ty-0.0.84-py3-none-macosx_10_12_x86_64.whl", hash = "sha256:b00f39b521f0b559cde0f6abf8dd47213b513b04f33d58efeb9b98486a16bb17", size = 13511328, upload-time = "2026-09-24T13:15:34.446Z" }, + { url = "https://files.pythonhosted.org/packages/fc/b6/b12dd130e7f30141f0645614f454c925991cd389b1cc6c8284fd6bf32326/ty-0.0.84-py3-none-macosx_11_0_arm64.whl", hash = "sha256:8f78c2915567f2ee62ce0c24d86091bc76f81577677755e242d09951db588b9a", size = 13445445, upload-time = "2026-09-24T13:15:37.476Z" }, + { url = "https://files.pythonhosted.org/packages/44/96/f012de4b3c1d9a4773326621380d768b4d48c47e620dc012775c5e01c7d7/ty-0.0.84-py3-none-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:61155cad0921e890d86d1d911a40b3a0e924ef4bdc3b4f3865a6c89c7e3012dc", size = 13471235, upload-time = "2026-09-24T13:15:40.64Z" }, + { url = "https://files.pythonhosted.org/packages/05/32/dca630f5f5353fa7bf67eebccf1549df99b16656d21071d0df6677c826a5/ty-0.0.84-py3-none-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:417c3bf14947ed16b0b92976ab333c27915590354e96d4f581709068d3811355", size = 13680704, upload-time = "2026-09-24T13:15:43.817Z" }, + { url = "https://files.pythonhosted.org/packages/44/8d/33702c8f62f05ad45b33644fb914d483de6a0a7bc8e51884e9b8773e6d69/ty-0.0.84-py3-none-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:e06ef117759f654f5379efe9ae79bda8a186016c016f1023d05c34f8e15c1546", size = 14518923, upload-time = "2026-09-24T13:15:46.802Z" }, + { url = "https://files.pythonhosted.org/packages/c8/a1/da3246e6ef2ae0e842bca4058b25c54bee55ab058ffff4ee54e72517ec01/ty-0.0.84-py3-none-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:a05d1d53db8b64410ab38c6e32c5898510584869d1d5dd9c0605db81c9aef7df", size = 15027731, upload-time = "2026-09-24T13:15:49.755Z" }, + { url = "https://files.pythonhosted.org/packages/20/17/b0099098a560aeb4f9fde6c912f9a81622c3e44e6555e1fbe322b368ec31/ty-0.0.84-py3-none-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:b92400494ab855cfb1204ac1101e3f53b2d09a8e21d171758b2cf9acdab1809e", size = 14748260, upload-time = "2026-09-24T13:15:52.713Z" }, + { url = "https://files.pythonhosted.org/packages/60/d5/32cd67aedb271f006e716ad093806c96c5d7e2af86be8255469dbd3d76c8/ty-0.0.84-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:af17eb391ae3027fed9be1780bea555c2b8a25ba68e184c627b4981b2c9eaab7", size = 14205858, upload-time = "2026-09-24T13:15:55.677Z" }, + { url = "https://files.pythonhosted.org/packages/43/c4/638ac62ca2e8eca1b92e3e16273085744fada728b993500819b5c258ca15/ty-0.0.84-py3-none-manylinux_2_31_riscv64.whl", hash = "sha256:44cf06af0e7ec78ea6d8ff33a3450fc19627756a4fba681f7410d4e3ef63047a", size = 14573196, upload-time = "2026-09-24T13:15:59.165Z" }, + { url = "https://files.pythonhosted.org/packages/a6/f6/9ee81c5d8921e5fd5f10889ae563c95cd3272b31df2ada528919ca1cef6b/ty-0.0.84-py3-none-musllinux_1_2_aarch64.whl", hash = "sha256:c6e56b443ff9ba462e34bfbc1e582004c2719ed35f8f93a433744b0ba0afb2d9", size = 13480760, upload-time = "2026-09-24T13:16:02.313Z" }, + { url = "https://files.pythonhosted.org/packages/f6/83/fbe5d1177667ae16d84c3242846b143809c9cc09568b333e5a9b4d5908f6/ty-0.0.84-py3-none-musllinux_1_2_armv7l.whl", hash = "sha256:8b6562fcc48820c1030e2a4509ef2c14ef93043ca4860280074314ac3a825b95", size = 13696346, upload-time = "2026-09-24T13:16:05.391Z" }, + { url = "https://files.pythonhosted.org/packages/86/9a/1b4ac9ddc34cc0ab4624bc219825d66af52f1ad87f4778ad08ad03fa06d5/ty-0.0.84-py3-none-musllinux_1_2_i686.whl", hash = "sha256:dee7421451efcbd70f03e95a8f2c8d9e5ec3d35edcd27b4095ee00ee062abfc4", size = 13959410, upload-time = "2026-09-24T13:16:08.661Z" }, + { url = "https://files.pythonhosted.org/packages/c9/04/70e289437a67dfd686056ab4fd91880261a8233033a6a95d29393741fb97/ty-0.0.84-py3-none-musllinux_1_2_x86_64.whl", hash = "sha256:31e96307cf621babe969307dd657783d672e2d8bdbd06493710da635bf63fa07", size = 14350294, upload-time = "2026-09-24T13:16:11.865Z" }, + { url = "https://files.pythonhosted.org/packages/9c/f4/ae7dbd118425c850981d23efb850f0d9d7f6b8d2b852bde2767434f668b4/ty-0.0.84-py3-none-win32.whl", hash = "sha256:4bdc85f91cbaae35825f2b17bb4ac5de0563367875cb0f42b2bf07a1bf7ff2ba", size = 13176439, upload-time = "2026-09-24T13:16:15.049Z" }, + { url = "https://files.pythonhosted.org/packages/29/ec/994d87252f065dda5a500bc8afc726d18f6582d81d6f4ff1feb8a267c800/ty-0.0.84-py3-none-win_amd64.whl", hash = "sha256:71cc9aa7a7c89d4e12c6a814e9c0c3310ad5bee44efedbd6f5a42d91e9ac0f0d", size = 13962879, upload-time = "2026-09-24T13:16:18.191Z" }, + { url = "https://files.pythonhosted.org/packages/6d/1c/e6e91852b72566b1ab9aca504712857dfb89324c12092a5bd61d42d34427/ty-0.0.84-py3-none-win_arm64.whl", hash = "sha256:6c9d3718c4a0d318bc9c92af9b2aa85c3294e7dff54cf1d69ad79568f56516da", size = 13656453, upload-time = "2026-09-24T13:16:21.51Z" }, ] [[package]]