diff --git a/CHANGELOG.md b/CHANGELOG.md index 0e4744c..eb20ea1 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,56 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 --- +## [0.4.0] — 2026-07-07 + +### Added — the raw layer (`contextifier.raw`) + +Contextifier now offers **two views of every document**: the existing +AI-friendly extraction pipeline, and `open_raw()` — a lossless, +addressable, **writable** view of OOXML files (`.xlsx` / `.docx` / +`.pptx`). + +- **`open_raw(source)`** (top-level export) and + `DocumentProcessor.open_raw()` — dispatch to `XlsxRawDocument` / + `DocxRawDocument` / `PptxRawDocument` by extension or package sniffing. +- **Byte-preservation contract** (`raw/opc.py`): saving writes untouched + parts back byte-identically (BOM, whitespace, extensions and all). + Charts, pivot tables, sparklines, custom XML, styles — everything the + usual Office libraries destroy on a load→save round-trip — survives. + Contract-tested against synthetic and real files, including a contrast + test showing openpyxl 3.1.5 dropping sparklines/customXml/chart styles + that the raw layer preserves. +- **XLSX** (`raw/xlsx.py`): surgical `set_cell` (style attribute + preserved, strings as `inlineStr`, rows/cells kept sorted), + `append_rows`, cached-formula reads + `get_formula`, merged ranges, + defined names, `charts`; overwriting a formula sets + `calcPr fullCalcOnLoad` so Excel recalculates on open. +- **DOCX** (`raw/docx.py`): python-docx-compatible paragraph/table + addressing (incl. gridSpan/vMerge grid semantics and nested tables), + **run-preserving** `set_paragraph_text` and `cell.set_text` (inline + images, drawings, bookmarks and hyperlink elements survive; hyperlinks + are emptied, never deleted — `strip_empty_hyperlinks()` when you want + them gone), `insert_paragraph_after`/`delete_paragraph`, + `insert_row`/`delete_row`, headers/footers, `body_order()`. +- **PPTX** (`raw/pptx.py`): slide/shape inventory (tables, charts, + diagrams, groups), format-preserving `set_text`, table cell edits + + row insert/delete, `notes_text`, **`replace_content(new_xml, + preserve_native=True)`** — swap a slide's XML while carrying over its + native charts/tables/pictures (id-renumbered), and `remove_slide()` + with transitive orphan-part sweeping (chart XML, embedded workbooks, + media) that provably spares shared layouts/masters. +- **Charts** (`raw/chart.py`): one `ChartModel` for all three formats — + reads kind/title/series from classic `c:` charts (incl. scatter/bubble + x/y, literals, multi-level caches) and chartEx; **writes** titles and + full data (`set_data`) with schema-valid cache/ref rewriting, series + grow/shrink keeping styling, and embedded-workbook regeneration — + verified by reopening with python-pptx/openpyxl. +- `lxml` promoted to an explicit dependency. + +99 new tests (598 total). + +--- + ## [0.3.0] — 2025-07-19 (released 2026-07-07) ### Added (RAG provenance & payload ingestion — 2026-07-07) diff --git a/README.md b/README.md index 22a4468..8d1f505 100644 --- a/README.md +++ b/README.md @@ -5,6 +5,7 @@ ## Key Features - **Broad Format Support**: PDF, DOCX, DOC, PPTX, PPT, XLSX, XLS, HWP, HWPX, RTF, CSV, TSV, TXT, MD, HTML, images, code files, and 80+ extensions +- **Two Views of Every Document**: the AI-friendly pipeline (lightweight, normalized text for LLMs) *and* `open_raw()` — a lossless, addressable, **writable** view of OOXML files where saving keeps untouched parts byte-identical - **Intelligent Text Extraction**: Preserves document structure (headings, tables, image positions) with automatic metadata extraction - **Table Processing**: Converts tables to HTML/Markdown/Text with `rowspan`/`colspan` support for merged cells - **OCR Integration**: 5 Vision LLM engines — OpenAI, Anthropic, Google Gemini, AWS Bedrock, vLLM @@ -35,7 +36,39 @@ text = processor.extract_text("document.pdf") print(text) ``` -### 2. Extract + Chunk in One Step +### 2. Raw Access — Read *and Write* OOXML Losslessly + +The extraction pipeline renders an AI-friendly view and discards the +rest. `open_raw()` is its lossless twin: the full package stays +available, edits are surgical, and **untouched parts round-trip +byte-identically** — charts, pivot tables, sparklines, styles and +custom XML all survive (unlike a load→save round-trip through the +usual Office libraries). + +```python +from contextifier import open_raw + +raw = open_raw("report.xlsx") # XlsxRawDocument +raw.sheets["Sales"].set_cell("B3", 142) # surgical edit +raw.charts[0].set_data( # real chart-data editing + categories=["Q1", "Q2", "Q3"], + series=[("Sales", [120, 135, 150])], +) +raw.save("report-edited.xlsx") + +raw = open_raw("paper.docx") # DocxRawDocument +raw.set_paragraph_text(3, "Revised text") # runs & inline images preserved +raw.tables[0].insert_row(2) + +raw = open_raw("deck.pptx") # PptxRawDocument +raw.slides[0].set_text(shape_id=2, new_text="New title") +raw.save("deck2.pptx") +``` + +Supported raw formats: `.xlsx` / `.docx` / `.pptx` (the OOXML trio). +Every model also exposes `.package` for part-level OPC access. + +### 3. Extract + Chunk in One Step ```python from contextifier import DocumentProcessor @@ -50,7 +83,7 @@ for i, chunk in enumerate(result.chunks, 1): result.save_to_md("output/chunks") ``` -### 3. Custom Configuration +### 4. Custom Configuration ```python from contextifier import DocumentProcessor @@ -65,7 +98,7 @@ processor = DocumentProcessor(config=config) text = processor.extract_text("report.xlsx") ``` -### 4. OCR Integration +### 5. OCR Integration ```python from contextifier import DocumentProcessor diff --git a/contextifier/__init__.py b/contextifier/__init__.py index 3236d2c..57135b6 100644 --- a/contextifier/__init__.py +++ b/contextifier/__init__.py @@ -24,11 +24,17 @@ └── OCR (optional vision-based extraction) Usage: - from contextifier import DocumentProcessor + from contextifier import DocumentProcessor, open_raw + # AI-friendly view (lightweight, normalized) processor = DocumentProcessor() text = processor.extract_text("document.pdf") chunks = processor.extract_chunks("document.pdf", chunk_size=1000) + + # Raw view (lossless, addressable, WRITABLE — xlsx/docx/pptx) + raw = open_raw("report.xlsx") + raw.sheets["Sales"].set_cell("B3", 142) + raw.save("report-edited.xlsx") # untouched parts stay byte-identical """ from importlib.metadata import version, PackageNotFoundError @@ -43,6 +49,7 @@ from contextifier.types import ExtractionResult, FileContext, Chunk, ChunkMetadata from contextifier.chunking.chunker import TextChunker from contextifier.errors import ContextifierError, UnsupportedFormatError +from contextifier.raw import open_raw from contextifier.async_processor import AsyncDocumentProcessor from contextifier.cached_processor import CachedDocumentProcessor @@ -63,6 +70,8 @@ "ChunkMetadata", # Chunking "TextChunker", + # Raw (lossless, writable) access + "open_raw", # Errors "ContextifierError", "UnsupportedFormatError", diff --git a/contextifier/cached_processor.py b/contextifier/cached_processor.py index 23313bd..cf50595 100644 --- a/contextifier/cached_processor.py +++ b/contextifier/cached_processor.py @@ -27,7 +27,7 @@ import json import logging from pathlib import Path -from typing import Any, Dict, List, Optional, Protocol, Union +from typing import Any, Dict, Optional, Protocol, Union from contextifier.config import ProcessingConfig from contextifier.document_processor import DocumentProcessor, ChunkResult @@ -47,6 +47,7 @@ # ── Cache backend protocol ──────────────────────────────────────────────── + class CacheBackend(Protocol): """Minimal interface for a cache backend.""" @@ -61,6 +62,7 @@ def put(self, key: str, value: str) -> None: # ── Built-in backends ──────────────────────────────────────────────────── + class MemoryCacheBackend: """ In-memory LRU cache using OrderedDict. @@ -125,6 +127,7 @@ def put(self, key: str, value: str) -> None: # ── Cached processor ───────────────────────────────────────────────────── + class CachedDocumentProcessor: """Document processor with transparent result caching.""" @@ -176,7 +179,10 @@ def process( ) -> ExtractionResult: """Process a file with cache lookup for the full ExtractionResult.""" cache_key = self._make_key( - str(file_path), extract_metadata, ocr_processing, suffix="process", + str(file_path), + extract_metadata, + ocr_processing, + suffix="process", ) cached = self._cache.get(cache_key) if cached is not None: @@ -208,9 +214,11 @@ def extract_chunks( ) -> ChunkResult: """Extract and chunk text with cache lookup.""" cache_key = self._make_key( - str(file_path), extract_metadata, ocr_processing, + str(file_path), + extract_metadata, + ocr_processing, suffix=f"chunks|cs={chunk_size}|co={chunk_overlap}" - f"|pt={preserve_tables}|pm={include_position_metadata}", + f"|pt={preserve_tables}|pm={include_position_metadata}", ) cached = self._cache.get(cache_key) if cached is not None: @@ -283,6 +291,7 @@ def __repr__(self) -> str: # ── Serialization helpers ───────────────────────────────────────────────── + def _serialize_extraction_result(result: ExtractionResult) -> str: """Serialize an ExtractionResult to a JSON string.""" data: Dict[str, Any] = { @@ -302,7 +311,9 @@ def _deserialize_extraction_result(s: str) -> ExtractionResult: data = json.loads(s) return ExtractionResult( text=data.get("text", ""), - metadata=DocumentMetadata.from_dict(data["metadata"]) if data.get("metadata") else None, + metadata=DocumentMetadata.from_dict(data["metadata"]) + if data.get("metadata") + else None, images=data.get("images", []), page_count=data.get("page_count", 0), warnings=data.get("warnings", []), @@ -325,7 +336,9 @@ def _serialize_chunk_result(result: ChunkResult) -> str: "line_end": c.metadata.line_end, "global_start": c.metadata.global_start, "global_end": c.metadata.global_end, - } if c.metadata else None, + } + if c.metadata + else None, } for c in result.chunks_with_metadata ] @@ -366,17 +379,19 @@ def _deserialize_chunk_result(s: str) -> ChunkResult: def _table_to_dict(table: TableData) -> Dict[str, Any]: rows = [] for row in table.rows: - rows.append([ - { - "content": c.content, - "row_span": c.row_span, - "col_span": c.col_span, - "is_header": c.is_header, - "row_index": c.row_index, - "col_index": c.col_index, - } - for c in row - ]) + rows.append( + [ + { + "content": c.content, + "row_span": c.row_span, + "col_span": c.col_span, + "is_header": c.is_header, + "row_index": c.row_index, + "col_index": c.col_index, + } + for c in row + ] + ) return { "rows": rows, "num_rows": table.num_rows, @@ -389,17 +404,19 @@ def _table_to_dict(table: TableData) -> Dict[str, Any]: def _dict_to_table(d: Dict[str, Any]) -> TableData: rows = [] for row_data in d.get("rows", []): - rows.append([ - TableCell( - content=c["content"], - row_span=c.get("row_span", 1), - col_span=c.get("col_span", 1), - is_header=c.get("is_header", False), - row_index=c.get("row_index", 0), - col_index=c.get("col_index", 0), - ) - for c in row_data - ]) + rows.append( + [ + TableCell( + content=c["content"], + row_span=c.get("row_span", 1), + col_span=c.get("col_span", 1), + is_header=c.get("is_header", False), + row_index=c.get("row_index", 0), + col_index=c.get("col_index", 0), + ) + for c in row_data + ] + ) return TableData( rows=rows, num_rows=d.get("num_rows", 0), @@ -414,10 +431,7 @@ def _chart_to_dict(chart: ChartData) -> Dict[str, Any]: "chart_type": chart.chart_type, "title": chart.title, "categories": chart.categories, - "series": [ - {"name": s.name, "values": s.values} - for s in chart.series - ], + "series": [{"name": s.name, "values": s.values} for s in chart.series], } diff --git a/contextifier/chunking/chunker.py b/contextifier/chunking/chunker.py index ef152bd..fa2a5ff 100644 --- a/contextifier/chunking/chunker.py +++ b/contextifier/chunking/chunker.py @@ -29,7 +29,9 @@ from contextifier.chunking.strategies.base import BaseChunkingStrategy from contextifier.chunking.strategies.page_strategy import PageChunkingStrategy from contextifier.chunking.strategies.table_strategy import TableChunkingStrategy -from contextifier.chunking.strategies.protected_strategy import ProtectedChunkingStrategy +from contextifier.chunking.strategies.protected_strategy import ( + ProtectedChunkingStrategy, +) from contextifier.chunking.strategies.plain_strategy import PlainChunkingStrategy @@ -69,12 +71,14 @@ def __init__( self._strategies.extend(custom_strategies) # Built-in strategies - self._strategies.extend([ - TableChunkingStrategy(), - PageChunkingStrategy(), - ProtectedChunkingStrategy(), - PlainChunkingStrategy(), - ]) + self._strategies.extend( + [ + TableChunkingStrategy(), + PageChunkingStrategy(), + ProtectedChunkingStrategy(), + PlainChunkingStrategy(), + ] + ) # Sort by priority (lower = higher priority) self._strategies.sort(key=lambda s: s.priority) @@ -152,11 +156,14 @@ def chunk( from contextifier.chunking.metadata_enricher import ( enrich_chunk_metadata, ) + result = enrich_chunk_metadata(result) return result except NotImplementedError: # Strategy not yet implemented — try next - return self._fallback_chunk(text, effective_config, ext, include_meta, **kwargs) + return self._fallback_chunk( + text, effective_config, ext, include_meta, **kwargs + ) except Exception as e: raise ChunkingError( f"Chunking failed with strategy '{strategy.strategy_name}': {e}", @@ -209,7 +216,8 @@ def _fallback_chunk( try: if strategy.can_handle(text, config, file_extension=ext, **context): return strategy.chunk( - text, config, + text, + config, file_extension=ext, include_position_metadata=include_meta, **context, @@ -220,7 +228,9 @@ def _fallback_chunk( continue # Ultimate fallback: return text as single chunk - self._logger.warning("All chunking strategies failed, returning text as single chunk") + self._logger.warning( + "All chunking strategies failed, returning text as single chunk" + ) if include_meta: return [Chunk(text=text, metadata=ChunkMetadata(chunk_index=0))] return [text] diff --git a/contextifier/chunking/constants.py b/contextifier/chunking/constants.py index 15a2297..bf50254 100644 --- a/contextifier/chunking/constants.py +++ b/contextifier/chunking/constants.py @@ -61,7 +61,9 @@ HTML_TABLE_PATTERN = re.compile(r"]*>.*?", re.DOTALL | re.IGNORECASE) # Textbox block — [textbox]...[/textbox] — config-independent (no TagConfig entry) -TEXTBOX_BLOCK_PATTERN = re.compile(r"\[textbox\].*?\[/textbox\]", re.DOTALL | re.IGNORECASE) +TEXTBOX_BLOCK_PATTERN = re.compile( + r"\[textbox\].*?\[/textbox\]", re.DOTALL | re.IGNORECASE +) # Markdown tables — config-independent MARKDOWN_TABLE_PATTERN = re.compile( @@ -76,12 +78,16 @@ # Pattern Builder Functions — Config-Driven # ============================================================================ + def _build_block_pattern( - open_tag: str, close_tag: str, flags: int = re.DOTALL, + open_tag: str, + close_tag: str, + flags: int = re.DOTALL, ) -> re.Pattern[str]: """Build regex for block-level tags like [chart]...[/chart].""" return re.compile( - rf"{re.escape(open_tag)}.*?{re.escape(close_tag)}", flags, + rf"{re.escape(open_tag)}.*?{re.escape(close_tag)}", + flags, ) @@ -114,7 +120,9 @@ def build_protected_patterns(config: "ProcessingConfig") -> list[re.Pattern[str] tags = config.tags return [ # Block-level protected regions - _build_block_pattern(tags.chart_prefix, tags.chart_suffix, re.DOTALL | re.IGNORECASE), + _build_block_pattern( + tags.chart_prefix, tags.chart_suffix, re.DOTALL | re.IGNORECASE + ), TEXTBOX_BLOCK_PATTERN, HTML_TABLE_PATTERN, _build_block_pattern(tags.metadata_prefix, tags.metadata_suffix, re.DOTALL), @@ -148,10 +156,10 @@ def build_image_capture_pattern(config: "ProcessingConfig") -> re.Pattern[str]: # Table Chunking Thresholds # ============================================================================ -TABLE_WRAPPER_OVERHEAD: int = 30 # \n
-ROW_OVERHEAD: int = 12 # \n -CELL_OVERHEAD: int = 10 # or -CHUNK_INDEX_OVERHEAD: int = 30 # [Table chunk 1/10]\n +TABLE_WRAPPER_OVERHEAD: int = 30 # \n
+ROW_OVERHEAD: int = 12 # \n +CELL_OVERHEAD: int = 10 # or +CHUNK_INDEX_OVERHEAD: int = 30 # [Table chunk 1/10]\n TABLE_SIZE_THRESHOLD_MULTIPLIER: float = 1.2 # 1.2× of chunk_size # Extensions that are inherently table-based @@ -162,6 +170,7 @@ def build_image_capture_pattern(config: "ProcessingConfig") -> re.Pattern[str]: # Dataclasses # ============================================================================ + @dataclass(frozen=True) class TableRow: """A single table row (HTML or Markdown).""" diff --git a/contextifier/chunking/metadata_enricher.py b/contextifier/chunking/metadata_enricher.py index 8cd2dff..a32d110 100644 --- a/contextifier/chunking/metadata_enricher.py +++ b/contextifier/chunking/metadata_enricher.py @@ -32,6 +32,7 @@ # collides because extraction renders headings with markdown hashes. _HEADING_RE = re.compile(r"^(#{1,6})\s+(.+?)\s*$", re.MULTILINE) + def _heading_breadcrumb(stack: Dict[int, str]) -> Optional[str]: if not stack: return None diff --git a/contextifier/chunking/strategies/__init__.py b/contextifier/chunking/strategies/__init__.py index 2cafb9d..85e54d4 100644 --- a/contextifier/chunking/strategies/__init__.py +++ b/contextifier/chunking/strategies/__init__.py @@ -10,7 +10,9 @@ from contextifier.chunking.strategies.base import BaseChunkingStrategy from contextifier.chunking.strategies.page_strategy import PageChunkingStrategy from contextifier.chunking.strategies.table_strategy import TableChunkingStrategy -from contextifier.chunking.strategies.protected_strategy import ProtectedChunkingStrategy +from contextifier.chunking.strategies.protected_strategy import ( + ProtectedChunkingStrategy, +) from contextifier.chunking.strategies.plain_strategy import PlainChunkingStrategy __all__ = [ diff --git a/contextifier/chunking/strategies/base.py b/contextifier/chunking/strategies/base.py index cf3e100..05d67f5 100644 --- a/contextifier/chunking/strategies/base.py +++ b/contextifier/chunking/strategies/base.py @@ -16,10 +16,10 @@ import logging from abc import ABC, abstractmethod -from typing import Any, Dict, List, Optional, Union +from typing import Any, List, Union -from contextifier.config import ProcessingConfig, ChunkingConfig -from contextifier.types import Chunk, ChunkMetadata +from contextifier.config import ProcessingConfig +from contextifier.types import Chunk class BaseChunkingStrategy(ABC): diff --git a/contextifier/chunking/strategies/page_strategy.py b/contextifier/chunking/strategies/page_strategy.py index 9ca3835..7307fa4 100644 --- a/contextifier/chunking/strategies/page_strategy.py +++ b/contextifier/chunking/strategies/page_strategy.py @@ -24,7 +24,7 @@ import re import logging -from typing import Any, List, Optional, Tuple, Union +from typing import Any, List, Tuple, Union from contextifier.config import ProcessingConfig from contextifier.types import Chunk, ChunkMetadata @@ -86,9 +86,13 @@ def chunk( if not pages: # No page markers found → fall back to plain splitting - from contextifier.chunking.strategies.plain_strategy import PlainChunkingStrategy + from contextifier.chunking.strategies.plain_strategy import ( + PlainChunkingStrategy, + ) + return PlainChunkingStrategy().chunk( - text, config, + text, + config, file_extension=file_extension, include_position_metadata=include_position_metadata, **context, @@ -99,7 +103,11 @@ def chunk( # 2. Greedy page merging max_size = int(chunk_size * _MAX_SIZE_FACTOR) raw = self._merge_pages( - pages, chunk_size, max_size, chunk_overlap, config, + pages, + chunk_size, + max_size, + chunk_overlap, + config, ) if not raw: @@ -144,7 +152,8 @@ def _build_marker_patterns(config: ProcessingConfig) -> List[re.Pattern]: @staticmethod def _split_into_pages( - text: str, patterns: List[re.Pattern], + text: str, + patterns: List[re.Pattern], ) -> List[Tuple[int, str]]: """ Split *text* by the first matching marker pattern. @@ -256,7 +265,9 @@ def _merge_pages( extreme = int(max_size * 1.5) for chunk in chunks: if len(chunk) > extreme: - sub = self._sub_split_large_chunk(chunk, chunk_size, chunk_overlap, config) + sub = self._sub_split_large_chunk( + chunk, chunk_size, chunk_overlap, config + ) final.extend(sub) else: final.append(chunk) @@ -288,7 +299,10 @@ def _sub_split_large_chunk( Re-split an oversized merged chunk using ProtectedChunkingStrategy so embedded tables/charts remain intact. """ - from contextifier.chunking.strategies.protected_strategy import ProtectedChunkingStrategy + from contextifier.chunking.strategies.protected_strategy import ( + ProtectedChunkingStrategy, + ) + strategy = ProtectedChunkingStrategy() if strategy.can_handle(text, config): result = strategy.chunk(text, config, include_position_metadata=False) @@ -296,6 +310,7 @@ def _sub_split_large_chunk( return result # type: ignore[return-value] # Fallback to plain recursive split from langchain_text_splitters import RecursiveCharacterTextSplitter + splitter = RecursiveCharacterTextSplitter( chunk_size=chunk_size, chunk_overlap=chunk_overlap, diff --git a/contextifier/chunking/strategies/plain_strategy.py b/contextifier/chunking/strategies/plain_strategy.py index 7b93ee1..380aec8 100644 --- a/contextifier/chunking/strategies/plain_strategy.py +++ b/contextifier/chunking/strategies/plain_strategy.py @@ -17,7 +17,7 @@ from __future__ import annotations import re -from typing import Any, List, Optional, Union +from typing import Any, List, Union from langchain_text_splitters import RecursiveCharacterTextSplitter @@ -107,11 +107,15 @@ def priority(self) -> int: @staticmethod def _split_code( - text: str, lang: str, chunk_size: int, chunk_overlap: int, + text: str, + lang: str, + chunk_size: int, + chunk_overlap: int, ) -> List[str]: """Use langchain's language-aware splitter.""" try: from langchain_text_splitters import Language + lang_enum = getattr(Language, lang, None) if lang_enum is not None: splitter = RecursiveCharacterTextSplitter.from_language( @@ -133,7 +137,9 @@ def _split_code( @staticmethod def _split_plain( - text: str, chunk_size: int, chunk_overlap: int, + text: str, + chunk_size: int, + chunk_overlap: int, ) -> List[str]: """ Paragraph-aware recursive splitting. diff --git a/contextifier/chunking/strategies/protected_strategy.py b/contextifier/chunking/strategies/protected_strategy.py index f91252d..6dde85f 100644 --- a/contextifier/chunking/strategies/protected_strategy.py +++ b/contextifier/chunking/strategies/protected_strategy.py @@ -32,8 +32,6 @@ HTML_TABLE_PATTERN, MARKDOWN_TABLE_PATTERN, TEXTBOX_BLOCK_PATTERN, - build_protected_patterns, - TABLE_SIZE_THRESHOLD_MULTIPLIER, ) from contextifier.chunking.table_chunker import chunk_large_table from contextifier.chunking.table_parser import is_markdown_table @@ -120,8 +118,8 @@ def chunk( md_table_positions = self._find_markdown_tables(text, regions) # Build lookup structures - block_regions = [(s, e) for s, e, _ in regions] - region_type_map: Dict[Tuple[int, int], str] = {(s, e): t for s, e, t in regions} + [(s, e) for s, e, _ in regions] + {(s, e): t for s, e, t in regions} # Merge force-chunking table regions with block regions all_typed: List[_Region] = list(regions) @@ -139,7 +137,9 @@ def chunk( merged_typed.append((s, e, t)) all_block = [(s, e) for s, e, _ in merged_typed] - all_type_map: Dict[Tuple[int, int], str] = {(s, e): t for s, e, t in merged_typed} + all_type_map: Dict[Tuple[int, int], str] = { + (s, e): t for s, e, t in merged_typed + } # 4. Walk and split raw = self._walk_and_split( @@ -256,7 +256,8 @@ def _find_protected_regions( @staticmethod def _find_html_tables( - text: str, existing_regions: List[_Region], + text: str, + existing_regions: List[_Region], ) -> List[_Region]: """Find HTML tables not already covered by *existing_regions*.""" existing_set = {(s, e) for s, e, _ in existing_regions} @@ -269,7 +270,8 @@ def _find_html_tables( @staticmethod def _find_markdown_tables( - text: str, existing_regions: List[_Region], + text: str, + existing_regions: List[_Region], ) -> List[_Region]: """Find Markdown tables not already covered by *existing_regions*.""" existing_set = {(s, e) for s, e, _ in existing_regions} @@ -283,16 +285,26 @@ def _find_markdown_tables( @staticmethod def _extract_no_overlap_positions( - text: str, config: ProcessingConfig, + text: str, + config: ProcessingConfig, ) -> List[_Region]: """Page/slide/sheet/chart/metadata tag positions — never overlap.""" tags = config.tags positions: List[_Region] = [] for pat_str, rtype in [ - (rf"{re.escape(tags.page_prefix)}\d+(?:\s*\(OCR(?:\+Ref)?\))?{re.escape(tags.page_suffix)}", "page_tag"), - (rf"{re.escape(tags.slide_prefix)}\d+(?:\s*\(OCR\))?{re.escape(tags.slide_suffix)}", "slide_tag"), - (rf"{re.escape(tags.sheet_prefix)}.+?{re.escape(tags.sheet_suffix)}", "sheet_tag"), + ( + rf"{re.escape(tags.page_prefix)}\d+(?:\s*\(OCR(?:\+Ref)?\))?{re.escape(tags.page_suffix)}", + "page_tag", + ), + ( + rf"{re.escape(tags.slide_prefix)}\d+(?:\s*\(OCR\))?{re.escape(tags.slide_suffix)}", + "slide_tag", + ), + ( + rf"{re.escape(tags.sheet_prefix)}.+?{re.escape(tags.sheet_suffix)}", + "sheet_tag", + ), ]: for m in re.finditer(pat_str, text): positions.append((m.start(), m.end(), rtype)) @@ -315,7 +327,8 @@ def _extract_no_overlap_positions( @staticmethod def _extract_image_positions( - text: str, config: ProcessingConfig, + text: str, + config: ProcessingConfig, ) -> List[Tuple[int, int]]: """Image tag positions — never split mid-tag, never overlap.""" tags = config.tags @@ -373,7 +386,9 @@ def _walk_and_split( # We're at / inside the block if block_size > chunk_size: block_content = text[b_start:b_end].strip() - if force_chunking and self._is_table_type(b_type, block_content): + if force_chunking and self._is_table_type( + b_type, block_content + ): table_chunks = chunk_large_table(block_content, chunk_size) chunks.extend(table_chunks) else: @@ -383,12 +398,18 @@ def _walk_and_split( else: # Block fits → try to extend past it end_pos = min(b_end + (chunk_size - block_size), text_len) - end_pos = self._clamp_to_next_block(end_pos, b_end, block_regions) - end_pos = self._adjust_image_boundary(end_pos, image_positions, text_len) + end_pos = self._clamp_to_next_block( + end_pos, b_end, block_regions + ) + end_pos = self._adjust_image_boundary( + end_pos, image_positions, text_len + ) chunk = text[pos:end_pos].strip() if chunk: chunks.append(chunk) - if self._ends_with_no_overlap(end_pos, no_overlap_positions, image_positions): + if self._ends_with_no_overlap( + end_pos, no_overlap_positions, image_positions + ): pos = end_pos else: pos = max(b_end, end_pos - chunk_overlap) @@ -403,12 +424,18 @@ def _walk_and_split( extra = chunk_size - space_with if extra > 0: end_pos = min(b_end + extra, text_len) - end_pos = self._clamp_to_next_block(end_pos, b_end, block_regions) - end_pos = self._adjust_image_boundary(end_pos, image_positions, text_len) + end_pos = self._clamp_to_next_block( + end_pos, b_end, block_regions + ) + end_pos = self._adjust_image_boundary( + end_pos, image_positions, text_len + ) chunk = text[pos:end_pos].strip() if chunk: chunks.append(chunk) - if self._ends_with_no_overlap(end_pos, no_overlap_positions, image_positions): + if self._ends_with_no_overlap( + end_pos, no_overlap_positions, image_positions + ): pos = end_pos else: pos = max(b_end, end_pos - chunk_overlap) @@ -417,19 +444,29 @@ def _walk_and_split( if space_before > chunk_overlap: # Cut before the block end_pos = b_start - end_pos = self._adjust_image_boundary(end_pos, image_positions, text_len) + end_pos = self._adjust_image_boundary( + end_pos, image_positions, text_len + ) chunk = text[pos:end_pos].strip() if chunk: chunks.append(chunk) - if self._ends_with_no_overlap(end_pos, no_overlap_positions, image_positions): + if self._ends_with_no_overlap( + end_pos, no_overlap_positions, image_positions + ): pos = end_pos else: pos = max(pos + 1, b_start - chunk_overlap) else: # Block is too close → emit block directly block_content = text[b_start:b_end].strip() - if block_size > chunk_size and force_chunking and self._is_table_type(b_type, block_content): - table_chunks = chunk_large_table(block_content, chunk_size) + if ( + block_size > chunk_size + and force_chunking + and self._is_table_type(b_type, block_content) + ): + table_chunks = chunk_large_table( + block_content, chunk_size + ) chunks.extend(table_chunks) else: if block_content: @@ -437,14 +474,20 @@ def _walk_and_split( pos = b_end # No overlap for blocks else: # ── No block in range → find best natural split ─────────── - best_split = self._find_natural_split(text, pos, tentative_end, chunk_size) - best_split = self._adjust_image_boundary(best_split, image_positions, text_len) + best_split = self._find_natural_split( + text, pos, tentative_end, chunk_size + ) + best_split = self._adjust_image_boundary( + best_split, image_positions, text_len + ) chunk = text[pos:best_split].strip() if chunk: chunks.append(chunk) - if self._ends_with_no_overlap(best_split, no_overlap_positions, image_positions): + if self._ends_with_no_overlap( + best_split, no_overlap_positions, image_positions + ): pos = best_split else: pos = max(pos + 1, best_split - chunk_overlap) @@ -457,7 +500,9 @@ def _walk_and_split( @staticmethod def _find_block_in_range( - pos: int, end: int, blocks: List[Tuple[int, int]], + pos: int, + end: int, + blocks: List[Tuple[int, int]], ) -> Optional[Tuple[int, int]]: """Return the first block region overlapping ``[pos, end)``.""" for b_start, b_end in blocks: @@ -467,7 +512,9 @@ def _find_block_in_range( @staticmethod def _clamp_to_next_block( - end_pos: int, after: int, blocks: List[Tuple[int, int]], + end_pos: int, + after: int, + blocks: List[Tuple[int, int]], ) -> int: """Ensure *end_pos* doesn't reach into the next block.""" for b_start, _ in blocks: @@ -514,7 +561,10 @@ def _is_table_type(region_type: str, content: str) -> bool: @staticmethod def _find_natural_split( - text: str, start: int, end: int, chunk_size: int, + text: str, + start: int, + end: int, + chunk_size: int, ) -> int: """Find the best natural split point in ``[start, end)``.""" # Prefer paragraph break diff --git a/contextifier/chunking/strategies/table_strategy.py b/contextifier/chunking/strategies/table_strategy.py index 0f68767..d903262 100644 --- a/contextifier/chunking/strategies/table_strategy.py +++ b/contextifier/chunking/strategies/table_strategy.py @@ -19,7 +19,7 @@ from __future__ import annotations import re -from typing import Any, FrozenSet, List, Optional, Tuple, Union +from typing import Any, List, Optional, Tuple, Union from contextifier.config import ProcessingConfig from contextifier.types import Chunk, ChunkMetadata @@ -96,11 +96,19 @@ def chunk( if sheets: raw = self._chunk_multi_sheet( - sheets, context_prefix, chunk_size, chunk_overlap, config, + sheets, + context_prefix, + chunk_size, + chunk_overlap, + config, ) else: raw = self._chunk_single_table( - body, context_prefix, chunk_size, chunk_overlap, config, + body, + context_prefix, + chunk_size, + chunk_overlap, + config, ) if not raw: @@ -157,7 +165,9 @@ def _chunk_multi_sheet( all_chunks.append(f"{sheet_prefix}\n{seg_content}".strip()) else: table_chunks = chunk_large_table( - seg_content, chunk_size, sheet_prefix, + seg_content, + chunk_size, + sheet_prefix, ) all_chunks.extend(table_chunks) @@ -170,7 +180,9 @@ def _chunk_multi_sheet( if len(sheet_prefix) + len(seg_content) <= chunk_size: all_chunks.append(f"{sheet_prefix}\n{seg_content}".strip()) else: - text_chunks = self._split_plain(seg_content, chunk_size, chunk_overlap) + text_chunks = self._split_plain( + seg_content, chunk_size, chunk_overlap + ) for tc in text_chunks: all_chunks.append(f"{sheet_prefix}\n{tc}".strip()) @@ -199,11 +211,17 @@ def _chunk_single_table( for _, _, table_content in tables: total = len(context_prefix) + len(table_content) if total <= chunk_size: - chunk = f"{context_prefix}\n\n{table_content}".strip() if context_prefix else table_content + chunk = ( + f"{context_prefix}\n\n{table_content}".strip() + if context_prefix + else table_content + ) all_chunks.append(chunk) else: table_chunks = chunk_large_table( - table_content, chunk_size, context_prefix, + table_content, + chunk_size, + context_prefix, ) all_chunks.extend(table_chunks) @@ -213,7 +231,8 @@ def _chunk_single_table( @staticmethod def _extract_metadata_block( - text: str, tags: Any, + text: str, + tags: Any, ) -> Tuple[Optional[str], str]: """Extract ``[Document-Metadata]…[/Document-Metadata]`` block.""" pattern = re.compile( @@ -248,7 +267,8 @@ def _extract_sheets( @staticmethod def _extract_segments( - content: str, config: ProcessingConfig, + content: str, + config: ProcessingConfig, ) -> List[Tuple[str, str]]: """ Extract typed segments (table / textbox / chart / image / text) @@ -264,10 +284,13 @@ def _extract_segments( ) patterns: List[Tuple[str, re.Pattern]] = [ - ("table", re.compile( - r"(?:\[Table\s*\d+\]\s*)?]*>.*?", - re.DOTALL | re.IGNORECASE, - )), + ( + "table", + re.compile( + r"(?:\[Table\s*\d+\]\s*)?]*>.*?", + re.DOTALL | re.IGNORECASE, + ), + ), ("table", MARKDOWN_TABLE_PATTERN), ("textbox", TEXTBOX_BLOCK_PATTERN), ("chart", chart_pattern), diff --git a/contextifier/chunking/table_chunker.py b/contextifier/chunking/table_chunker.py index 160166b..8e6eccb 100644 --- a/contextifier/chunking/table_chunker.py +++ b/contextifier/chunking/table_chunker.py @@ -18,19 +18,16 @@ import re import logging -from typing import List, Optional, Tuple +from typing import List from contextifier.chunking.constants import ( TABLE_WRAPPER_OVERHEAD, - ROW_OVERHEAD, CHUNK_INDEX_OVERHEAD, ParsedTable, - ParsedMarkdownTable, TableRow, ) from contextifier.chunking.table_parser import ( parse_html_table, - extract_cell_spans, has_complex_spans, parse_markdown_table, is_markdown_table, @@ -43,6 +40,7 @@ # HTML Table Chunking # ═══════════════════════════════════════════════════════════════════════════════ + def chunk_html_table( table_html: str, chunk_size: int, @@ -168,7 +166,9 @@ def _split_table_preserving_rowspan( chunks: List[str] = [] for idx, group in enumerate(groups, start=1): - adjusted_rows = [_adjust_row_rowspan(r, len(group), i) for i, r in enumerate(group)] + adjusted_rows = [ + _adjust_row_rowspan(r, len(group), i) for i, r in enumerate(group) + ] body_rows = "\n".join(adjusted_rows) header_part = f"\n{parsed.header_html}" if parsed.header_html else "" table_str = f"{header_part}\n{body_rows}\n
" @@ -191,13 +191,16 @@ def _clamp_rowspan(match: re.Match) -> str: return "" # Remove rowspan="1" return f'rowspan="{clamped}"' - return re.sub(r'rowspan\s*=\s*["\']?(\d+)["\']?', _clamp_rowspan, html, flags=re.IGNORECASE) + return re.sub( + r'rowspan\s*=\s*["\']?(\d+)["\']?', _clamp_rowspan, html, flags=re.IGNORECASE + ) # ═══════════════════════════════════════════════════════════════════════════════ # Markdown Table Chunking # ═══════════════════════════════════════════════════════════════════════════════ + def chunk_markdown_table( table_text: str, chunk_size: int, @@ -269,6 +272,7 @@ def chunk_markdown_table( # Unified Entry Point # ═══════════════════════════════════════════════════════════════════════════════ + def chunk_large_table( table_text: str, chunk_size: int, diff --git a/contextifier/chunking/table_parser.py b/contextifier/chunking/table_parser.py index 25ed747..28a0c99 100644 --- a/contextifier/chunking/table_parser.py +++ b/contextifier/chunking/table_parser.py @@ -32,6 +32,7 @@ # HTML Table Parsing # ═══════════════════════════════════════════════════════════════════════════════ + def parse_html_table(html: str) -> ParsedTable: """ Parse an HTML table string into a structured ParsedTable. @@ -49,16 +50,22 @@ def parse_html_table(html: str) -> ParsedTable: data_rows: List[TableRow] = [] # Extract rows from section - thead_match = re.search(r"]*>(.*?)", html, re.DOTALL | re.IGNORECASE) + thead_match = re.search( + r"]*>(.*?)", html, re.DOTALL | re.IGNORECASE + ) thead_rows_html: List[str] = [] if thead_match: - thead_rows_html = re.findall(r"]*>.*?", thead_match.group(1), re.DOTALL | re.IGNORECASE) + thead_rows_html = re.findall( + r"]*>.*?", thead_match.group(1), re.DOTALL | re.IGNORECASE + ) # Extract ALL rows from the full table all_rows_html = re.findall(r"]*>.*?", html, re.DOTALL | re.IGNORECASE) for row_html in all_rows_html: - cells = re.findall(r"]*>.*?", row_html, re.DOTALL | re.IGNORECASE) + cells = re.findall( + r"]*>.*?", row_html, re.DOTALL | re.IGNORECASE + ) cell_count = len(cells) char_length = len(row_html) @@ -134,6 +141,7 @@ def has_complex_spans(html: str) -> bool: # Markdown Table Parsing # ═══════════════════════════════════════════════════════════════════════════════ + def parse_markdown_table(text: str) -> Optional[ParsedMarkdownTable]: """ Parse a Markdown pipe-table into a structured ParsedMarkdownTable. diff --git a/contextifier/config.py b/contextifier/config.py index 468837e..2393a77 100644 --- a/contextifier/config.py +++ b/contextifier/config.py @@ -27,7 +27,7 @@ import types from dataclasses import dataclass, field, replace -from typing import Any, Dict, FrozenSet, Literal, Optional +from typing import Any, Dict, Literal, Optional from contextifier.errors import ConfigurationError from contextifier.types import ( @@ -39,6 +39,7 @@ # ─── Tag Configuration ──────────────────────────────────────────────────────── + @dataclass(frozen=True) class TagConfig: """ @@ -47,6 +48,7 @@ class TagConfig: Controls the prefix/suffix wrapping for page numbers, slide numbers, sheet names, image references, and chart blocks. """ + # Page tags: e.g., "[Page Number: 1]" page_prefix: str = "[Page Number: " page_suffix: str = "]" @@ -74,11 +76,13 @@ class TagConfig: # ─── Image Configuration ────────────────────────────────────────────────────── + @dataclass(frozen=True) class ImageConfig: """ Configuration for image extraction and storage. """ + directory_path: str = "temp/images" naming_strategy: NamingStrategy = NamingStrategy.HASH default_format: str = "png" @@ -90,11 +94,13 @@ class ImageConfig: # ─── Chart Configuration ────────────────────────────────────────────────────── + @dataclass(frozen=True) class ChartConfig: """ Configuration for chart data formatting. """ + use_html_table: bool = True include_chart_type: bool = True include_chart_title: bool = True @@ -102,23 +108,27 @@ class ChartConfig: # ─── Metadata Configuration ────────────────────────────────────────────────── + @dataclass(frozen=True) class MetadataConfig: """ Configuration for metadata formatting. """ - language: str = "ko" # "ko" or "en" + + language: str = "ko" # "ko" or "en" date_format: str = "%Y-%m-%d %H:%M:%S" indent: str = " " # ─── Table Configuration ───────────────────────────────────────────────────── + @dataclass(frozen=True) class TableConfig: """ Configuration for table output formatting. """ + output_format: OutputFormat = OutputFormat.HTML clean_whitespace: bool = True preserve_merged_cells: bool = True @@ -126,11 +136,13 @@ class TableConfig: # ─── Chunking Configuration ────────────────────────────────────────────────── + @dataclass(frozen=True) class ChunkingConfig: """ Configuration for text chunking. """ + chunk_size: int = 1000 chunk_overlap: int = 200 preserve_tables: bool = True @@ -149,19 +161,24 @@ def __post_init__(self) -> None: # ─── OCR Configuration ─────────────────────────────────────────────────────── + @dataclass(frozen=True) class OCRConfig: """ Configuration for OCR processing. """ + enabled: bool = False - provider: Optional[str] = None # "openai", "anthropic", "gemini", "vllm", "bedrock", "tesseract" - prompt: Optional[str] = None # Custom OCR prompt (None = use default) - prompt_language: str = "ko" # Prompt output language: "ko", "en", etc. + provider: Optional[str] = ( + None # "openai", "anthropic", "gemini", "vllm", "bedrock", "tesseract" + ) + prompt: Optional[str] = None # Custom OCR prompt (None = use default) + prompt_language: str = "ko" # Prompt output language: "ko", "en", etc. # ─── Encoding Configuration ────────────────────────────────────────────────── + @dataclass(frozen=True) class EncodingConfig: """ @@ -179,8 +196,14 @@ class EncodingConfig: detection. If confidence is below this, fallback_encodings are tried instead. (Reserved for future chardet integration.) """ + fallback_encodings: tuple[str, ...] = ( - "utf-8", "utf-8-sig", "cp949", "euc-kr", "latin-1", "ascii", + "utf-8", + "utf-8-sig", + "cp949", + "euc-kr", + "latin-1", + "ascii", ) force_encoding: Optional[str] = None min_confidence: float = 0.7 @@ -188,6 +211,7 @@ class EncodingConfig: # ─── Root Configuration ────────────────────────────────────────────────────── + @dataclass(frozen=True) class ProcessingConfig: """ @@ -210,6 +234,7 @@ class ProcessingConfig: # Modify existing config config2 = config.with_tags(page_prefix="") """ + tags: TagConfig = field(default_factory=TagConfig) images: ImageConfig = field(default_factory=ImageConfig) charts: ChartConfig = field(default_factory=ChartConfig) @@ -225,7 +250,11 @@ class ProcessingConfig: def __post_init__(self) -> None: # Deep-freeze format_options so external code cannot mutate # the internal state of this frozen dataclass. - raw = self.format_options if isinstance(self.format_options, dict) else dict(self.format_options) + raw = ( + self.format_options + if isinstance(self.format_options, dict) + else dict(self.format_options) + ) frozen = types.MappingProxyType( {k: types.MappingProxyType(dict(v)) for k, v in raw.items()} ) @@ -282,7 +311,8 @@ def get_format_option(self, format_name: str, key: str, default: Any = None) -> def to_dict(self) -> Dict[str, Any]: """Serialize config to dictionary.""" - from dataclasses import asdict, fields + from dataclasses import asdict + # asdict() fails on MappingProxyType (not deepcopy-able). # Temporarily swap format_options to a plain dict for serialization. plain_opts = dict(self.format_options) @@ -290,11 +320,11 @@ def to_dict(self) -> Dict[str, Any]: try: d = asdict(self) finally: - object.__setattr__(self, "format_options", types.MappingProxyType(plain_opts)) + object.__setattr__( + self, "format_options", types.MappingProxyType(plain_opts) + ) # Ensure format_options values are plain dicts - d["format_options"] = { - k: dict(v) for k, v in plain_opts.items() - } + d["format_options"] = {k: dict(v) for k, v in plain_opts.items()} return d @classmethod @@ -302,23 +332,48 @@ def from_dict(cls, data: Dict[str, Any]) -> "ProcessingConfig": """Deserialize config from dictionary.""" tags = TagConfig(**data.get("tags", {})) if "tags" in data else TagConfig() images_data = data.get("images", {}) - if "naming_strategy" in images_data and isinstance(images_data["naming_strategy"], str): - images_data["naming_strategy"] = NamingStrategy(images_data["naming_strategy"]) - if "storage_type" in images_data and isinstance(images_data["storage_type"], str): + if "naming_strategy" in images_data and isinstance( + images_data["naming_strategy"], str + ): + images_data["naming_strategy"] = NamingStrategy( + images_data["naming_strategy"] + ) + if "storage_type" in images_data and isinstance( + images_data["storage_type"], str + ): images_data["storage_type"] = StorageType(images_data["storage_type"]) images = ImageConfig(**images_data) if "images" in data else ImageConfig() - charts = ChartConfig(**data.get("charts", {})) if "charts" in data else ChartConfig() - meta = MetadataConfig(**data.get("metadata", {})) if "metadata" in data else MetadataConfig() + charts = ( + ChartConfig(**data.get("charts", {})) if "charts" in data else ChartConfig() + ) + meta = ( + MetadataConfig(**data.get("metadata", {})) + if "metadata" in data + else MetadataConfig() + ) tables_data = data.get("tables", {}) - if "output_format" in tables_data and isinstance(tables_data["output_format"], str): + if "output_format" in tables_data and isinstance( + tables_data["output_format"], str + ): tables_data["output_format"] = OutputFormat(tables_data["output_format"]) tables = TableConfig(**tables_data) if "tables" in data else TableConfig() - chunking = ChunkingConfig(**data.get("chunking", {})) if "chunking" in data else ChunkingConfig() + chunking = ( + ChunkingConfig(**data.get("chunking", {})) + if "chunking" in data + else ChunkingConfig() + ) ocr = OCRConfig(**data.get("ocr", {})) if "ocr" in data else OCRConfig() encoding_data = data.get("encoding", {}) - if "fallback_encodings" in encoding_data and isinstance(encoding_data["fallback_encodings"], list): - encoding_data = {**encoding_data, "fallback_encodings": tuple(encoding_data["fallback_encodings"])} - encoding = EncodingConfig(**encoding_data) if "encoding" in data else EncodingConfig() + if "fallback_encodings" in encoding_data and isinstance( + encoding_data["fallback_encodings"], list + ): + encoding_data = { + **encoding_data, + "fallback_encodings": tuple(encoding_data["fallback_encodings"]), + } + encoding = ( + EncodingConfig(**encoding_data) if "encoding" in data else EncodingConfig() + ) fmt_opts = data.get("format_options", {}) return cls( diff --git a/contextifier/document_processor.py b/contextifier/document_processor.py index 0c0ac44..3a90709 100644 --- a/contextifier/document_processor.py +++ b/contextifier/document_processor.py @@ -44,7 +44,6 @@ from __future__ import annotations -import io import logging import os from dataclasses import dataclass, field @@ -61,8 +60,6 @@ from contextifier.errors import ( FileNotFoundError as ContextifyFileNotFoundError, UnsupportedFormatError, - HandlerNotFoundError, - ContextifierError, ) from contextifier.handlers.registry import HandlerRegistry from contextifier.chunking.chunker import TextChunker @@ -80,6 +77,7 @@ # ── ChunkResult ────────────────────────────────────────────────────────── + @dataclass class ChunkResult: """ @@ -97,7 +95,9 @@ class ChunkResult: @property def has_metadata(self) -> bool: """Whether position metadata is available.""" - return self.chunks_with_metadata is not None and len(self.chunks_with_metadata) > 0 + return ( + self.chunks_with_metadata is not None and len(self.chunks_with_metadata) > 0 + ) def __len__(self) -> int: return len(self.chunks) @@ -154,6 +154,7 @@ def save_to_md( # ── DocumentProcessor ──────────────────────────────────────────────────── + class DocumentProcessor: """ Main entry point for document processing. @@ -201,6 +202,7 @@ def __init__( self._ocr_processor = None if self._ocr_engine is not None: from contextifier.ocr.processor import OCRProcessor + self._ocr_processor = OCRProcessor( engine=self._ocr_engine, config=self._config, @@ -260,7 +262,9 @@ def extract_text( # Build FileContext file_context = self._create_file_context( - file_path_str, ext, max_file_size=self._MAX_FILE_SIZE, + file_path_str, + ext, + max_file_size=self._MAX_FILE_SIZE, password=password, ) @@ -330,7 +334,9 @@ def process( ) file_context = self._create_file_context( - file_path_str, ext, max_file_size=self._MAX_FILE_SIZE, + file_path_str, + ext, + max_file_size=self._MAX_FILE_SIZE, password=password, ) @@ -354,6 +360,30 @@ def process( # Public API — Chunking # ═══════════════════════════════════════════════════════════════════════ + def open_raw(self, file_path, *, file_extension=None): + """Open a document for lossless, writable access (raw twin of process()). + + Unlike :meth:`extract_text` — which renders an AI-friendly view and + discards the rest — this returns a raw document model over the full + OPC package: every part stays available, edits are surgical, and + saving keeps untouched parts byte-identical. + + Args: + file_path: path, bytes, or binary stream of the document. + file_extension: override format detection (e.g. "xlsx"). + + Returns: + XlsxRawDocument | DocxRawDocument | PptxRawDocument + (see :mod:`contextifier.raw`). + + Raises: + RawUnsupportedError: the format has no raw model yet + (supported today: xlsx, docx, pptx). + """ + from contextifier.raw import open_raw as _open_raw + + return _open_raw(file_path, extension=file_extension) + def chunk_text( self, text: str, @@ -447,7 +477,11 @@ def extract_chunks( # Build ChunkResult strategy_name = self._chunker.last_strategy_name - if include_position_metadata and raw_chunks and isinstance(raw_chunks[0], Chunk): + if ( + include_position_metadata + and raw_chunks + and isinstance(raw_chunks[0], Chunk) + ): return ChunkResult( chunks=[c.text for c in raw_chunks], # type: ignore chunks_with_metadata=raw_chunks, # type: ignore @@ -548,7 +582,10 @@ def _resolve_extension(file_path: str, override: Optional[str]) -> str: @staticmethod def _create_file_context( - file_path: str, extension: str, *, max_file_size: int = 0, + file_path: str, + extension: str, + *, + max_file_size: int = 0, password: Optional[str] = None, ) -> FileContext: """Create a standardised FileContext dict from a file path. @@ -563,17 +600,22 @@ def _create_file_context( file_size = os.path.getsize(file_path) if max_file_size and file_size > max_file_size: from contextifier.errors import FileReadError + raise FileReadError( f"File size ({file_size:,} bytes) exceeds the limit " f"({max_file_size:,} bytes). Consider processing in " f"streaming mode or increase the limit.", - context={"file_path": file_path, "file_size": file_size, - "max_file_size": max_file_size}, + context={ + "file_path": file_path, + "file_size": file_size, + "max_file_size": max_file_size, + }, ) file_data = Path(file_path).read_bytes() # Decrypt password-protected MS Office files from contextifier.services.crypto_service import decrypt_if_encrypted + file_data = decrypt_if_encrypted(file_data, password=password) return FileContext( diff --git a/contextifier/errors.py b/contextifier/errors.py index 8a59714..d7982a2 100644 --- a/contextifier/errors.py +++ b/contextifier/errors.py @@ -88,35 +88,43 @@ def with_context(self, **kwargs: Any) -> "ContextifierError": # ─── Configuration Errors ───────────────────────────────────────────────────── + class ConfigurationError(ContextifierError): """Invalid or missing configuration.""" + pass # ─── File Errors ────────────────────────────────────────────────────────────── + class FileError(ContextifierError): """Base class for file-related errors.""" + pass class FileNotFoundError(FileError): """File does not exist at the specified path.""" + pass class FileReadError(FileError): """Cannot read the file (permission, corruption, etc.).""" + pass class UnsupportedFormatError(FileError): """File extension is not supported by any registered handler.""" + pass # ─── Pipeline Errors ───────────────────────────────────────────────────────── + class PipelineError(ContextifierError): """Base class for processing pipeline failures.""" @@ -138,67 +146,82 @@ def __init__( class ConversionError(PipelineError): """Failed to convert binary data to format-specific object.""" + pass class PreprocessingError(PipelineError): """Failed during preprocessing stage.""" + pass class ExtractionError(PipelineError): """Failed during content or metadata extraction.""" + pass class PostprocessingError(PipelineError): """Failed during postprocessing / final assembly.""" + pass # ─── Handler Errors ────────────────────────────────────────────────────────── + class HandlerError(ContextifierError): """Base class for handler-level failures.""" + pass class HandlerNotFoundError(HandlerError): """No handler registered for the requested file extension.""" + pass class HandlerExecutionError(HandlerError): """Handler encountered an error during execution.""" + pass # ─── Service Errors ────────────────────────────────────────────────────────── + class ServiceError(ContextifierError): """Base class for service-level failures.""" + pass class ImageServiceError(ServiceError): """Image processing or storage failure.""" + pass class StorageError(ServiceError): """Storage backend operation failure.""" + pass class OCRError(ServiceError): """OCR processing failure.""" + pass # ─── Chunking Errors ───────────────────────────────────────────────────────── + class ChunkingError(ContextifierError): """Text chunking operation failure.""" + pass diff --git a/contextifier/handlers/base.py b/contextifier/handlers/base.py index d09801c..fee063e 100644 --- a/contextifier/handlers/base.py +++ b/contextifier/handlers/base.py @@ -333,9 +333,11 @@ def _get_timeout_executor(cls) -> concurrent.futures.ThreadPoolExecutor: if BaseHandler._timeout_executor is None: with BaseHandler._timeout_executor_lock: if BaseHandler._timeout_executor is None: - BaseHandler._timeout_executor = concurrent.futures.ThreadPoolExecutor( - max_workers=4, - thread_name_prefix="handler-timeout", + BaseHandler._timeout_executor = ( + concurrent.futures.ThreadPoolExecutor( + max_workers=4, + thread_name_prefix="handler-timeout", + ) ) atexit.register(BaseHandler._shutdown_timeout_executor) return BaseHandler._timeout_executor @@ -359,9 +361,7 @@ def _execute_pipeline( # Stage 0: Delegation check (e.g., .doc that is actually RTF) delegation_result = self._check_delegation(file_context, **kwargs) if delegation_result is not None: - self._logger.info( - f"{self.handler_name} delegated to another handler" - ) + self._logger.info(f"{self.handler_name} delegated to another handler") return delegation_result converted = None @@ -385,9 +385,7 @@ def _execute_pipeline( metadata = None if include_metadata: try: - metadata = self._metadata_extractor.extract( - preprocessed.content - ) + metadata = self._metadata_extractor.extract(preprocessed.content) except Exception as e: self._logger.warning(f"Metadata extraction failed: {e}") metadata = None @@ -412,7 +410,12 @@ def _execute_pipeline( result.text = final_text return result - except (ConversionError, PreprocessingError, ExtractionError, PostprocessingError): + except ( + ConversionError, + PreprocessingError, + ExtractionError, + PostprocessingError, + ): raise except Exception as e: raise HandlerExecutionError( diff --git a/contextifier/handlers/csv/content_extractor.py b/contextifier/handlers/csv/content_extractor.py index adde1df..eece414 100644 --- a/contextifier/handlers/csv/content_extractor.py +++ b/contextifier/handlers/csv/content_extractor.py @@ -35,7 +35,7 @@ from contextifier.handlers.csv.preprocessor import CsvParsedData if TYPE_CHECKING: - from contextifier.services.table_service import TableService + pass class CsvContentExtractor(BaseContentExtractor): @@ -167,14 +167,16 @@ def _build_simple( for col_idx in range(col_count): content = row[col_idx].strip() if col_idx < len(row) else "" - cells.append(TableCell( - content=content, - row_span=1, - col_span=1, - is_header=is_header_row, - row_index=row_idx, - col_index=col_idx, - )) + cells.append( + TableCell( + content=content, + row_span=1, + col_span=1, + is_header=is_header_row, + row_index=row_idx, + col_index=col_idx, + ) + ) table_rows.append(cells) return TableData( @@ -207,12 +209,14 @@ def _build_with_merges( row_info: List[Dict[str, Any]] = [] for c_idx in range(col_count): value = row[c_idx].strip() if c_idx < len(row) else "" - row_info.append({ - "value": value, - "colspan": 1, - "rowspan": 1, - "skip": False, - }) + row_info.append( + { + "value": value, + "colspan": 1, + "rowspan": 1, + "skip": False, + } + ) grid.append(row_info) # Pass 1: Horizontal merge (colspan) @@ -272,14 +276,16 @@ def _build_with_merges( if cell["skip"]: continue - cells.append(TableCell( - content=cell["value"], - row_span=cell["rowspan"], - col_span=cell["colspan"], - is_header=is_header_row, - row_index=r_idx, - col_index=c_idx, - )) + cells.append( + TableCell( + content=cell["value"], + row_span=cell["rowspan"], + col_span=cell["colspan"], + is_header=is_header_row, + row_index=r_idx, + col_index=c_idx, + ) + ) table_rows.append(cells) return TableData( diff --git a/contextifier/handlers/csv/converter.py b/contextifier/handlers/csv/converter.py index 2c61110..5871ca8 100644 --- a/contextifier/handlers/csv/converter.py +++ b/contextifier/handlers/csv/converter.py @@ -20,12 +20,10 @@ from __future__ import annotations -import logging from typing import Any, List, NamedTuple, Optional from contextifier.pipeline.converter import BaseConverter from contextifier.types import FileContext -from contextifier.errors import ConversionError class CsvConvertedData(NamedTuple): @@ -34,6 +32,7 @@ class CsvConvertedData(NamedTuple): Carries decoded text and metadata needed by downstream stages. """ + text: str encoding: str file_extension: str @@ -45,9 +44,9 @@ class CsvConvertedData(NamedTuple): _BOM_TABLE: List[tuple[bytes, str]] = [ (b"\xff\xfe\x00\x00", "utf-32-le"), (b"\x00\x00\xfe\xff", "utf-32-be"), - (b"\xef\xbb\xbf", "utf-8-sig"), - (b"\xff\xfe", "utf-16-le"), - (b"\xfe\xff", "utf-16-be"), + (b"\xef\xbb\xbf", "utf-8-sig"), + (b"\xff\xfe", "utf-16-le"), + (b"\xfe\xff", "utf-16-be"), ] # Default encoding priority for CSV/TSV files. @@ -129,7 +128,8 @@ def convert( text = file_data.decode(bom_encoding) self._logger.debug( "BOM detected: %s, decoded %d bytes", - bom_encoding, len(file_data), + bom_encoding, + len(file_data), ) return CsvConvertedData( text=text, @@ -137,7 +137,9 @@ def convert( file_extension=file_ext, ) except (UnicodeDecodeError, LookupError): - self._logger.debug("BOM %s detected but decode failed", bom_encoding) + self._logger.debug( + "BOM %s detected but decode failed", bom_encoding + ) # Phase 2: Try encoding candidates extra: List[str] = kwargs.get("encodings", []) diff --git a/contextifier/handlers/csv/metadata_extractor.py b/contextifier/handlers/csv/metadata_extractor.py index 2ae8bf9..89da965 100644 --- a/contextifier/handlers/csv/metadata_extractor.py +++ b/contextifier/handlers/csv/metadata_extractor.py @@ -20,8 +20,7 @@ from __future__ import annotations -import logging -from typing import Any, Dict, List +from typing import Any, Dict from contextifier.pipeline.metadata_extractor import BaseMetadataExtractor from contextifier.types import DocumentMetadata @@ -30,10 +29,10 @@ # Delimiter display names (human-readable) _DELIMITER_NAMES: Dict[str, str] = { - ",": "Comma (,)", + ",": "Comma (,)", "\t": "Tab (\\t)", - ";": "Semicolon (;)", - "|": "Pipe (|)", + ";": "Semicolon (;)", + "|": "Pipe (|)", } # Maximum number of column names to include in metadata diff --git a/contextifier/handlers/csv/preprocessor.py b/contextifier/handlers/csv/preprocessor.py index 57761d7..9223710 100644 --- a/contextifier/handlers/csv/preprocessor.py +++ b/contextifier/handlers/csv/preprocessor.py @@ -53,6 +53,7 @@ # ── Parsed data structure ──────────────────────────────────────────────── + class CsvParsedData(NamedTuple): """ Structured output of the CsvPreprocessor. @@ -60,6 +61,7 @@ class CsvParsedData(NamedTuple): Stored in PreprocessedData.content and consumed by CsvMetadataExtractor and CsvContentExtractor. """ + rows: List[List[str]] has_header: bool delimiter: str @@ -72,6 +74,7 @@ class CsvParsedData(NamedTuple): # ── Preprocessor ───────────────────────────────────────────────────────── + class CsvPreprocessor(BasePreprocessor): """ Preprocessor for CSV and TSV files. @@ -120,10 +123,7 @@ def preprocess( text = text.replace("\r\n", "\n").replace("\r", "\n") # Determine delimiter - delimiter = ( - kwargs.get("delimiter") - or self._default_delimiter - ) + delimiter = kwargs.get("delimiter") or self._default_delimiter delimiter_confidence = 1.0 # forced delimiter = full confidence if delimiter is None: delimiter, delimiter_confidence = _detect_delimiter( @@ -318,7 +318,8 @@ def _parse_csv_content( for i, row in enumerate(reader): if i >= max_rows: _logger.warning( - "CSV row limit reached: %d", max_rows, + "CSV row limit reached: %d", + max_rows, ) truncated = True break @@ -389,12 +390,8 @@ def _detect_header(rows: List[List[str]]) -> bool: first_row = rows[0] second_row = rows[1] - first_all_text = all( - not _is_numeric(cell) for cell in first_row if cell.strip() - ) - second_has_numbers = any( - _is_numeric(cell) for cell in second_row if cell.strip() - ) + first_all_text = all(not _is_numeric(cell) for cell in first_row if cell.strip()) + second_has_numbers = any(_is_numeric(cell) for cell in second_row if cell.strip()) first_unique = len(set(first_row)) == len(first_row) return first_all_text and (second_has_numbers or first_unique) @@ -402,12 +399,12 @@ def _detect_header(rows: List[List[str]]) -> bool: # Compiled numeric patterns for performance _NUMERIC_PATTERNS = [ - re.compile(r"^-?\d+$"), # Integer - re.compile(r"^-?\d+\.\d+$"), # Float - re.compile(r"^-?\d{1,3}(,\d{3})*(\.\d+)?$"), # Thousands separators - re.compile(r"^-?\d+(\.\d+)?%$"), # Percentage - re.compile(r"^\$-?\d+(\.\d+)?$"), # USD - re.compile(r"^₩-?\d+(,\d{3})*$"), # KRW + re.compile(r"^-?\d+$"), # Integer + re.compile(r"^-?\d+\.\d+$"), # Float + re.compile(r"^-?\d{1,3}(,\d{3})*(\.\d+)?$"), # Thousands separators + re.compile(r"^-?\d+(\.\d+)?%$"), # Percentage + re.compile(r"^\$-?\d+(\.\d+)?$"), # USD + re.compile(r"^₩-?\d+(,\d{3})*$"), # KRW ] diff --git a/contextifier/handlers/doc/_constants.py b/contextifier/handlers/doc/_constants.py index a9b7e52..2ea42f1 100644 --- a/contextifier/handlers/doc/_constants.py +++ b/contextifier/handlers/doc/_constants.py @@ -28,27 +28,29 @@ TABLE_STREAM_NAMES: tuple[str, ...] = ("1Table", "0Table") # Streams that may contain embedded images -IMAGE_STREAM_KEYWORDS: frozenset[str] = frozenset({ - "pictures", - "data", - "object", - "oleobject", - "objectpool", -}) +IMAGE_STREAM_KEYWORDS: frozenset[str] = frozenset( + { + "pictures", + "data", + "object", + "oleobject", + "objectpool", + } +) # ── Image format detection ──────────────────────────────────────────────── # Binary signatures for detecting image format in OLE streams. # Mapping: format_name → (header_bytes, min_data_length) IMAGE_SIGNATURES: dict[str, tuple[bytes, int]] = { - "png": (b"\x89PNG\r\n\x1a\n", 8), - "jpeg": (b"\xff\xd8", 2), - "gif87": (b"GIF87a", 6), - "gif89": (b"GIF89a", 6), - "bmp": (b"BM", 2), - "tiff_le": (b"II\x2a\x00", 4), - "tiff_be": (b"MM\x00\x2a", 4), - "emf": (b"\x01\x00\x00\x00", 4), # EMF header (Enhanced Metafile) + "png": (b"\x89PNG\r\n\x1a\n", 8), + "jpeg": (b"\xff\xd8", 2), + "gif87": (b"GIF87a", 6), + "gif89": (b"GIF89a", 6), + "bmp": (b"BM", 2), + "tiff_le": (b"II\x2a\x00", 4), + "tiff_be": (b"MM\x00\x2a", 4), + "emf": (b"\x01\x00\x00\x00", 4), # EMF header (Enhanced Metafile) } # ── Encoding fallback list ──────────────────────────────────────────────── @@ -76,8 +78,8 @@ # AC00-D7AF: Hangul Syllables # 3000-4DFF: CJK Symbols, Hiragana, Katakana, Bopomofo, CJK Unified CJK_HIGH_BYTE_RANGES: tuple[tuple[int, int], ...] = ( - (0xAC, 0xD7), # Hangul Syllables - (0x30, 0x4E), # CJK range (partial) + (0xAC, 0xD7), # Hangul Syllables + (0x30, 0x4E), # CJK range (partial) ) diff --git a/contextifier/handlers/doc/_fib.py b/contextifier/handlers/doc/_fib.py index b791544..be4c16e 100644 --- a/contextifier/handlers/doc/_fib.py +++ b/contextifier/handlers/doc/_fib.py @@ -21,20 +21,20 @@ import re import struct import logging -from typing import List, NamedTuple, Optional, Tuple +from typing import List, NamedTuple, Optional logger = logging.getLogger(__name__) # ── FIB well-known offsets (Word 97+ / nFib >= 0x00C1) ────────────────────── # FibBase fields -_OFFSET_WIDENT = 0x0000 # uint16 — magic (0xA5EC / 0xA5DC) -_OFFSET_NFIB = 0x0002 # uint16 — version number -_OFFSET_FLAGS = 0x000A # uint16 — flags (bit 9 = fWhichTblStm) +_OFFSET_WIDENT = 0x0000 # uint16 — magic (0xA5EC / 0xA5DC) +_OFFSET_NFIB = 0x0002 # uint16 — version number +_OFFSET_FLAGS = 0x000A # uint16 — flags (bit 9 = fWhichTblStm) # FibRgLw97 fields (absolute offsets assuming standard csw=14, cslw=22) -_OFFSET_CCPTEXT = 0x004C # uint32 — character count of main body text -_OFFSET_CCPFTN = 0x0050 # uint32 — character count of footnote text +_OFFSET_CCPTEXT = 0x004C # uint32 — character count of main body text +_OFFSET_CCPFTN = 0x0050 # uint32 — character count of footnote text # FibRgFcLcb97: fcClx / lcbClx # fcClx is at absolute offset 0x01A2 (entry index 33 of FibRgFcLcb97) @@ -43,18 +43,19 @@ _OFFSET_LCB_CLX = 0x01A6 # Minimum FIB size to attempt piece-table parsing -_MIN_FIB_SIZE = 0x01AA # must be able to read lcbClx +_MIN_FIB_SIZE = 0x01AA # must be able to read lcbClx # nFib threshold — piece table parsing only reliable for Word 97+ -_MIN_NFIB_FOR_PIECE_TABLE = 0x00C1 # Word 97 +_MIN_NFIB_FOR_PIECE_TABLE = 0x00C1 # Word 97 class PieceDescriptor(NamedTuple): """A single entry from the PlcPcd piece table.""" - cp_start: int # Starting character position (in document order) - cp_end: int # Ending character position (exclusive) - fc: int # File Character offset in WordDocument stream - is_compressed: bool # True → cp1252 (1 byte/char), False → UTF-16LE + + cp_start: int # Starting character position (in document order) + cp_end: int # Ending character position (exclusive) + fc: int # File Character offset in WordDocument stream + is_compressed: bool # True → cp1252 (1 byte/char), False → UTF-16LE def parse_fib_text( @@ -73,7 +74,9 @@ def parse_fib_text( (caller should fall back to heuristic extraction). """ if len(word_data) < _MIN_FIB_SIZE: - logger.debug("WordDocument stream too short for FIB parsing (%d bytes)", len(word_data)) + logger.debug( + "WordDocument stream too short for FIB parsing (%d bytes)", len(word_data) + ) return None if table_stream is None or len(table_stream) == 0: @@ -103,7 +106,9 @@ def parse_fib_text( if fc_clx + lcb_clx > len(table_stream): logger.debug( "Clx extends beyond table stream (fcClx=%d, lcbClx=%d, stream=%d)", - fc_clx, lcb_clx, len(table_stream), + fc_clx, + lcb_clx, + len(table_stream), ) return None @@ -152,7 +157,9 @@ def _parse_clx(clx_data: bytes) -> List[PieceDescriptor]: break else: # Unknown type — abort - logger.debug("Unknown Clx entry type: 0x%02X at offset %d", clx_data[offset], offset) + logger.debug( + "Unknown Clx entry type: 0x%02X at offset %d", clx_data[offset], offset + ) return [] if offset >= length or clx_data[offset] != 0x02: @@ -216,12 +223,14 @@ def _parse_plc_pcd(plc_data: bytes) -> List[PieceDescriptor]: # ANSI: actual byte offset is fc / 2 fc = fc // 2 - pieces.append(PieceDescriptor( - cp_start=cps[i], - cp_end=cps[i + 1], - fc=fc, - is_compressed=is_compressed, - )) + pieces.append( + PieceDescriptor( + cp_start=cps[i], + cp_end=cps[i + 1], + fc=fc, + is_compressed=is_compressed, + ) + ) return pieces @@ -256,7 +265,9 @@ def _read_pieces( if byte_offset + byte_len > stream_len: logger.debug( "Piece extends beyond stream (offset=%d, len=%d, stream=%d)", - byte_offset, byte_len, stream_len, + byte_offset, + byte_len, + stream_len, ) continue raw = word_data[byte_offset : byte_offset + byte_len] @@ -271,7 +282,9 @@ def _read_pieces( if byte_offset + byte_len > stream_len: logger.debug( "Piece extends beyond stream (offset=%d, len=%d, stream=%d)", - byte_offset, byte_len, stream_len, + byte_offset, + byte_len, + stream_len, ) continue raw = word_data[byte_offset : byte_offset + byte_len] diff --git a/contextifier/handlers/doc/content_extractor.py b/contextifier/handlers/doc/content_extractor.py index d60a070..2620b1d 100644 --- a/contextifier/handlers/doc/content_extractor.py +++ b/contextifier/handlers/doc/content_extractor.py @@ -32,7 +32,7 @@ import re import logging -from typing import Any, Dict, List, Optional, Set +from typing import Any, List, Optional, Set from contextifier.pipeline.content_extractor import BaseContentExtractor from contextifier.types import PreprocessedData, TableData, TableCell @@ -64,7 +64,8 @@ def __init__(self, **kwargs: Any) -> None: self._min_unicode_bytes = MIN_UNICODE_BYTES if self._config is not None: frag_len = self._config.get_format_option( - "doc", "min_text_fragment_length", + "doc", + "min_text_fragment_length", self._min_text_fragment_length, ) self._min_text_fragment_length = int(frag_len) @@ -146,19 +147,23 @@ def extract_tables( for row_idx, row_cells in enumerate(raw_table): row: List[TableCell] = [] for col_idx, cell_text in enumerate(row_cells): - row.append(TableCell( - content=cell_text, - row_index=row_idx, - col_index=col_idx, - )) + row.append( + TableCell( + content=cell_text, + row_index=row_idx, + col_index=col_idx, + ) + ) rows.append(row) max_cols = max(max_cols, len(row_cells)) - tables.append(TableData( - rows=rows, - num_rows=len(rows), - num_cols=max_cols, - )) + tables.append( + TableData( + rows=rows, + num_rows=len(rows), + num_cols=max_cols, + ) + ) return tables @@ -224,15 +229,20 @@ def _extract_text_from_word_stream(self, data: bytes) -> str: if len(run_bytes) >= self._min_unicode_bytes: try: - fragment = bytes(run_bytes).decode( - "utf-16-le", errors="ignore" - ).strip() - if ( - len(fragment) >= self._min_text_fragment_length - and not fragment.startswith("\\") + fragment = ( + bytes(run_bytes) + .decode("utf-16-le", errors="ignore") + .strip() + ) + if len( + fragment + ) >= self._min_text_fragment_length and not fragment.startswith( + "\\" ): # Clean control characters - fragment = fragment.replace("\r\n", "\n").replace("\r", "\n") + fragment = fragment.replace("\r\n", "\n").replace( + "\r", "\n" + ) fragment = re.sub( r"[\x00-\x08\x0b\x0c\x0e-\x1f]", "", fragment ) @@ -274,9 +284,7 @@ def _is_run_start_pair(low: int, high: int) -> bool: of actual text). """ # ASCII printable + common whitespace (high byte = 0x00) - if high == 0x00 and ( - 0x20 <= low <= 0x7E or low in (0x0D, 0x0A, 0x09) - ): + if high == 0x00 and (0x20 <= low <= 0x7E or low in (0x0D, 0x0A, 0x09)): return True # Hangul / CJK — allowed to start a run, but ONLY if low != 0x00 @@ -294,9 +302,7 @@ def _is_run_start_pair(low: int, high: int) -> bool: def _is_text_pair(low: int, high: int) -> bool: """Check if a UTF-16LE byte pair represents printable text.""" # ASCII printable + common whitespace, high byte = 0x00 - if high == 0x00 and ( - 0x20 <= low <= 0x7E or low in (0x0D, 0x0A, 0x09) - ): + if high == 0x00 and (0x20 <= low <= 0x7E or low in (0x0D, 0x0A, 0x09)): return True # Hangul Syllables (U+AC00 – U+D7AF) @@ -351,6 +357,7 @@ def _save_images(self, preprocessed: PreprocessedData) -> List[str]: # Simple deduplication by content hash import hashlib + content_hash = hashlib.md5(data).hexdigest() if content_hash in processed_hashes: continue @@ -380,7 +387,7 @@ def _detect_image_format(data: bytes) -> Optional[str]: return None for fmt_name, (signature, min_len) in IMAGE_SIGNATURES.items(): - if len(data) >= min_len and data[:len(signature)] == signature: + if len(data) >= min_len and data[: len(signature)] == signature: # Normalise names that have variants if fmt_name.startswith("gif"): return "gif" diff --git a/contextifier/handlers/doc/converter.py b/contextifier/handlers/doc/converter.py index cf02fee..5081334 100644 --- a/contextifier/handlers/doc/converter.py +++ b/contextifier/handlers/doc/converter.py @@ -31,8 +31,9 @@ class DocConvertedData(NamedTuple): """Result of the DOC conversion stage.""" - ole: olefile.OleFileIO # Opened OLE2 compound file - file_extension: str # Original extension (always "doc") + + ole: olefile.OleFileIO # Opened OLE2 compound file + file_extension: str # Original extension (always "doc") class DocConverter(BaseConverter): diff --git a/contextifier/handlers/doc/handler.py b/contextifier/handlers/doc/handler.py index ac86164..c56bb06 100644 --- a/contextifier/handlers/doc/handler.py +++ b/contextifier/handlers/doc/handler.py @@ -50,8 +50,8 @@ from contextifier.handlers.doc.content_extractor import DocContentExtractor # Magic bytes for format detection -_ZIP_MAGIC = b"PK" # ZIP/OOXML (DOCX) -_RTF_MAGIC = b"{\\rtf" # RTF +_ZIP_MAGIC = b"PK" # ZIP/OOXML (DOCX) +_RTF_MAGIC = b"{\\rtf" # RTF _OLE2_MAGIC = b"\xd0\xcf\x11\xe0\xa1\xb1\x1a\xe1" # OLE2/CFBF (genuine DOC) _HTML_MARKERS = (b" DocStreamData: for entry in ole.listdir(): entry_path = "/".join(entry) if any( - kw in part.lower() - for part in entry - for kw in IMAGE_STREAM_KEYWORDS + kw in part.lower() for part in entry for kw in IMAGE_STREAM_KEYWORDS ): image_streams.append(entry_path) except Exception as exc: diff --git a/contextifier/handlers/docx/_constants.py b/contextifier/handlers/docx/_constants.py index 151f3c9..748b378 100644 --- a/contextifier/handlers/docx/_constants.py +++ b/contextifier/handlers/docx/_constants.py @@ -39,6 +39,7 @@ # ── Element type enum ──────────────────────────────────────────────────── + @unique class ElementType(str, Enum): """Classification of document body elements.""" diff --git a/contextifier/handlers/docx/_paragraph.py b/contextifier/handlers/docx/_paragraph.py index 210101c..8c7f120 100644 --- a/contextifier/handlers/docx/_paragraph.py +++ b/contextifier/handlers/docx/_paragraph.py @@ -20,9 +20,8 @@ from enum import Enum, unique from typing import Any, List, Optional, Tuple -from lxml import etree -from contextifier.handlers.docx._constants import NAMESPACES, ElementType +from contextifier.handlers.docx._constants import NAMESPACES logger = logging.getLogger(__name__) @@ -52,9 +51,11 @@ # ── Drawing descriptor ──────────────────────────────────────────────────── + @unique class DrawingKind(str, Enum): """Kind of drawing element detected.""" + IMAGE = "image" CHART = "chart" DIAGRAM = "diagram" @@ -64,26 +65,30 @@ class DrawingKind(str, Enum): @dataclass class DrawingInfo: """Information about a drawing element found in a run.""" + kind: DrawingKind - rel_id: Optional[str] = None # r:embed or r:link for images - uri: Optional[str] = None # graphicData URI - graphic_data: Any = None # lxml element (for diagrams) + rel_id: Optional[str] = None # r:embed or r:link for images + uri: Optional[str] = None # graphicData URI + graphic_data: Any = None # lxml element (for diagrams) @dataclass class PictInfo: """Information about a legacy VML pict element in a run.""" + rel_id: Optional[str] = None # ── Run element descriptors ─────────────────────────────────────────────── + @dataclass class RunContent: """Content extracted from a single run (or hyperlink child).""" + text: str = "" drawings: List[DrawingInfo] = None # type: ignore[assignment] - picts: List[PictInfo] = None # type: ignore[assignment] + picts: List[PictInfo] = None # type: ignore[assignment] has_page_break: bool = False def __post_init__(self) -> None: @@ -95,7 +100,10 @@ def __post_init__(self) -> None: # ── Paragraph processing ───────────────────────────────────────────────── -def process_paragraph(paragraph_element: Any) -> Tuple[str, List[DrawingInfo], List[PictInfo], bool]: + +def process_paragraph( + paragraph_element: Any, +) -> Tuple[str, List[DrawingInfo], List[PictInfo], bool]: """ Process a ```` element and extract its content. @@ -115,7 +123,7 @@ def process_paragraph(paragraph_element: Any) -> Tuple[str, List[DrawingInfo], L has_page_break = False for child in paragraph_element: - tag = _local_name(child) + _local_name(child) if child.tag == _QN_R: rc = _process_run(child) @@ -163,6 +171,7 @@ def has_page_break(paragraph_element: Any) -> bool: # ── Run processing (internal) ───────────────────────────────────────────── + def _process_run(run_element: Any) -> RunContent: """ Process a single ```` element. @@ -319,6 +328,7 @@ def extract_diagram_text(graphic_data: Any) -> str: # ── Utilities ───────────────────────────────────────────────────────────── + def _local_name(element: Any) -> str: """Get the local name of an lxml element (without namespace).""" tag = element.tag diff --git a/contextifier/handlers/docx/_table_extractor.py b/contextifier/handlers/docx/_table_extractor.py index 8fdae73..bd8343d 100644 --- a/contextifier/handlers/docx/_table_extractor.py +++ b/contextifier/handlers/docx/_table_extractor.py @@ -55,7 +55,9 @@ def extract_table(table_element: Any) -> Optional[TableData]: try: # 1. Column widths col_widths = _calculate_column_widths(table_element) - num_cols = len(col_widths) if col_widths else _estimate_column_count(table_element) + num_cols = ( + len(col_widths) if col_widths else _estimate_column_count(table_element) + ) if num_cols == 0: return None @@ -134,6 +136,7 @@ def extract_table(table_element: Any) -> Optional[TableData]: # ── Internal data structures ────────────────────────────────────────────── + class _RawCell: """Temporary cell representation during parsing.""" @@ -154,6 +157,7 @@ def __init__( # ── Row parsing ─────────────────────────────────────────────────────────── + def _parse_row(tr_element: Any, expected_cols: int) -> List[_RawCell]: """Parse a ```` into a list of ``_RawCell``.""" cells: List[_RawCell] = [] @@ -184,18 +188,21 @@ def _parse_row(tr_element: Any, expected_cols: int) -> List[_RawCell]: # Empty val or "continue" means this cell continues a merge v_merge_continue = True - cells.append(_RawCell( - text=text, - col_span=col_span, - v_merge_restart=v_merge_restart, - v_merge_continue=v_merge_continue, - )) + cells.append( + _RawCell( + text=text, + col_span=col_span, + v_merge_restart=v_merge_restart, + v_merge_continue=v_merge_continue, + ) + ) return cells # ── Cell text extraction ────────────────────────────────────────────────── + def _extract_cell_text(tc_element: Any) -> str: """ Extract text from a ```` (table cell) element. @@ -219,6 +226,7 @@ def _extract_cell_text(tc_element: Any) -> str: # ── Column width calculation ────────────────────────────────────────────── + def _calculate_column_widths(table_element: Any) -> List[int]: """ Calculate column widths from ``/``. @@ -265,6 +273,7 @@ def _estimate_column_count(table_element: Any) -> int: # ── Row span (vMerge) calculation ───────────────────────────────────────── + def _calculate_rowspans( raw_rows: List[List[_RawCell]], num_rows: int, @@ -328,9 +337,7 @@ def _find_cell_at_col( return None -def _find_col_start( - entries: List[Tuple[int, _RawCell]], target_col: int -) -> int: +def _find_col_start(entries: List[Tuple[int, _RawCell]], target_col: int) -> int: """Find the starting column index of the cell covering target_col.""" for col_start, cell in entries: if col_start <= target_col < col_start + cell.col_span: @@ -340,6 +347,7 @@ def _find_col_start( # ── Header detection ────────────────────────────────────────────────────── + def _detect_header(raw_rows: List[List[_RawCell]]) -> bool: """ Detect if the table has a header row. diff --git a/contextifier/handlers/docx/content_extractor.py b/contextifier/handlers/docx/content_extractor.py index 51bbf79..09cec12 100644 --- a/contextifier/handlers/docx/content_extractor.py +++ b/contextifier/handlers/docx/content_extractor.py @@ -21,17 +21,16 @@ import hashlib import logging import re -from typing import Any, Dict, List, Optional, Set +from typing import Any, Dict, List, Optional from contextifier.pipeline.content_extractor import BaseContentExtractor from contextifier.services.table_service import TableService from contextifier.types import ( - ChartData, PreprocessedData, TableData, ) -from contextifier.handlers.docx._constants import NAMESPACES, ElementType +from contextifier.handlers.docx._constants import NAMESPACES from contextifier.handlers.docx._paragraph import ( process_paragraph, extract_diagram_text, @@ -123,8 +122,12 @@ def extract_text( # Process drawings (images, charts, diagrams) for drawing in drawings: content = self._process_drawing( - drawing, doc, charts_by_rel, positional_charts, - chart_index, processed_images, + drawing, + doc, + charts_by_rel, + positional_charts, + chart_index, + processed_images, ) if drawing.kind == DrawingKind.CHART: chart_index += 1 @@ -212,15 +215,11 @@ def extract_images( _, drawings, picts, _ = process_paragraph(element) for drawing in drawings: if drawing.kind == DrawingKind.IMAGE: - tag = self._extract_image_by_rel( - drawing.rel_id, doc, processed - ) + tag = self._extract_image_by_rel(drawing.rel_id, doc, processed) if tag: tags.append(tag) for pict in picts: - tag = self._extract_image_by_rel( - pict.rel_id, doc, processed - ) + tag = self._extract_image_by_rel(pict.rel_id, doc, processed) if tag: tags.append(tag) @@ -243,9 +242,7 @@ def _process_drawing( """Process a drawing element → return content string.""" if drawing.kind == DrawingKind.IMAGE: - tag = self._extract_image_by_rel( - drawing.rel_id, doc, processed_images - ) + tag = self._extract_image_by_rel(drawing.rel_id, doc, processed_images) return tag or "" if drawing.kind == DrawingKind.CHART: @@ -422,7 +419,10 @@ def _extract_supplementary(doc: Any) -> str: if fn_part is not None: fn_element = fn_part.element for fn in fn_element: - fn_type = fn.get("{http://schemas.openxmlformats.org/wordprocessingml/2006/main}type", "") + fn_type = fn.get( + "{http://schemas.openxmlformats.org/wordprocessingml/2006/main}type", + "", + ) if fn_type in ("separator", "continuationSeparator"): continue paras = fn.findall( @@ -444,9 +444,8 @@ def _extract_supplementary(doc: Any) -> str: if footnotes: sections.append( - "[Footnotes]\n" + "\n".join( - f"[{i}] {fn}" for i, fn in enumerate(footnotes, 1) - ) + "[Footnotes]\n" + + "\n".join(f"[{i}] {fn}" for i, fn in enumerate(footnotes, 1)) ) return "\n\n".join(sections) diff --git a/contextifier/handlers/docx/preprocessor.py b/contextifier/handlers/docx/preprocessor.py index c5f020e..46e990b 100644 --- a/contextifier/handlers/docx/preprocessor.py +++ b/contextifier/handlers/docx/preprocessor.py @@ -112,7 +112,9 @@ def _extract_charts_by_rel(self, doc: Any) -> Dict[str, str]: chart_xml_map: Dict[str, bytes] = {} with zipfile.ZipFile(zip_stream, "r") as zf: for name in zf.namelist(): - if name.startswith("word/charts/chart") and name.endswith(".xml"): + if name.startswith("word/charts/chart") and name.endswith( + ".xml" + ): chart_xml_map[name] = zf.read(name) for rel_id, rel in rels.items(): @@ -162,8 +164,7 @@ def _extract_charts_from_zip(self, doc: Any) -> List[str]: chart_files = sorted( name for name in zf.namelist() - if name.startswith("word/charts/chart") - and name.endswith(".xml") + if name.startswith("word/charts/chart") and name.endswith(".xml") ) for chart_file in chart_files: @@ -224,7 +225,11 @@ def _parse_chart_xml(chart_xml: bytes) -> str: # Extract title title_elem = chart_elem.find(".//c:title//c:tx//c:rich//a:t", ns) - title = title_elem.text.strip() if title_elem is not None and title_elem.text else None + title = ( + title_elem.text.strip() + if title_elem is not None and title_elem.text + else None + ) # Detect chart type chart_type = "Chart" diff --git a/contextifier/handlers/html/content_extractor.py b/contextifier/handlers/html/content_extractor.py index 3a1df2e..5c6c45b 100644 --- a/contextifier/handlers/html/content_extractor.py +++ b/contextifier/handlers/html/content_extractor.py @@ -18,16 +18,15 @@ import logging import re -from typing import Any, Dict, List, Optional, Set +from typing import Any, List, Optional -from bs4 import BeautifulSoup, NavigableString, Tag +from bs4 import NavigableString, Tag from contextifier.pipeline.content_extractor import BaseContentExtractor from contextifier.types import ( PreprocessedData, TableData, TableCell, - ChartData, ) from contextifier.handlers.html.preprocessor import HtmlParsedData @@ -35,9 +34,22 @@ _HEADING_TAGS = {"h1", "h2", "h3", "h4", "h5", "h6"} _BLOCK_TAGS = { - "p", "div", "section", "article", "aside", "main", "header", - "footer", "nav", "blockquote", "figure", "figcaption", "details", - "summary", "address", "hgroup", + "p", + "div", + "section", + "article", + "aside", + "main", + "header", + "footer", + "nav", + "blockquote", + "figure", + "figcaption", + "details", + "summary", + "address", + "hgroup", } _LIST_TAGS = {"ul", "ol"} _SKIP_TAGS = {"script", "style", "noscript", "template", "svg", "math"} @@ -242,9 +254,9 @@ def _parse_table(self, table_tag: Tag) -> Optional[TableData]: if not rows_data: return None - num_cols = max( - sum(c.col_span for c in r) for r in rows_data - ) if rows_data else 0 + num_cols = ( + max(sum(c.col_span for c in r) for r in rows_data) if rows_data else 0 + ) has_header = bool(rows_data) and all(c.is_header for c in rows_data[0]) diff --git a/contextifier/handlers/html/converter.py b/contextifier/handlers/html/converter.py index 8ae39c5..8f70ddd 100644 --- a/contextifier/handlers/html/converter.py +++ b/contextifier/handlers/html/converter.py @@ -11,7 +11,7 @@ import logging import re -from typing import Any, NamedTuple, Optional +from typing import Any, NamedTuple from contextifier.pipeline.converter import BaseConverter from contextifier.types import FileContext @@ -45,6 +45,7 @@ class HtmlConvertedData(NamedTuple): """Result of Stage 1 conversion for HTML.""" + html_text: str encoding: str file_extension: str @@ -53,9 +54,7 @@ class HtmlConvertedData(NamedTuple): class HtmlConverter(BaseConverter): """Validate HTML input and decode to string.""" - def convert( - self, file_context: FileContext, **kwargs: Any - ) -> HtmlConvertedData: + def convert(self, file_context: FileContext, **kwargs: Any) -> HtmlConvertedData: file_data: bytes = file_context.get("file_data", b"") if not file_data: raise ConversionError( diff --git a/contextifier/handlers/html/metadata_extractor.py b/contextifier/handlers/html/metadata_extractor.py index 9ce70fc..7debdee 100644 --- a/contextifier/handlers/html/metadata_extractor.py +++ b/contextifier/handlers/html/metadata_extractor.py @@ -86,7 +86,12 @@ def _unpack(source: Any): @staticmethod def _try_parse_date(value: str) -> Optional[datetime]: - for fmt in ("%Y-%m-%d", "%Y-%m-%dT%H:%M:%S", "%Y-%m-%dT%H:%M:%SZ", "%Y-%m-%d %H:%M:%S"): + for fmt in ( + "%Y-%m-%d", + "%Y-%m-%dT%H:%M:%S", + "%Y-%m-%dT%H:%M:%SZ", + "%Y-%m-%d %H:%M:%S", + ): try: return datetime.strptime(value.strip(), fmt) except ValueError: diff --git a/contextifier/handlers/html/preprocessor.py b/contextifier/handlers/html/preprocessor.py index 70afb30..1952a0e 100644 --- a/contextifier/handlers/html/preprocessor.py +++ b/contextifier/handlers/html/preprocessor.py @@ -13,9 +13,9 @@ import base64 import logging import re -from typing import Any, Dict, List, NamedTuple, Optional +from typing import Any, Dict, List, NamedTuple -from bs4 import BeautifulSoup, Tag +from bs4 import BeautifulSoup from contextifier.pipeline.preprocessor import BasePreprocessor from contextifier.types import PreprocessedData @@ -28,13 +28,12 @@ # memory exhaustion from maliciously crafted HTML. _MAX_IMAGE_DECODE_BYTES = 50 * 1024 * 1024 -_DATA_URI_RE = re.compile( - r"data:image/([a-zA-Z0-9+.-]+);base64,([A-Za-z0-9+/=\s]+)" -) +_DATA_URI_RE = re.compile(r"data:image/([a-zA-Z0-9+.-]+);base64,([A-Za-z0-9+/=\s]+)") class HtmlParsedData(NamedTuple): """Preprocessed HTML payload.""" + soup: BeautifulSoup title: str encoding: str @@ -43,9 +42,7 @@ class HtmlParsedData(NamedTuple): class HtmlPreprocessor(BasePreprocessor): """Parse HTML, strip scripts/styles, extract embedded images.""" - def preprocess( - self, converted_data: Any, **kwargs: Any - ) -> PreprocessedData: + def preprocess(self, converted_data: Any, **kwargs: Any) -> PreprocessedData: html_text, encoding, file_ext = self._unpack(converted_data) if not html_text: @@ -69,6 +66,7 @@ def preprocess( # Remove HTML comments from bs4 import Comment + for comment in soup.find_all(string=lambda t: isinstance(t, Comment)): comment.extract() @@ -127,7 +125,11 @@ def validate(self, data: Any) -> bool: @staticmethod def _unpack(converted_data: Any): if isinstance(converted_data, HtmlConvertedData): - return converted_data.html_text, converted_data.encoding, converted_data.file_extension + return ( + converted_data.html_text, + converted_data.encoding, + converted_data.file_extension, + ) if isinstance(converted_data, str): return converted_data, "utf-8", "html" return "", "utf-8", "html" diff --git a/contextifier/handlers/hwp/_constants.py b/contextifier/handlers/hwp/_constants.py index 654f20c..b949ea6 100644 --- a/contextifier/handlers/hwp/_constants.py +++ b/contextifier/handlers/hwp/_constants.py @@ -20,35 +20,35 @@ HWPTAG_BEGIN = 0x10 -HWPTAG_DOCUMENT_PROPERTIES = HWPTAG_BEGIN + 0 # 16 -HWPTAG_ID_MAPPINGS = HWPTAG_BEGIN + 1 # 17 -HWPTAG_BIN_DATA = HWPTAG_BEGIN + 2 # 18 -HWPTAG_FACE_NAME = HWPTAG_BEGIN + 3 # 19 -HWPTAG_BORDER_FILL = HWPTAG_BEGIN + 4 # 20 -HWPTAG_CHAR_SHAPE = HWPTAG_BEGIN + 5 # 21 -HWPTAG_TAB_DEF = HWPTAG_BEGIN + 6 # 22 -HWPTAG_NUMBERING = HWPTAG_BEGIN + 7 # 23 -HWPTAG_BULLET = HWPTAG_BEGIN + 8 # 24 -HWPTAG_PARA_SHAPE = HWPTAG_BEGIN + 9 # 25 -HWPTAG_STYLE = HWPTAG_BEGIN + 10 # 26 -HWPTAG_DOC_DATA = HWPTAG_BEGIN + 11 # 27 +HWPTAG_DOCUMENT_PROPERTIES = HWPTAG_BEGIN + 0 # 16 +HWPTAG_ID_MAPPINGS = HWPTAG_BEGIN + 1 # 17 +HWPTAG_BIN_DATA = HWPTAG_BEGIN + 2 # 18 +HWPTAG_FACE_NAME = HWPTAG_BEGIN + 3 # 19 +HWPTAG_BORDER_FILL = HWPTAG_BEGIN + 4 # 20 +HWPTAG_CHAR_SHAPE = HWPTAG_BEGIN + 5 # 21 +HWPTAG_TAB_DEF = HWPTAG_BEGIN + 6 # 22 +HWPTAG_NUMBERING = HWPTAG_BEGIN + 7 # 23 +HWPTAG_BULLET = HWPTAG_BEGIN + 8 # 24 +HWPTAG_PARA_SHAPE = HWPTAG_BEGIN + 9 # 25 +HWPTAG_STYLE = HWPTAG_BEGIN + 10 # 26 +HWPTAG_DOC_DATA = HWPTAG_BEGIN + 11 # 27 HWPTAG_DISTRIBUTE_DOC_DATA = HWPTAG_BEGIN + 12 # 28 HWPTAG_COMPATIBLE_DOCUMENT = HWPTAG_BEGIN + 38 # 54 # Body text record tags -HWPTAG_PARA_HEADER = HWPTAG_BEGIN + 50 # 66 -HWPTAG_PARA_TEXT = HWPTAG_BEGIN + 51 # 67 -HWPTAG_PARA_CHAR_SHAPE = HWPTAG_BEGIN + 52 # 68 -HWPTAG_PARA_LINE_SEG = HWPTAG_BEGIN + 53 # 69 -HWPTAG_PARA_RANGE_TAG = HWPTAG_BEGIN + 54 # 70 -HWPTAG_CTRL_HEADER = HWPTAG_BEGIN + 55 # 71 -HWPTAG_LIST_HEADER = HWPTAG_BEGIN + 56 # 72 -HWPTAG_PAGE_DEF = HWPTAG_BEGIN + 57 # 73 -HWPTAG_FOOTNOTE_SHAPE = HWPTAG_BEGIN + 58 # 74 -HWPTAG_PAGE_BORDER_FILL = HWPTAG_BEGIN + 59 # 75 -HWPTAG_SHAPE_COMPONENT = HWPTAG_BEGIN + 60 # 76 -HWPTAG_TABLE = HWPTAG_BEGIN + 61 # 77 +HWPTAG_PARA_HEADER = HWPTAG_BEGIN + 50 # 66 +HWPTAG_PARA_TEXT = HWPTAG_BEGIN + 51 # 67 +HWPTAG_PARA_CHAR_SHAPE = HWPTAG_BEGIN + 52 # 68 +HWPTAG_PARA_LINE_SEG = HWPTAG_BEGIN + 53 # 69 +HWPTAG_PARA_RANGE_TAG = HWPTAG_BEGIN + 54 # 70 +HWPTAG_CTRL_HEADER = HWPTAG_BEGIN + 55 # 71 +HWPTAG_LIST_HEADER = HWPTAG_BEGIN + 56 # 72 +HWPTAG_PAGE_DEF = HWPTAG_BEGIN + 57 # 73 +HWPTAG_FOOTNOTE_SHAPE = HWPTAG_BEGIN + 58 # 74 +HWPTAG_PAGE_BORDER_FILL = HWPTAG_BEGIN + 59 # 75 +HWPTAG_SHAPE_COMPONENT = HWPTAG_BEGIN + 60 # 76 +HWPTAG_TABLE = HWPTAG_BEGIN + 61 # 77 HWPTAG_SHAPE_COMPONENT_LINE = HWPTAG_BEGIN + 62 # 78 HWPTAG_SHAPE_COMPONENT_RECTANGLE = HWPTAG_BEGIN + 63 # 79 HWPTAG_SHAPE_COMPONENT_ELLIPSE = HWPTAG_BEGIN + 64 # 80 @@ -58,11 +58,11 @@ HWPTAG_SHAPE_COMPONENT_OLE = HWPTAG_BEGIN + 68 # 84 HWPTAG_SHAPE_COMPONENT_PICTURE = HWPTAG_BEGIN + 69 # 85 HWPTAG_SHAPE_COMPONENT_CONTAINER = HWPTAG_BEGIN + 70 # 86 -HWPTAG_CTRL_DATA = HWPTAG_BEGIN + 71 # 87 -HWPTAG_EQEDIT = HWPTAG_BEGIN + 72 # 88 +HWPTAG_CTRL_DATA = HWPTAG_BEGIN + 71 # 87 +HWPTAG_EQEDIT = HWPTAG_BEGIN + 72 # 88 -HWPTAG_MEMO_SHAPE = HWPTAG_BEGIN + 76 # 92 -HWPTAG_MEMO_LIST = HWPTAG_BEGIN + 77 # 93 +HWPTAG_MEMO_SHAPE = HWPTAG_BEGIN + 76 # 92 +HWPTAG_MEMO_LIST = HWPTAG_BEGIN + 77 # 93 # ── Special control characters found in PARA_TEXT ───────────────────────── diff --git a/contextifier/handlers/hwp/_decoder.py b/contextifier/handlers/hwp/_decoder.py index 2975bdb..55a3ac5 100644 --- a/contextifier/handlers/hwp/_decoder.py +++ b/contextifier/handlers/hwp/_decoder.py @@ -46,9 +46,7 @@ def is_compressed(ole: olefile.OleFileIO) -> bool: header = stream.read() if len(header) < FILE_HEADER_FLAGS_OFFSET + 4: return False - flags = struct.unpack_from( - "= 6: ext_len = struct.unpack_from("= 6 + ext_len * 2: - ext = payload[6: 6 + ext_len * 2].decode( + ext = payload[6 : 6 + ext_len * 2].decode( "utf-16le", errors="ignore" ) by_id[sid] = (sid, ext) @@ -101,7 +101,7 @@ def parse_doc_info( if len(payload) >= 6: ext_len = struct.unpack_from("= 6 + ext_len * 2: - ext = payload[6: 6 + ext_len * 2].decode( + ext = payload[6 : 6 + ext_len * 2].decode( "utf-16le", errors="ignore" ) if sid > 0: diff --git a/contextifier/handlers/hwp/_record.py b/contextifier/handlers/hwp/_record.py index 8da48ff..b32f946 100644 --- a/contextifier/handlers/hwp/_record.py +++ b/contextifier/handlers/hwp/_record.py @@ -178,7 +178,7 @@ def get_next_siblings(self, count: int) -> List["HwpRecord"]: except ValueError: return [] - return siblings[idx + 1: idx + 1 + count] + return siblings[idx + 1 : idx + 1 + count] __all__ = ["HwpRecord"] diff --git a/contextifier/handlers/hwp/_recovery.py b/contextifier/handlers/hwp/_recovery.py index a7d231d..46ca6a8 100644 --- a/contextifier/handlers/hwp/_recovery.py +++ b/contextifier/handlers/hwp/_recovery.py @@ -35,12 +35,12 @@ def extract_text_raw(data: bytes) -> str: val = struct.unpack_from(" str: return "".join(p for p in parts if p.strip()) -def find_zlib_streams( - data: bytes, min_size: int = 50 -) -> List[Tuple[int, bytes]]: +def find_zlib_streams(data: bytes, min_size: int = 50) -> List[Tuple[int, bytes]]: """ Scan raw data for zlib-compressed streams and decompress them. diff --git a/contextifier/handlers/hwp/content_extractor.py b/contextifier/handlers/hwp/content_extractor.py index bdd4aac..894d73a 100644 --- a/contextifier/handlers/hwp/content_extractor.py +++ b/contextifier/handlers/hwp/content_extractor.py @@ -23,7 +23,6 @@ from contextifier.pipeline.content_extractor import BaseContentExtractor from contextifier.types import ( ChartData, - ExtractionResult, PreprocessedData, TableData, ) @@ -41,11 +40,7 @@ from contextifier.handlers.hwp._recovery import extract_text_raw if TYPE_CHECKING: - import olefile - from contextifier.services.image_service import ImageService - from contextifier.services.tag_service import TagService - from contextifier.services.chart_service import ChartService - from contextifier.services.table_service import TableService + pass logger = logging.getLogger(__name__) @@ -57,9 +52,7 @@ class HwpContentExtractor(BaseContentExtractor): # ── BaseContentExtractor interface ──────────────────────────────── - def extract_text( - self, preprocessed: PreprocessedData, **kwargs: Any - ) -> str: + def extract_text(self, preprocessed: PreprocessedData, **kwargs: Any) -> str: ole, bin_data_map = self._unwrap(preprocessed) if ole is None: return "" @@ -177,9 +170,13 @@ def _process_paragraph( for i, seg in enumerate(segments): parts.append(seg) if i < len(controls): - parts.append(self._traverse(controls[i], ole, bin_data_map, processed_images)) + parts.append( + self._traverse(controls[i], ole, bin_data_map, processed_images) + ) for k in range(len(segments) - 1, len(controls)): - parts.append(self._traverse(controls[k], ole, bin_data_map, processed_images)) + parts.append( + self._traverse(controls[k], ole, bin_data_map, processed_images) + ) else: parts.append(text_content) for c in controls: @@ -258,15 +255,11 @@ def _process_picture( idx = _extract_bindata_index(record.payload, len(bin_list)) if idx is not None and 0 < idx <= len(bin_list): - return self._save_bindata_image( - ole, bin_list[idx - 1], processed_images - ) + return self._save_bindata_image(ole, bin_list[idx - 1], processed_images) # Fallback: if only one BinData, use it if len(bin_list) == 1: - return self._save_bindata_image( - ole, bin_list[0], processed_images - ) + return self._save_bindata_image(ole, bin_list[0], processed_images) return None @@ -291,6 +284,7 @@ def _save_bindata_image( try: import zlib + raw = ole.openstream(stream_path).read() # Try decompress @@ -320,14 +314,19 @@ def _save_bindata_image( def _unwrap(preprocessed: PreprocessedData): """Return (ole, bin_data_map) from PreprocessedData.""" ole = preprocessed.content - bdm = preprocessed.resources.get("bin_data_map") if preprocessed.resources else None + bdm = ( + preprocessed.resources.get("bin_data_map") + if preprocessed.resources + else None + ) return ole, bdm @staticmethod def _sorted_sections(ole: Any) -> List: """Return BodyText/Section* entries sorted by number.""" sections = [ - e for e in ole.listdir() + e + for e in ole.listdir() if len(e) >= 2 and e[0] == STREAM_BODY_TEXT and e[1].startswith("Section") ] sections.sort(key=lambda x: int(x[1].replace("Section", ""))) @@ -360,7 +359,27 @@ def _extract_bindata_index(payload: bytes, list_len: int) -> Optional[int]: return val # Strategy 3: scan common offsets - for off in (4, 6, 10, 12, 14, 16, 18, 20, 40, 44, 48, 52, 56, 60, 64, 68, 72, 76, 80): + for off in ( + 4, + 6, + 10, + 12, + 14, + 16, + 18, + 20, + 40, + 44, + 48, + 52, + 56, + 60, + 64, + 68, + 72, + 76, + 80, + ): if len(payload) >= off + 2: val = struct.unpack_from(" HwpConvertedData: file_data: bytes = file_context.get("file_data", b"") if not file_data: raise ConversionError( - "Empty file data", stage="convert", handler="hwp", + "Empty file data", + stage="convert", + handler="hwp", ) try: diff --git a/contextifier/handlers/hwp/metadata_extractor.py b/contextifier/handlers/hwp/metadata_extractor.py index 08c18da..fa7a038 100644 --- a/contextifier/handlers/hwp/metadata_extractor.py +++ b/contextifier/handlers/hwp/metadata_extractor.py @@ -70,7 +70,8 @@ def extract(self, source: Any) -> DocumentMetadata: # Count sections as pages section_count = sum( - 1 for e in ole.listdir() + 1 + for e in ole.listdir() if len(e) >= 2 and e[0] == "BodyText" and e[1].startswith("Section") ) @@ -81,8 +82,12 @@ def extract(self, source: Any) -> DocumentMetadata: keywords=_str_or_none(meta_dict.get("keywords")), comments=_str_or_none(meta_dict.get("comments")), last_saved_by=_str_or_none(meta_dict.get("last_saved_by")), - create_time=meta_dict.get("create_time") if isinstance(meta_dict.get("create_time"), datetime) else None, - last_saved_time=meta_dict.get("last_saved_time") if isinstance(meta_dict.get("last_saved_time"), datetime) else None, + create_time=meta_dict.get("create_time") + if isinstance(meta_dict.get("create_time"), datetime) + else None, + last_saved_time=meta_dict.get("last_saved_time") + if isinstance(meta_dict.get("last_saved_time"), datetime) + else None, page_count=section_count or None, ) @@ -163,7 +168,7 @@ def _parse_hwp_summary(data: bytes) -> Dict[str, Any]: if voff + 4 < len(data): slen = struct.unpack_from(" Dict[str, Any]: slen = struct.unpack_from(" PreprocessedData: # Count sections sections = [ - e for e in ole.listdir() + e + for e in ole.listdir() if len(e) >= 2 and e[0] == STREAM_BODY_TEXT and e[1].startswith("Section") ] diff --git a/contextifier/handlers/hwpx/_constants.py b/contextifier/handlers/hwpx/_constants.py index 2f354b4..1f45773 100644 --- a/contextifier/handlers/hwpx/_constants.py +++ b/contextifier/handlers/hwpx/_constants.py @@ -40,11 +40,11 @@ # Standard Paths Inside the ZIP Archive # ═══════════════════════════════════════════════════════════════════════════════ -HPF_PATH = "Contents/content.hpf" # OPF manifest — bin_item_id → file path -HEADER_PATH = "Contents/header.xml" # Document metadata / docInfo -VERSION_PATH = "version.xml" # Version information -MANIFEST_PATH = "META-INF/manifest.xml" # MIME type information -MIMETYPE_PATH = "mimetype" # Plain text MIME type file +HPF_PATH = "Contents/content.hpf" # OPF manifest — bin_item_id → file path +HEADER_PATH = "Contents/header.xml" # Document metadata / docInfo +VERSION_PATH = "version.xml" # Version information +MANIFEST_PATH = "META-INF/manifest.xml" # MIME type information +MIMETYPE_PATH = "mimetype" # Plain text MIME type file # Alternative header paths (for older/variant HWPX files) HEADER_FILE_PATHS: list[str] = [ @@ -54,7 +54,7 @@ ] # Section XML pattern prefix -SECTION_PREFIX = "Contents/section" # Sections are Contents/section0.xml, etc. +SECTION_PREFIX = "Contents/section" # Sections are Contents/section0.xml, etc. # BinData directory prefix BINDATA_PREFIX = "BinData/" @@ -70,13 +70,25 @@ # Image Extensions # ═══════════════════════════════════════════════════════════════════════════════ -SUPPORTED_IMAGE_EXTENSIONS: frozenset[str] = frozenset({ - ".png", ".jpg", ".jpeg", ".gif", ".bmp", ".tiff", ".tif", ".webp", -}) - -SKIP_IMAGE_EXTENSIONS: frozenset[str] = frozenset({ - ".wmf", ".emf", -}) +SUPPORTED_IMAGE_EXTENSIONS: frozenset[str] = frozenset( + { + ".png", + ".jpg", + ".jpeg", + ".gif", + ".bmp", + ".tiff", + ".tif", + ".webp", + } +) + +SKIP_IMAGE_EXTENSIONS: frozenset[str] = frozenset( + { + ".wmf", + ".emf", + } +) # ═══════════════════════════════════════════════════════════════════════════════ # OOXML Chart Namespaces (for embedded charts) diff --git a/contextifier/handlers/hwpx/_section.py b/contextifier/handlers/hwpx/_section.py index 9dcad17..f1b8407 100644 --- a/contextifier/handlers/hwpx/_section.py +++ b/contextifier/handlers/hwpx/_section.py @@ -33,18 +33,15 @@ from __future__ import annotations import logging -import os import xml.etree.ElementTree as ET import zipfile from typing import TYPE_CHECKING, Dict, List, Optional, Set from contextifier.handlers.hwpx._constants import ( - BINDATA_PREFIX, CHART_PREFIXES, HWPX_NAMESPACES, OOXML_CHART_NS, CHART_TYPE_MAP, - SUPPORTED_IMAGE_EXTENSIONS, ) from contextifier.handlers.hwpx._table import parse_hwpx_table @@ -59,6 +56,7 @@ # Public API # ═══════════════════════════════════════════════════════════════════════════════ + def parse_hwpx_section( section_xml: bytes, zf: zipfile.ZipFile, @@ -103,7 +101,10 @@ def parse_hwpx_section( paragraphs = root.findall(".//hp:p", ns) for para in paragraphs: para_text = _process_paragraph( - para, zf, bin_item_map, ns, + para, + zf, + bin_item_map, + ns, image_service=image_service, chart_service=chart_service, processed_images=processed_images, @@ -118,6 +119,7 @@ def parse_hwpx_section( # Paragraph Processing # ═══════════════════════════════════════════════════════════════════════════════ + def _process_paragraph( para: ET.Element, zf: zipfile.ZipFile, @@ -136,7 +138,10 @@ def _process_paragraph( if tag == "run": run_text = _process_run( - child, zf, bin_item_map, ns, + child, + zf, + bin_item_map, + ns, image_service=image_service, processed_images=processed_images, ) @@ -150,7 +155,10 @@ def _process_paragraph( elif tag == "switch": switch_text = _process_switch( - child, zf, bin_item_map, ns, + child, + zf, + bin_item_map, + ns, chart_service=chart_service, image_service=image_service, processed_images=processed_images, @@ -160,7 +168,10 @@ def _process_paragraph( elif tag == "ctrl": ctrl_text = _process_ctrl( - child, zf, bin_item_map, ns, + child, + zf, + bin_item_map, + ns, image_service=image_service, processed_images=processed_images, ) @@ -169,7 +180,10 @@ def _process_paragraph( elif tag == "pic": img_text = _process_picture_element( - child, zf, bin_item_map, ns, + child, + zf, + bin_item_map, + ns, image_service=image_service, processed_images=processed_images, ) @@ -178,7 +192,8 @@ def _process_paragraph( elif tag == "chart": chart_text = _process_chart_ref( - child, zf, + child, + zf, chart_service=chart_service, ) if chart_text: @@ -191,6 +206,7 @@ def _process_paragraph( # Run / Control / Switch Processing # ═══════════════════════════════════════════════════════════════════════════════ + def _process_run( run: ET.Element, zf: zipfile.ZipFile, @@ -216,7 +232,10 @@ def _process_run( elif tag == "ctrl": ctrl_text = _process_ctrl( - child, zf, bin_item_map, ns, + child, + zf, + bin_item_map, + ns, image_service=image_service, processed_images=processed_images, ) @@ -225,7 +244,10 @@ def _process_run( elif tag == "pic": img_text = _process_picture_element( - child, zf, bin_item_map, ns, + child, + zf, + bin_item_map, + ns, image_service=image_service, processed_images=processed_images, ) @@ -249,7 +271,10 @@ def _process_ctrl( # hc:pic is in the 'hc' namespace for pic in ctrl.findall("hc:pic", ns): img_text = _process_picture_element( - pic, zf, bin_item_map, ns, + pic, + zf, + bin_item_map, + ns, image_service=image_service, processed_images=processed_images, ) @@ -258,7 +283,10 @@ def _process_ctrl( # Also direct hp:pic for pic in ctrl.findall("hp:pic", ns): img_text = _process_picture_element( - pic, zf, bin_item_map, ns, + pic, + zf, + bin_item_map, + ns, image_service=image_service, processed_images=processed_images, ) @@ -290,14 +318,19 @@ def _process_switch( if tag == "chart": chart_text = _process_chart_ref( - child, zf, chart_service=chart_service, + child, + zf, + chart_service=chart_service, ) if chart_text: parts.append(chart_text) elif tag == "p": para_text = _process_paragraph( - child, zf, bin_item_map, ns, + child, + zf, + bin_item_map, + ns, image_service=image_service, chart_service=chart_service, processed_images=processed_images, @@ -307,7 +340,10 @@ def _process_switch( elif tag == "pic": img_text = _process_picture_element( - child, zf, bin_item_map, ns, + child, + zf, + bin_item_map, + ns, image_service=image_service, processed_images=processed_images, ) @@ -320,7 +356,10 @@ def _process_switch( tag = child.tag.split("}")[-1] if "}" in child.tag else child.tag if tag == "p": para_text = _process_paragraph( - child, zf, bin_item_map, ns, + child, + zf, + bin_item_map, + ns, image_service=image_service, chart_service=chart_service, processed_images=processed_images, @@ -335,6 +374,7 @@ def _process_switch( # Image Processing # ═══════════════════════════════════════════════════════════════════════════════ + def _process_picture_element( pic: ET.Element, zf: zipfile.ZipFile, @@ -411,6 +451,7 @@ def _resolve_zip_path(zf: zipfile.ZipFile, href: str) -> Optional[str]: # Chart Processing # ═══════════════════════════════════════════════════════════════════════════════ + def _process_chart_ref( chart_elem: ET.Element, zf: zipfile.ZipFile, diff --git a/contextifier/handlers/hwpx/_table.py b/contextifier/handlers/hwpx/_table.py index 7a217fe..75d2523 100644 --- a/contextifier/handlers/hwpx/_table.py +++ b/contextifier/handlers/hwpx/_table.py @@ -26,6 +26,7 @@ # Public API # ═══════════════════════════════════════════════════════════════════════════════ + def parse_hwpx_table( tbl_elem: ET.Element, ns: Optional[Dict[str, str]] = None, @@ -86,6 +87,7 @@ def parse_hwpx_table( # Internal Grid Builder # ═══════════════════════════════════════════════════════════════════════════════ + def _build_grid( tbl_elem: ET.Element, ns: Dict[str, str], @@ -178,6 +180,7 @@ def _extract_cell_text(tc: ET.Element, ns: Dict[str, str]) -> str: # HTML Renderer # ═══════════════════════════════════════════════════════════════════════════════ + def _render_html( grid: Dict[Tuple[int, int], Dict], total_rows: int, @@ -216,7 +219,12 @@ def _render_html( attrs.append(f'colspan="{info["colspan"]}"') attr_str = (" " + " ".join(attrs)) if attrs else "" - text = info["text"].replace("&", "&").replace("<", "<").replace(">", ">") + text = ( + info["text"] + .replace("&", "&") + .replace("<", "<") + .replace(">", ">") + ) cells_html.append(f"{text}") if cells_html: diff --git a/contextifier/handlers/hwpx/content_extractor.py b/contextifier/handlers/hwpx/content_extractor.py index 2f9343f..b123ad5 100644 --- a/contextifier/handlers/hwpx/content_extractor.py +++ b/contextifier/handlers/hwpx/content_extractor.py @@ -14,14 +14,13 @@ import logging import os -from typing import TYPE_CHECKING, Any, Dict, List, Optional, Set +from typing import TYPE_CHECKING, Any, Dict, List, Set import zipfile from contextifier.pipeline.content_extractor import BaseContentExtractor from contextifier.types import ( ChartData, - ExtractionResult, PreprocessedData, TableData, ) @@ -33,10 +32,7 @@ from contextifier.handlers.hwpx.preprocessor import find_section_paths if TYPE_CHECKING: - from contextifier.services.image_service import ImageService - from contextifier.services.tag_service import TagService - from contextifier.services.chart_service import ChartService - from contextifier.services.table_service import TableService + pass logger = logging.getLogger(__name__) @@ -48,14 +44,14 @@ class HwpxContentExtractor(BaseContentExtractor): # ── BaseContentExtractor interface ──────────────────────────────── - def extract_text( - self, preprocessed: PreprocessedData, **kwargs: Any - ) -> str: + def extract_text(self, preprocessed: PreprocessedData, **kwargs: Any) -> str: zf, bin_item_map = self._unwrap(preprocessed) if zf is None: return "" - section_paths = preprocessed.properties.get("section_paths") or find_section_paths(zf) + section_paths = preprocessed.properties.get( + "section_paths" + ) or find_section_paths(zf) processed_images: Set[str] = set() parts: List[str] = [] @@ -80,7 +76,9 @@ def extract_text( # Process remaining BinData images not referenced inline remaining = self._process_remaining_images( - zf, bin_item_map, processed_images, + zf, + bin_item_map, + processed_images, ) if remaining: parts.append(remaining) diff --git a/contextifier/handlers/hwpx/converter.py b/contextifier/handlers/hwpx/converter.py index 37622a8..e79bf41 100644 --- a/contextifier/handlers/hwpx/converter.py +++ b/contextifier/handlers/hwpx/converter.py @@ -24,8 +24,9 @@ class HwpxConvertedData(NamedTuple): """Result of the HWPX conversion stage.""" - zf: zipfile.ZipFile # Opened ZIP archive - file_data: bytes # Original raw bytes + + zf: zipfile.ZipFile # Opened ZIP archive + file_data: bytes # Original raw bytes class HwpxConverter(BaseConverter): @@ -40,7 +41,9 @@ def convert(self, file_context: FileContext, **kwargs: Any) -> HwpxConvertedData file_data: bytes = file_context.get("file_data", b"") if not file_data: raise ConversionError( - "Empty file data", stage="convert", handler="hwpx", + "Empty file data", + stage="convert", + handler="hwpx", ) try: diff --git a/contextifier/handlers/hwpx/metadata_extractor.py b/contextifier/handlers/hwpx/metadata_extractor.py index ba9ece4..48326ec 100644 --- a/contextifier/handlers/hwpx/metadata_extractor.py +++ b/contextifier/handlers/hwpx/metadata_extractor.py @@ -20,7 +20,6 @@ from contextifier.handlers.hwpx._constants import ( HEADER_FILE_PATHS, HWPX_NAMESPACES, - MANIFEST_NAMESPACES, MANIFEST_PATH, VERSION_PATH, ) @@ -53,8 +52,14 @@ def extract(self, source: Any) -> DocumentMetadata: # Map standard fields standard_keys = { - "title", "subject", "author", "keywords", "comments", - "last_saved_by", "create_time", "last_saved_time", + "title", + "subject", + "author", + "keywords", + "comments", + "last_saved_by", + "create_time", + "last_saved_time", } custom = {k: v for k, v in raw.items() if k not in standard_keys} @@ -124,20 +129,14 @@ def _read_manifest(self, zf: zipfile.ZipFile, raw: Dict[str, Any]) -> None: for child in root: tag = child.tag.split("}")[-1] if "}" in child.tag else child.tag if tag == "file-entry": - full_path = ( - child.get("full-path") - or child.get( - "{urn:oasis:names:tc:opendocument:xmlns:manifest:1.0}full-path", - "", - ) + full_path = child.get("full-path") or child.get( + "{urn:oasis:names:tc:opendocument:xmlns:manifest:1.0}full-path", + "", ) if full_path == "/": - media_type = ( - child.get("media-type") - or child.get( - "{urn:oasis:names:tc:opendocument:xmlns:manifest:1.0}media-type", - "", - ) + media_type = child.get("media-type") or child.get( + "{urn:oasis:names:tc:opendocument:xmlns:manifest:1.0}media-type", + "", ) if media_type: raw["media_type"] = media_type diff --git a/contextifier/handlers/hwpx/preprocessor.py b/contextifier/handlers/hwpx/preprocessor.py index eb0509a..873de14 100644 --- a/contextifier/handlers/hwpx/preprocessor.py +++ b/contextifier/handlers/hwpx/preprocessor.py @@ -21,7 +21,6 @@ from contextifier.handlers.hwpx._constants import ( HPF_PATH, OPF_NAMESPACES, - SECTION_PREFIX, ) from contextifier.handlers.hwpx.converter import HwpxConvertedData @@ -107,8 +106,8 @@ def preprocess(self, converted_data: Any, **kwargs: Any) -> PreprocessedData: sections = find_section_paths(zf) return PreprocessedData( - content=zf, # ZipFile for downstream stages - raw_content=file_data, # Original bytes + content=zf, # ZipFile for downstream stages + raw_content=file_data, # Original bytes encoding="utf-8", resources={ "file_data": file_data, diff --git a/contextifier/handlers/image/_constants.py b/contextifier/handlers/image/_constants.py index c9470d9..91cfc16 100644 --- a/contextifier/handlers/image/_constants.py +++ b/contextifier/handlers/image/_constants.py @@ -15,7 +15,7 @@ MAGIC_GIF89 = b"GIF89a" MAGIC_BMP = b"BM" MAGIC_WEBP_RIFF = b"RIFF" -MAGIC_WEBP_TAG = b"WEBP" # bytes 8..12 +MAGIC_WEBP_TAG = b"WEBP" # bytes 8..12 MAGIC_TIFF_LE = b"II\x2a\x00" MAGIC_TIFF_BE = b"MM\x00\x2a" @@ -40,16 +40,36 @@ # Supported extensions (from handler.py skeleton) # ═══════════════════════════════════════════════════════════════════════════════ -IMAGE_EXTENSIONS: frozenset[str] = frozenset({ - "jpg", "jpeg", "png", "gif", "bmp", "webp", - "tiff", "tif", "ico", "heic", "heif", - # Note: SVG is handled by TextHandler as it is XML-based text. -}) +IMAGE_EXTENSIONS: frozenset[str] = frozenset( + { + "jpg", + "jpeg", + "png", + "gif", + "bmp", + "webp", + "tiff", + "tif", + "ico", + "heic", + "heif", + # Note: SVG is handled by TextHandler as it is XML-based text. + } +) # Subset for which we can validate magic bytes -MAGIC_VALIDATED_EXTENSIONS: frozenset[str] = frozenset({ - "jpg", "jpeg", "png", "gif", "bmp", "webp", "tiff", "tif", -}) +MAGIC_VALIDATED_EXTENSIONS: frozenset[str] = frozenset( + { + "jpg", + "jpeg", + "png", + "gif", + "bmp", + "webp", + "tiff", + "tif", + } +) def detect_image_format(data: bytes) -> str | None: diff --git a/contextifier/handlers/image/content_extractor.py b/contextifier/handlers/image/content_extractor.py index 62d0655..36ae2fc 100644 --- a/contextifier/handlers/image/content_extractor.py +++ b/contextifier/handlers/image/content_extractor.py @@ -25,7 +25,6 @@ from __future__ import annotations import logging -import os from typing import Any, Optional from contextifier.pipeline.content_extractor import BaseContentExtractor @@ -49,7 +48,9 @@ def extract_text( preprocessed: PreprocessedData, **kwargs: Any, ) -> str: - image_data: bytes = preprocessed.content if isinstance(preprocessed.content, bytes) else b"" + image_data: bytes = ( + preprocessed.content if isinstance(preprocessed.content, bytes) else b"" + ) if not image_data: return "" diff --git a/contextifier/handlers/image/converter.py b/contextifier/handlers/image/converter.py index d3c3025..ddfbcc1 100644 --- a/contextifier/handlers/image/converter.py +++ b/contextifier/handlers/image/converter.py @@ -26,8 +26,9 @@ class ImageConvertedData(NamedTuple): """Result of the image conversion stage.""" - image_data: bytes # Raw image bytes (unchanged) - detected_format: Optional[str] # Detected format from magic bytes (or None) + + image_data: bytes # Raw image bytes (unchanged) + detected_format: Optional[str] # Detected format from magic bytes (or None) class ImageConverter(BaseConverter): @@ -43,7 +44,9 @@ def convert(self, file_context: FileContext, **kwargs: Any) -> ImageConvertedDat file_data: bytes = file_context.get("file_data", b"") if not file_data: raise ConversionError( - "Empty file data", stage="convert", handler="image", + "Empty file data", + stage="convert", + handler="image", ) detected = detect_image_format(file_data) @@ -54,7 +57,8 @@ def convert(self, file_context: FileContext, **kwargs: Any) -> ImageConvertedDat # but we let the content stage decide). if ext in MAGIC_VALIDATED_EXTENSIONS and detected is None: logger.warning( - "Image magic bytes not recognised for extension '%s'", ext, + "Image magic bytes not recognised for extension '%s'", + ext, ) return ImageConvertedData(image_data=file_data, detected_format=detected) diff --git a/contextifier/handlers/pdf/_constants.py b/contextifier/handlers/pdf/_constants.py index c16cc2e..7d0a56e 100644 --- a/contextifier/handlers/pdf/_constants.py +++ b/contextifier/handlers/pdf/_constants.py @@ -20,9 +20,7 @@ # ── PDF Date Parsing ───────────────────────────────────────────────────────── -_PDF_DATE_RE = re.compile( - r"D:(\d{4})(\d{2})?(\d{2})?(\d{2})?(\d{2})?(\d{2})?" -) +_PDF_DATE_RE = re.compile(r"D:(\d{4})(\d{2})?(\d{2})?(\d{2})?(\d{2})?(\d{2})?") def parse_pdf_date(date_str: Optional[str]) -> Optional[datetime]: @@ -37,9 +35,13 @@ def parse_pdf_date(date_str: Optional[str]) -> Optional[datetime]: if not m: return None try: - parts = [int(g) if g else d for g, d in zip( - m.groups(), (0, 1, 1, 0, 0, 0), - )] + parts = [ + int(g) if g else d + for g, d in zip( + m.groups(), + (0, 1, 1, 0, 0, 0), + ) + ] return datetime(*parts) # type: ignore[arg-type] except (ValueError, TypeError): return None diff --git a/contextifier/handlers/pdf/converter.py b/contextifier/handlers/pdf/converter.py index 30fb77f..c6a4fe0 100644 --- a/contextifier/handlers/pdf/converter.py +++ b/contextifier/handlers/pdf/converter.py @@ -28,7 +28,8 @@ class PdfConvertedData(NamedTuple): """Result of PDF conversion: a fitz.Document + the original bytes.""" - doc: Any # fitz.Document (typed as Any to avoid import issues in tests) + + doc: Any # fitz.Document (typed as Any to avoid import issues in tests) file_data: bytes diff --git a/contextifier/handlers/pdf/handler.py b/contextifier/handlers/pdf/handler.py index dc6f51c..873b515 100644 --- a/contextifier/handlers/pdf/handler.py +++ b/contextifier/handlers/pdf/handler.py @@ -73,12 +73,15 @@ def create_metadata_extractor(self) -> BaseMetadataExtractor: def create_content_extractor(self) -> BaseContentExtractor: mode = self._config.get_format_option( - PDF_FORMAT_OPTION_KEY, PDF_MODE_OPTION, PDF_MODE_PLUS, + PDF_FORMAT_OPTION_KEY, + PDF_MODE_OPTION, + PDF_MODE_PLUS, ) logger.debug("[PDFHandler] PDF mode = %s", mode) if mode not in PDF_VALID_MODES: from contextifier.errors import ConfigurationError + raise ConfigurationError( f"Invalid PDF mode '{mode}'. Valid options: {sorted(PDF_VALID_MODES)}", context={"mode": mode}, @@ -88,6 +91,7 @@ def create_content_extractor(self) -> BaseContentExtractor: from contextifier.handlers.pdf_default import ( PdfDefaultContentExtractor, ) + return PdfDefaultContentExtractor( image_service=self._image_service, tag_service=self._tag_service, @@ -99,6 +103,7 @@ def create_content_extractor(self) -> BaseContentExtractor: from contextifier.handlers.pdf_plus import ( PdfPlusContentExtractor, ) + return PdfPlusContentExtractor( image_service=self._image_service, tag_service=self._tag_service, @@ -113,4 +118,5 @@ def create_postprocessor(self) -> BasePostprocessor: tag_service=self._tag_service, ) + __all__ = ["PDFHandler"] diff --git a/contextifier/handlers/pdf/preprocessor.py b/contextifier/handlers/pdf/preprocessor.py index cb979d8..e307cb7 100644 --- a/contextifier/handlers/pdf/preprocessor.py +++ b/contextifier/handlers/pdf/preprocessor.py @@ -45,8 +45,8 @@ def preprocess(self, converted_data: Any, **kwargs: Any) -> PreprocessedData: needs_ocr = self._detect_scan(doc, page_count) return PreprocessedData( - content=doc, # fitz.Document - raw_content=file_data, # original bytes + content=doc, # fitz.Document + raw_content=file_data, # original bytes encoding="binary", resources={"document": doc}, properties={ @@ -89,7 +89,8 @@ def _detect_scan(doc: Any, page_count: int) -> bool: if is_scan: logger.info( "PDF appears to be a scan (avg %.1f chars/page across %d sample pages)", - avg, sample_size, + avg, + sample_size, ) return is_scan diff --git a/contextifier/handlers/pdf_default/content_extractor.py b/contextifier/handlers/pdf_default/content_extractor.py index 2901633..6c93941 100644 --- a/contextifier/handlers/pdf_default/content_extractor.py +++ b/contextifier/handlers/pdf_default/content_extractor.py @@ -18,7 +18,7 @@ import html as html_mod import logging -from typing import Any, Dict, List, Optional, Set, Tuple +from typing import Any, List, Optional, Set, Tuple try: import fitz # PyMuPDF (AGPL-3.0) — optional dependency @@ -27,7 +27,6 @@ from contextifier.pipeline.content_extractor import BaseContentExtractor from contextifier.types import ( - ExtractionResult, PreprocessedData, TableData, ) @@ -70,10 +69,7 @@ def _table_to_html(data: List[List[Optional[str]]]) -> str: rows_out: list[str] = [] for ri, row in enumerate(data): tag = "th" if ri == 0 else "td" - cells = "".join( - f"<{tag}>{_escape(cell or '')}" - for cell in row - ) + cells = "".join(f"<{tag}>{_escape(cell or '')}" for cell in row) rows_out.append(f"{cells}") return "" + "".join(rows_out) + "
" @@ -144,7 +140,10 @@ def extract_text( for page_num in range(page_count): page = doc.load_page(page_num) page_text = self._process_page( - doc, page, page_num, processed_images, + doc, + page, + page_num, + processed_images, ) if page_text.strip(): tag = self._make_page_tag(page_num + 1) @@ -220,7 +219,9 @@ def _extract_scan_pages(self, doc: Any) -> str: render_dpi = 150 if self._config is not None: render_dpi = self._config.get_format_option( - "pdf", "render_dpi", render_dpi, + "pdf", + "render_dpi", + render_dpi, ) zoom = render_dpi / 72.0 mat = fitz.Matrix(zoom, zoom) @@ -245,7 +246,9 @@ def _extract_scan_pages(self, doc: Any) -> str: parts.append(f"{page_tag}\n[Image: scan_page_{page_num + 1}.png]") except Exception as exc: logger.warning( - "Failed to render scan page %d: %s", page_num + 1, exc, + "Failed to render scan page %d: %s", + page_num + 1, + exc, ) return "\n\n".join(parts) @@ -260,7 +263,7 @@ def _process_page( processed_images: Set[int], ) -> str: """Process a single page: tables → text → images → merge.""" - elements: list[Tuple[float, str]] = [] # (y_pos, content) + elements: list[Tuple[float, str]] = [] # (y_pos, content) # 1. Tables tables = self._detect_tables(page) @@ -281,7 +284,11 @@ def _process_page( # 3. Images image_elements = self._extract_images( - doc, page, page_num, processed_images, table_bboxes, + doc, + page, + page_num, + processed_images, + table_bboxes, ) elements.extend(image_elements) @@ -292,7 +299,8 @@ def _process_page( # ── Table detection ────────────────────────────────────────────────── def _detect_tables( - self, page: Any, + self, + page: Any, ) -> List[Tuple[List[List[Optional[str]]], Tuple]]: """Detect tables on *page* via PyMuPDF find_tables().""" results: list[Tuple[List[List[Optional[str]]], Tuple]] = [] @@ -362,10 +370,14 @@ def _extract_images( # Override defaults from format_options if available if self._config is not None: min_size = self._config.get_format_option( - "pdf", "min_image_size", min_size, + "pdf", + "min_image_size", + min_size, ) min_area = self._config.get_format_option( - "pdf", "min_image_area", min_area, + "pdf", + "min_image_area", + min_area, ) elements: list[Tuple[float, str]] = [] diff --git a/contextifier/handlers/pdf_plus/_block_image_engine.py b/contextifier/handlers/pdf_plus/_block_image_engine.py index ae5d996..64bfcdc 100644 --- a/contextifier/handlers/pdf_plus/_block_image_engine.py +++ b/contextifier/handlers/pdf_plus/_block_image_engine.py @@ -113,7 +113,9 @@ def process_region( except Exception as exc: logger.error( "[BlockImageEngine] Error processing %s on page %d: %s", - region_type, self.page_num + 1, exc, + region_type, + self.page_num + 1, + exc, ) return BlockResult(success=False, bbox=bbox, error=str(exc)) @@ -167,6 +169,7 @@ def _try_semantic(self) -> MultiBlockResult: from contextifier.handlers.pdf_plus._layout_block_detector import ( LayoutBlockDetector, ) + detector = LayoutBlockDetector(self.page, self.page_num) layout = detector.detect() @@ -175,15 +178,14 @@ def _try_semantic(self) -> MultiBlockResult: results: List[BlockResult] = [] for block in layout.blocks: - area = ( - (block.bbox[2] - block.bbox[0]) - * (block.bbox[3] - block.bbox[1]) - ) + area = (block.bbox[2] - block.bbox[0]) * (block.bbox[3] - block.bbox[1]) if area < CFG.BLOCK_MIN_AREA: continue br = self.process_region( block.bbox, - region_type=block.block_type.name if block.block_type else "unknown", + region_type=block.block_type.name + if block.block_type + else "unknown", ) if br.success: results.append(br) @@ -201,12 +203,15 @@ def _try_semantic(self) -> MultiBlockResult: except Exception as exc: logger.warning( "[BlockImageEngine] Semantic strategy failed on page %d: %s", - self.page_num + 1, exc, + self.page_num + 1, + exc, ) return self._fullpage_fallback() def _try_grid( - self, rows: int = 2, cols: int = 2, + self, + rows: int = 2, + cols: int = 2, ) -> MultiBlockResult: try: cell_w = self.page_width / cols @@ -238,7 +243,8 @@ def _try_grid( except Exception as exc: logger.warning( "[BlockImageEngine] Grid strategy failed on page %d: %s", - self.page_num + 1, exc, + self.page_num + 1, + exc, ) return self._fullpage_fallback() @@ -297,7 +303,8 @@ def _save_image(self, img_bytes: bytes) -> Optional[str]: return None def _is_empty_region( - self, bbox: Tuple[float, float, float, float], + self, + bbox: Tuple[float, float, float, float], ) -> bool: """Check if a region is nearly all white (skip rendering).""" try: @@ -311,9 +318,7 @@ def _is_empty_region( if not pixels: return True thr = CFG.BLOCK_EMPTY_PIXEL_MIN # 240 - whites = sum( - 1 for p in pixels if p[0] > thr and p[1] > thr and p[2] > thr - ) + whites = sum(1 for p in pixels if p[0] > thr and p[1] > thr and p[2] > thr) return whites / len(pixels) >= CFG.BLOCK_EMPTY_THRESHOLD # 0.95 except Exception: return False diff --git a/contextifier/handlers/pdf_plus/_cell_analysis.py b/contextifier/handlers/pdf_plus/_cell_analysis.py index 75751f8..f98c4f0 100644 --- a/contextifier/handlers/pdf_plus/_cell_analysis.py +++ b/contextifier/handlers/pdf_plus/_cell_analysis.py @@ -16,7 +16,7 @@ from __future__ import annotations import logging -from typing import Any, Dict, List, Optional, Set, Tuple +from typing import Any, Dict, List, Optional from contextifier.handlers.pdf_plus._types import CellInfo, PdfPlusConfig @@ -31,7 +31,13 @@ class CellAnalysisEngine: GRID_TOLERANCE: float = CFG.CELL_GRID_TOLERANCE OVERLAP_THRESHOLD: float = CFG.CELL_OVERLAP_THRESHOLD - def __init__(self, data: List[List], cells_info: List[CellInfo], bbox: Optional[tuple] = None, page: Any = None) -> None: + def __init__( + self, + data: List[List], + cells_info: List[CellInfo], + bbox: Optional[tuple] = None, + page: Any = None, + ) -> None: self.data = data or [] self.cells_info = cells_info or [] self.bbox = bbox @@ -82,12 +88,8 @@ def analyze(self) -> List[Dict]: def _has_valid_spans(self) -> bool: if not self.cells_info: return False - has_span = any( - (c.rowspan > 1 or c.colspan > 1) for c in self.cells_info - ) - has_pos = all( - c.row is not None and c.col is not None for c in self.cells_info - ) + has_span = any((c.rowspan > 1 or c.colspan > 1) for c in self.cells_info) + has_pos = all(c.row is not None and c.col is not None for c in self.cells_info) return has_span or has_pos def _use_existing(self, nr: int, nc: int) -> List[Dict]: @@ -101,7 +103,9 @@ def _use_existing(self, nr: int, nc: int) -> List[Dict]: cs = min(max(1, c.colspan), nc - co) if (r, co) in covered: continue - out.append({"row": r, "col": co, "rowspan": rs, "colspan": cs, "bbox": c.bbox}) + out.append( + {"row": r, "col": co, "rowspan": rs, "colspan": cs, "bbox": c.bbox} + ) for ri in range(r, r + rs): for ci in range(co, co + cs): covered.add((ri, ci)) @@ -109,7 +113,9 @@ def _use_existing(self, nr: int, nc: int) -> List[Dict]: for ri in range(nr): for ci in range(nc): if (ri, ci) not in covered: - out.append({"row": ri, "col": ci, "rowspan": 1, "colspan": 1, "bbox": None}) + out.append( + {"row": ri, "col": ci, "rowspan": 1, "colspan": 1, "bbox": None} + ) return out # ------------------------------------------------------------------ @@ -118,7 +124,13 @@ def _use_existing(self, nr: int, nc: int) -> List[Dict]: def _cells_as_dicts(self) -> List[Dict]: return [ - {"row": c.row, "col": c.col, "rowspan": c.rowspan, "colspan": c.colspan, "bbox": c.bbox} + { + "row": c.row, + "col": c.col, + "rowspan": c.rowspan, + "colspan": c.colspan, + "bbox": c.bbox, + } for c in self.cells_info ] @@ -166,7 +178,9 @@ def _analyze_with_bbox(self, cells: list[dict], nr: int, nc: int) -> list[dict]: cspan = min(max(1, ce_idx - cs_idx), nc - dc) if (dr, dc) in covered: continue - out.append({"row": dr, "col": dc, "rowspan": rspan, "colspan": cspan, "bbox": bb}) + out.append( + {"row": dr, "col": dc, "rowspan": rspan, "colspan": cspan, "bbox": bb} + ) for ri in range(dr, min(dr + rspan, nr)): for ci in range(dc, min(dc + cspan, nc)): covered.add((ri, ci)) @@ -174,7 +188,9 @@ def _analyze_with_bbox(self, cells: list[dict], nr: int, nc: int) -> list[dict]: for ri in range(nr): for ci in range(nc): if (ri, ci) not in covered: - out.append({"row": ri, "col": ci, "rowspan": 1, "colspan": 1, "bbox": None}) + out.append( + {"row": ri, "col": ci, "rowspan": 1, "colspan": 1, "bbox": None} + ) return out # ------------------------------------------------------------------ @@ -192,14 +208,24 @@ def _validate_and_enhance(self, cells: list[dict], nr: int, nc: int) -> list[dic cs = min(c.get("colspan", 1), nc - co) if (r, co) in covered: continue - out.append({"row": r, "col": co, "rowspan": max(1, rs), "colspan": max(1, cs), "bbox": c.get("bbox")}) + out.append( + { + "row": r, + "col": co, + "rowspan": max(1, rs), + "colspan": max(1, cs), + "bbox": c.get("bbox"), + } + ) for ri in range(r, min(r + rs, nr)): for ci in range(co, min(co + cs, nc)): covered.add((ri, ci)) for ri in range(nr): for ci in range(nc): if (ri, ci) not in covered: - out.append({"row": ri, "col": ci, "rowspan": 1, "colspan": 1, "bbox": None}) + out.append( + {"row": ri, "col": ci, "rowspan": 1, "colspan": 1, "bbox": None} + ) return out # ------------------------------------------------------------------ @@ -210,7 +236,8 @@ def _validate_and_enhance(self, cells: list[dict], nr: int, nc: int) -> list[dic def _default_cells(nr: int, nc: int) -> list[dict]: return [ {"row": r, "col": c, "rowspan": 1, "colspan": 1, "bbox": None} - for r in range(nr) for c in range(nc) + for r in range(nr) + for c in range(nc) ] # ------------------------------------------------------------------ diff --git a/contextifier/handlers/pdf_plus/_complexity_analyzer.py b/contextifier/handlers/pdf_plus/_complexity_analyzer.py index fa10447..c3aa6e6 100644 --- a/contextifier/handlers/pdf_plus/_complexity_analyzer.py +++ b/contextifier/handlers/pdf_plus/_complexity_analyzer.py @@ -69,9 +69,7 @@ def analyze(self) -> PageComplexity: text_dict = self._get_text_dict() images = self._get_images() - text_blocks = [ - b for b in text_dict.get("blocks", []) if b.get("type") == 0 - ] + text_blocks = [b for b in text_dict.get("blocks", []) if b.get("type") == 0] column_count = self._count_columns(text_blocks) d_score = self._drawing_score(drawings) @@ -85,7 +83,8 @@ def analyze(self) -> PageComplexity: # region-wise regions = self._region_analysis(drawings, text_blocks, images) complex_regions = [ - r.bbox for r in regions + r.bbox + for r in regions if r.complexity_level in (ComplexityLevel.COMPLEX, ComplexityLevel.EXTREME) ] @@ -106,7 +105,11 @@ def analyze(self) -> PageComplexity: ) logger.debug( "[ComplexityAnalyzer] Page %d: %s (%.2f) → %s cols=%d", - self.page_num + 1, level.name, overall, strategy.name, column_count, + self.page_num + 1, + level.name, + overall, + strategy.name, + column_count, ) return result @@ -196,7 +199,11 @@ def _layout_score(self, col_count: int, text_blocks: List[Dict]) -> float: # ───────────────────────────────────────────────────────────────────── def _weighted_score( - self, drawing: float, image: float, text_quality: float, layout: float, + self, + drawing: float, + image: float, + text_quality: float, + layout: float, ) -> float: if layout >= 0.95: return 0.90 # cap — other factors needed for EXTREME @@ -240,11 +247,13 @@ def _region_analysis( area = (x1 - x0) * (y1 - y0) rd = [ - d for d in drawings + d + for d in drawings if d.get("rect") and self._overlaps(bbox, tuple(d["rect"])) ] rt = [ - b for b in text_blocks + b + for b in text_blocks if self._overlaps(bbox, b.get("bbox", (0, 0, 0, 0))) ] @@ -266,14 +275,16 @@ def _region_analysis( lev = ComplexityLevel.SIMPLE strat = ProcessingStrategy.TEXT_EXTRACTION - regions.append(RegionComplexity( - bbox=bbox, - complexity_level=lev, - complexity_score=score, - drawing_density=dd, - text_quality=tq, - recommended_strategy=strat, - )) + regions.append( + RegionComplexity( + bbox=bbox, + complexity_level=lev, + complexity_score=score, + drawing_density=dd, + text_quality=tq, + recommended_strategy=strat, + ) + ) return regions # ───────────────────────────────────────────────────────────────────── @@ -297,9 +308,7 @@ def _strategy( # large complex area + poor text → full OCR if complex_regions: - total_area = sum( - (r[2] - r[0]) * (r[3] - r[1]) for r in complex_regions - ) + total_area = sum((r[2] - r[0]) * (r[3] - r[1]) for r in complex_regions) if total_area / max(1, self.page_area) > 0.5 and text_quality < 0.7: return ProcessingStrategy.FULL_PAGE_OCR diff --git a/contextifier/handlers/pdf_plus/_graphic_detector.py b/contextifier/handlers/pdf_plus/_graphic_detector.py index cf7a581..0d9a1a4 100644 --- a/contextifier/handlers/pdf_plus/_graphic_detector.py +++ b/contextifier/handlers/pdf_plus/_graphic_detector.py @@ -102,9 +102,13 @@ def _cluster_drawings(self, drawings: list[dict]) -> list[GraphicRegionInfo]: "colors": set(), } if fill: - rd["colors"].add(tuple(fill) if isinstance(fill, (list, tuple)) else fill) + rd["colors"].add( + tuple(fill) if isinstance(fill, (list, tuple)) else fill + ) if stroke: - rd["colors"].add(tuple(stroke) if isinstance(stroke, (list, tuple)) else stroke) + rd["colors"].add( + tuple(stroke) if isinstance(stroke, (list, tuple)) else stroke + ) merged = False for existing in raw: @@ -219,10 +223,16 @@ def _analyze_region(self, region: GraphicRegionInfo) -> None: # 5. Chart pattern (curves + fills) if region.curve_count >= 5 and region.fill_count >= 3: score += 0.3 - reasons.append(f"chart_pattern(c={region.curve_count},f={region.fill_count})") + reasons.append( + f"chart_pattern(c={region.curve_count},f={region.fill_count})" + ) # 6. Only-rectangles penalty (table cells, not graphics) - if region.rect_count >= 5 and region.curve_count == 0 and region.line_count == 0: + if ( + region.rect_count >= 5 + and region.curve_count == 0 + and region.line_count == 0 + ): if region.color_count >= 3: score += 0.2 reasons.append(f"colored_rects({region.rect_count})") diff --git a/contextifier/handlers/pdf_plus/_image_extractor.py b/contextifier/handlers/pdf_plus/_image_extractor.py index 3789a63..78e06aa 100644 --- a/contextifier/handlers/pdf_plus/_image_extractor.py +++ b/contextifier/handlers/pdf_plus/_image_extractor.py @@ -46,7 +46,6 @@ def extract_images( Each qualifying image is saved via *image_service* (if available), producing an ```` tag in the returned :class:`PageElement`. """ - import fitz # PyMuPDF elements: list[PageElement] = [] img_list = page.get_images(full=True) @@ -74,7 +73,9 @@ def extract_images( bbox = find_image_position(page, xref) # Table overlap filter - if bbox and is_inside_any_bbox(bbox, table_bboxes, threshold=overlap_threshold): + if bbox and is_inside_any_bbox( + bbox, table_bboxes, threshold=overlap_threshold + ): continue # Convert to PNG @@ -96,17 +97,22 @@ def extract_images( else: tag = f"[Image: page {page_num + 1}, index {img_idx}]" - y_pos = bbox[1] if bbox else 0.0 - x_pos = bbox[0] if bbox else 0.0 - elements.append(PageElement( - element_type=ElementType.IMAGE, - content=tag, - bbox=bbox or (0, 0, float(width), float(height)), - page_num=page_num, - )) + bbox[1] if bbox else 0.0 + bbox[0] if bbox else 0.0 + elements.append( + PageElement( + element_type=ElementType.IMAGE, + content=tag, + bbox=bbox or (0, 0, float(width), float(height)), + page_num=page_num, + ) + ) except Exception as exc: logger.debug( - "[ImgExtractor] page %d xref %d error: %s", page_num + 1, xref, exc, + "[ImgExtractor] page %d xref %d error: %s", + page_num + 1, + xref, + exc, ) return elements @@ -116,6 +122,7 @@ def extract_images( # Helpers # ────────────────────────────────────────────────────────────────────── + def _to_png(doc: Any, xref: int, base_img: dict) -> Optional[bytes]: """Convert extracted image data to PNG bytes.""" import fitz diff --git a/contextifier/handlers/pdf_plus/_layout_block_detector.py b/contextifier/handlers/pdf_plus/_layout_block_detector.py index ee0cf44..3ea4918 100644 --- a/contextifier/handlers/pdf_plus/_layout_block_detector.py +++ b/contextifier/handlers/pdf_plus/_layout_block_detector.py @@ -19,7 +19,7 @@ import logging from collections import defaultdict -from dataclasses import dataclass, field +from dataclasses import dataclass from typing import Any, Dict, List, Optional, Set, Tuple from contextifier.handlers.pdf_plus._types import ( @@ -42,7 +42,8 @@ @dataclass class _ContentElement: """One content element on a page.""" - element_type: str # 'text', 'image', 'drawing' + + element_type: str # 'text', 'image', 'drawing' bbox: Tuple[float, float, float, float] content: Optional[str] = None font_size: float = 0.0 @@ -54,6 +55,7 @@ class _ContentElement: @dataclass class _ColumnInfo: """Detected column region.""" + index: int x_start: float x_end: float @@ -81,12 +83,12 @@ class LayoutBlockDetector: """ # ── configuration shortcuts ────────────────────────────────────────── - _GAP = CFG.LBD_GAP_THRESHOLD # 20 pt column gap + _GAP = CFG.LBD_GAP_THRESHOLD # 20 pt column gap _HEADER_RATIO = CFG.LBD_HEADER_FOOTER_RATIO # 10 % margin - _HEADER_MAX_H = CFG.LBD_HEADER_MAX_HEIGHT # 60 pt - _VERT_DIST = CFG.LBD_VERT_CLUSTER_DIST # 40 pt merge distance - _HORIZ_DIST = CFG.LBD_HORIZ_CLUSTER_DIST # 15 pt - _MIN_BOX_AREA = CFG.LBD_MIN_BOX_AREA # 10 000 pt² + _HEADER_MAX_H = CFG.LBD_HEADER_MAX_HEIGHT # 60 pt + _VERT_DIST = CFG.LBD_VERT_CLUSTER_DIST # 40 pt merge distance + _HORIZ_DIST = CFG.LBD_HORIZ_CLUSTER_DIST # 15 pt + _MIN_BOX_AREA = CFG.LBD_MIN_BOX_AREA # 10 000 pt² _MAX_BLOCKS = CFG.LBD_MAX_BLOCKS_BEFORE_MERGE # 15 _HEADLINE_MIN_SIZE = 14.0 @@ -95,9 +97,9 @@ class LayoutBlockDetector: _CAPTION_MAX_H = 50.0 _SEPARATOR_LENGTH_RATIO = 0.30 _SEPARATOR_MAX_THICK = 3.0 - _MIN_BLOCK_W = CFG.BLOCK_MIN_REGION_WIDTH # 80 pt + _MIN_BLOCK_W = CFG.BLOCK_MIN_REGION_WIDTH # 80 pt _MIN_BLOCK_H = CFG.BLOCK_MIN_REGION_HEIGHT # 60 pt - _MIN_BLOCK_AREA = CFG.BLOCK_MIN_AREA # 15 000 pt² + _MIN_BLOCK_AREA = CFG.BLOCK_MIN_AREA # 15 000 pt² _TARGET_MAX_BLOCKS = 10 _X_CLUSTER_TOL = CFG.COLUMN_CLUSTER_TOLERANCE # 50 pt @@ -143,7 +145,9 @@ def detect(self) -> LayoutAnalysisResult: # Phase 4 — semantic clustering try: blocks = self._cluster_into_blocks( - columns, header_region, footer_region, + columns, + header_region, + footer_region, ) except Exception: blocks = self._create_column_fallback(columns) @@ -157,7 +161,8 @@ def detect(self) -> LayoutAnalysisResult: except Exception as exc: logger.error( "[LayoutBlockDetector] Critical error on page %d: %s", - self.page_num + 1, exc, + self.page_num + 1, + exc, ) blocks = [ LayoutBlock( @@ -202,14 +207,16 @@ def _extract_elements(self) -> None: text += span.get("text", "") stripped = text.strip() if stripped: - self._elements.append(_ContentElement( - element_type="text", - bbox=bbox, - content=stripped, - font_size=max_fs, - is_bold=bold, - text_length=len(stripped), - )) + self._elements.append( + _ContentElement( + element_type="text", + bbox=bbox, + content=stripped, + font_size=max_fs, + is_bold=bold, + text_length=len(stripped), + ) + ) for img_info in self._get_images(): xref = img_info[0] @@ -217,11 +224,13 @@ def _extract_elements(self) -> None: for rect in self.page.get_image_rects(xref): b = (rect.x0, rect.y0, rect.x1, rect.y1) area = (b[2] - b[0]) * (b[3] - b[1]) - self._elements.append(_ContentElement( - element_type="image", - bbox=b, - image_area=area, - )) + self._elements.append( + _ContentElement( + element_type="image", + bbox=b, + image_area=area, + ) + ) break # first occurrence only except Exception: pass @@ -245,12 +254,16 @@ def _extract_separators_and_boxes(self) -> None: continue # horizontal separator - if (h <= self._SEPARATOR_MAX_THICK - and w >= self.page_width * self._SEPARATOR_LENGTH_RATIO): + if ( + h <= self._SEPARATOR_MAX_THICK + and w >= self.page_width * self._SEPARATOR_LENGTH_RATIO + ): self._separators.append((x0, y0, x1, y1)) # vertical separator - elif (w <= self._SEPARATOR_MAX_THICK - and h >= self.page_height * self._SEPARATOR_LENGTH_RATIO * 0.5): + elif ( + w <= self._SEPARATOR_MAX_THICK + and h >= self.page_height * self._SEPARATOR_LENGTH_RATIO * 0.5 + ): self._separators.append((x0, y0, x1, y1)) # box (advertisement / infobox candidate) elif w > 50 and h > 50 and w * h >= self._MIN_BOX_AREA: @@ -313,8 +326,10 @@ def _cluster_x_positions(self, xs: List[float]) -> List[List[float]]: def _detect_header_footer( self, - ) -> Tuple[Optional[Tuple[float, float, float, float]], - Optional[Tuple[float, float, float, float]]]: + ) -> Tuple[ + Optional[Tuple[float, float, float, float]], + Optional[Tuple[float, float, float, float]], + ]: h_bound = self.page_height * self._HEADER_RATIO f_bound = self.page_height * (1 - self._HEADER_RATIO) @@ -323,12 +338,16 @@ def _detect_header_footer( return header, footer def _region_for_band( - self, y_min: float, y_max: float, + self, + y_min: float, + y_max: float, ) -> Optional[Tuple[float, float, float, float]]: elems = [ - e for e in self._elements + e + for e in self._elements if e.element_type == "text" - and e.bbox[1] >= y_min - 5 and e.bbox[3] <= y_max + 5 + and e.bbox[1] >= y_min - 5 + and e.bbox[3] <= y_max + 5 ] if not elems: return None @@ -363,12 +382,14 @@ def _cluster_into_blocks( main_elems.append(e) if header_elems: - blocks.append(LayoutBlock( - block_type=LayoutBlockType.HEADER, - bbox=self._merge_bboxes([e.bbox for e in header_elems]), - elements=[], - confidence=0.9, - )) + blocks.append( + LayoutBlock( + block_type=LayoutBlockType.HEADER, + bbox=self._merge_bboxes([e.bbox for e in header_elems]), + elements=[], + confidence=0.9, + ) + ) # process per column for col in columns: @@ -387,45 +408,54 @@ def _cluster_into_blocks( h = bbox[3] - bbox[1] if w < self._MIN_BLOCK_W or h < self._MIN_BLOCK_H: continue - blocks.append(LayoutBlock( - block_type=LayoutBlockType.UNKNOWN, - bbox=bbox, - elements=[], - confidence=0.7, - column_index=col.index, - )) + blocks.append( + LayoutBlock( + block_type=LayoutBlockType.UNKNOWN, + bbox=bbox, + elements=[], + confidence=0.7, + column_index=col.index, + ) + ) if footer_elems: - blocks.append(LayoutBlock( - block_type=LayoutBlockType.FOOTER, - bbox=self._merge_bboxes([e.bbox for e in footer_elems]), - elements=[], - confidence=0.9, - )) + blocks.append( + LayoutBlock( + block_type=LayoutBlockType.FOOTER, + bbox=self._merge_bboxes([e.bbox for e in footer_elems]), + elements=[], + confidence=0.9, + ) + ) return blocks def _create_column_fallback( - self, columns: List[_ColumnInfo], + self, + columns: List[_ColumnInfo], ) -> List[LayoutBlock]: blocks: List[LayoutBlock] = [] for col in columns: col_elems = [e for e in self._elements if self._elem_in_col(e, col)] if col_elems: - blocks.append(LayoutBlock( - block_type=LayoutBlockType.COLUMN_BLOCK, - bbox=self._merge_bboxes([e.bbox for e in col_elems]), - elements=[], - confidence=0.5, - column_index=col.index, - )) + blocks.append( + LayoutBlock( + block_type=LayoutBlockType.COLUMN_BLOCK, + bbox=self._merge_bboxes([e.bbox for e in col_elems]), + elements=[], + confidence=0.5, + column_index=col.index, + ) + ) if not blocks: - blocks.append(LayoutBlock( - block_type=LayoutBlockType.UNKNOWN, - bbox=(0, 0, self.page_width, self.page_height), - elements=[], - confidence=0.1, - )) + blocks.append( + LayoutBlock( + block_type=LayoutBlockType.UNKNOWN, + bbox=(0, 0, self.page_width, self.page_height), + elements=[], + confidence=0.1, + ) + ) return blocks # ── splitting / clustering helpers ─────────────────────────────────── @@ -447,7 +477,8 @@ def _split_by_separators( groups: List[List[_ContentElement]] = [] for i in range(len(bounds) - 1): g = [ - e for e in elements + e + for e in elements if e.bbox[1] >= bounds[i] - 5 and e.bbox[3] <= bounds[i + 1] + 5 ] if g: @@ -455,7 +486,8 @@ def _split_by_separators( return groups or [elements] def _cluster_adjacent( - self, elements: List[_ContentElement], + self, + elements: List[_ContentElement], ) -> List[List[_ContentElement]]: if len(elements) <= 1: return [elements] @@ -508,11 +540,13 @@ def _classify_blocks(self, blocks: List[LayoutBlock]) -> None: def _determine_block_type(self, block: LayoutBlock) -> LayoutBlockType: # re-collect elements that fall inside the block bbox texts = [ - e for e in self._elements + e + for e in self._elements if e.element_type == "text" and self._is_inside(e.bbox, block.bbox, 5) ] images = [ - e for e in self._elements + e + for e in self._elements if e.element_type == "image" and self._is_inside(e.bbox, block.bbox, 5) ] @@ -525,10 +559,11 @@ def _determine_block_type(self, block: LayoutBlock) -> LayoutBlockType: return LayoutBlockType.STANDALONE_IMAGE if has_text: max_fs = max((t.font_size for t in texts), default=0) - avg_fs = ( - sum(t.font_size for t in texts) / len(texts) if texts else 0 - ) - if max_fs >= self._HEADLINE_MIN_SIZE and max_fs >= avg_fs * self._HEADLINE_FONT_RATIO: + avg_fs = sum(t.font_size for t in texts) / len(texts) if texts else 0 + if ( + max_fs >= self._HEADLINE_MIN_SIZE + and max_fs >= avg_fs * self._HEADLINE_FONT_RATIO + ): return LayoutBlockType.ARTICLE if self._bbox_in_box(block.bbox): total_len = sum(t.text_length for t in texts) @@ -559,7 +594,8 @@ def _optimize_and_sort( headers = [b for b in blocks if b.block_type == LayoutBlockType.HEADER] footers = [b for b in blocks if b.block_type == LayoutBlockType.FOOTER] main = [ - b for b in blocks + b + for b in blocks if b.block_type not in (LayoutBlockType.HEADER, LayoutBlockType.FOOTER) ] @@ -610,7 +646,10 @@ def _merge_small(self, blocks: List[LayoutBlock]) -> List[LayoutBlock]: return result def _should_merge( - self, b1: LayoutBlock, b2: LayoutBlock, aggressive: bool, + self, + b1: LayoutBlock, + b2: LayoutBlock, + aggressive: bool, ) -> bool: if not aggressive and b1.column_index != b2.column_index: return False @@ -639,7 +678,10 @@ def _force_merge(self, blocks: List[LayoutBlock]) -> List[LayoutBlock]: group: List[LayoutBlock] = [] for i, b in enumerate(cbs): group.append(b) - if len(group) >= per and len(result) < len(col_groups) * target_groups - 1: + if ( + len(group) >= per + and len(result) < len(col_groups) * target_groups - 1 + ): result.append(self._merge_group(group)) group = [] if group: diff --git a/contextifier/handlers/pdf_plus/_line_analysis.py b/contextifier/handlers/pdf_plus/_line_analysis.py index 4d5d6ce..6343028 100644 --- a/contextifier/handlers/pdf_plus/_line_analysis.py +++ b/contextifier/handlers/pdf_plus/_line_analysis.py @@ -62,12 +62,8 @@ def build_grid(self, tolerance: float | None = None) -> Optional[GridInfo]: if not self.h_lines and not self.v_lines: return None - h_positions = self._cluster_positions( - [ln.y0 for ln in self.h_lines], tolerance - ) - v_positions = self._cluster_positions( - [ln.x0 for ln in self.v_lines], tolerance - ) + h_positions = self._cluster_positions([ln.y0 for ln in self.h_lines], tolerance) + v_positions = self._cluster_positions([ln.x0 for ln in self.v_lines], tolerance) if len(h_positions) < 2 or len(v_positions) < 2: return None @@ -98,16 +94,28 @@ def reconstruct_incomplete_border(self, grid: GridInfo) -> GridInfo: x_min, x_max = min(v_lines), max(v_lines) reconstructed = False - if not any(abs(y - y_min) < CFG.LINE_MERGE_TOLERANCE for y in h_lines) and len(h_lines) >= 2: + if ( + not any(abs(y - y_min) < CFG.LINE_MERGE_TOLERANCE for y in h_lines) + and len(h_lines) >= 2 + ): h_lines.insert(0, y_min - CFG.BORDER_EXTENSION_MARGIN) reconstructed = True - if not any(abs(y - y_max) < CFG.LINE_MERGE_TOLERANCE for y in h_lines) and len(h_lines) >= 2: + if ( + not any(abs(y - y_max) < CFG.LINE_MERGE_TOLERANCE for y in h_lines) + and len(h_lines) >= 2 + ): h_lines.append(y_max + CFG.BORDER_EXTENSION_MARGIN) reconstructed = True - if not any(abs(x - x_min) < CFG.LINE_MERGE_TOLERANCE for x in v_lines) and len(v_lines) >= 2: + if ( + not any(abs(x - x_min) < CFG.LINE_MERGE_TOLERANCE for x in v_lines) + and len(v_lines) >= 2 + ): v_lines.insert(0, x_min - CFG.BORDER_EXTENSION_MARGIN) reconstructed = True - if not any(abs(x - x_max) < CFG.LINE_MERGE_TOLERANCE for x in v_lines) and len(v_lines) >= 2: + if ( + not any(abs(x - x_max) < CFG.LINE_MERGE_TOLERANCE for x in v_lines) + and len(v_lines) >= 2 + ): v_lines.append(x_max + CFG.BORDER_EXTENSION_MARGIN) reconstructed = True @@ -149,7 +157,10 @@ def _extract_all_lines(self) -> None: self.all_lines.append( LineInfo( - x0=x0, y0=y0, x1=x1, y1=y1, + x0=x0, + y0=y0, + x1=x1, + y1=y1, thickness=stroke_width, thickness_class=self._classify_thickness(stroke_width), is_horizontal=is_h, @@ -167,11 +178,14 @@ def _add_line_from_points(self, p1: Any, p2: Any, stroke_width: float) -> None: return self.all_lines.append( LineInfo( - x0=min(x0, x1), y0=min(y0, y1), - x1=max(x0, x1), y1=max(y0, y1), + x0=min(x0, x1), + y0=min(y0, y1), + x1=max(x0, x1), + y1=max(y0, y1), thickness=stroke_width, thickness_class=self._classify_thickness(stroke_width), - is_horizontal=is_h, is_vertical=is_v, + is_horizontal=is_h, + is_vertical=is_v, ) ) @@ -198,13 +212,15 @@ def _merge_double_lines(self) -> None: self.h_lines = self._merge_parallel(self.h_lines, is_horizontal=True) self.v_lines = self._merge_parallel(self.v_lines, is_horizontal=False) - def _merge_parallel(self, lines: List[LineInfo], *, is_horizontal: bool) -> List[LineInfo]: + def _merge_parallel( + self, lines: List[LineInfo], *, is_horizontal: bool + ) -> List[LineInfo]: if len(lines) < 2: return lines if is_horizontal: - sorted_lines = sorted(lines, key=lambda l: (l.y0, l.x0)) + sorted_lines = sorted(lines, key=lambda ln: (ln.y0, ln.x0)) else: - sorted_lines = sorted(lines, key=lambda l: (l.x0, l.y0)) + sorted_lines = sorted(lines, key=lambda ln: (ln.x0, ln.y0)) merged: list[LineInfo] = [] used: set[int] = set() @@ -241,19 +257,25 @@ def _merge_two(a: LineInfo, b: LineInfo, is_h: bool) -> LineInfo: if is_h: avg_y = (a.y0 + b.y0) / 2 return LineInfo( - x0=min(a.x0, b.x0), y0=avg_y, - x1=max(a.x1, b.x1), y1=avg_y, + x0=min(a.x0, b.x0), + y0=avg_y, + x1=max(a.x1, b.x1), + y1=avg_y, thickness=max(a.thickness, b.thickness), thickness_class=thicker.thickness_class, - is_horizontal=True, is_vertical=False, + is_horizontal=True, + is_vertical=False, ) avg_x = (a.x0 + b.x0) / 2 return LineInfo( - x0=avg_x, y0=min(a.y0, b.y0), - x1=avg_x, y1=max(a.y1, b.y1), + x0=avg_x, + y0=min(a.y0, b.y0), + x1=avg_x, + y1=max(a.y1, b.y1), thickness=max(a.thickness, b.thickness), thickness_class=thicker.thickness_class, - is_horizontal=False, is_vertical=True, + is_horizontal=False, + is_vertical=True, ) @staticmethod diff --git a/contextifier/handlers/pdf_plus/_page_analyzer.py b/contextifier/handlers/pdf_plus/_page_analyzer.py index 5ad64e6..5b591f1 100644 --- a/contextifier/handlers/pdf_plus/_page_analyzer.py +++ b/contextifier/handlers/pdf_plus/_page_analyzer.py @@ -9,7 +9,7 @@ from __future__ import annotations import logging -from typing import Any, Dict, List, Tuple +from typing import Any, Tuple from contextifier.handlers.pdf_plus._types import ( PageBorderInfo, @@ -44,8 +44,12 @@ def detect_page_border(page: Any) -> PageBorderInfo: h = y1 - y0 # Must be very thin — one dimension tiny - is_h_line = h <= CFG.PAGE_BORDER_LINE_MAX_SIZE and w >= pw * CFG.PAGE_SPANNING_RATIO - is_v_line = w <= CFG.PAGE_BORDER_LINE_MAX_SIZE and h >= ph * CFG.PAGE_SPANNING_RATIO + is_h_line = ( + h <= CFG.PAGE_BORDER_LINE_MAX_SIZE and w >= pw * CFG.PAGE_SPANNING_RATIO + ) + is_v_line = ( + w <= CFG.PAGE_BORDER_LINE_MAX_SIZE and h >= ph * CFG.PAGE_SPANNING_RATIO + ) if is_h_line: if y0 < margin_y: @@ -87,7 +91,7 @@ def is_table_likely_border( pw, ph = page.rect.width, page.rect.height tw = table_bbox[2] - table_bbox[0] th = table_bbox[3] - table_bbox[1] - return (tw / pw > CFG.PAGE_SPANNING_RATIO and th / ph > CFG.PAGE_SPANNING_RATIO) + return tw / pw > CFG.PAGE_SPANNING_RATIO and th / ph > CFG.PAGE_SPANNING_RATIO __all__ = ["detect_page_border", "is_table_likely_border"] diff --git a/contextifier/handlers/pdf_plus/_table_detection.py b/contextifier/handlers/pdf_plus/_table_detection.py index 759e9e6..26bd1c4 100644 --- a/contextifier/handlers/pdf_plus/_table_detection.py +++ b/contextifier/handlers/pdf_plus/_table_detection.py @@ -19,7 +19,7 @@ from __future__ import annotations import logging -from typing import Any, Dict, List, Optional, Set, Tuple +from typing import Any, List, Optional, Tuple from contextifier.handlers.pdf_plus._types import ( CellInfo, @@ -47,7 +47,7 @@ class TableDetectionEngine: def __init__(self, page: Any, page_num: int, file_data: bytes) -> None: self.page = page self.page_num = page_num - self.file_data = file_data # raw PDF bytes (for pdfplumber) + self.file_data = file_data # raw PDF bytes (for pdfplumber) self.page_width: float = page.rect.width self.page_height: float = page.rect.height @@ -112,11 +112,17 @@ def _detect_pymupdf(self) -> list[TableCandidate]: if conf < self.CONFIDENCE_THRESHOLD: continue cells = self._cells_from_pymupdf(table, col_map) - cands.append(TableCandidate( - strategy=TableDetectionStrategy.PYMUPDF_NATIVE, - confidence=conf, bbox=table.bbox, grid=None, - cells=cells, data=merged_data, raw_table=table, - )) + cands.append( + TableCandidate( + strategy=TableDetectionStrategy.PYMUPDF_NATIVE, + confidence=conf, + bbox=table.bbox, + grid=None, + cells=cells, + data=merged_data, + raw_table=table, + ) + ) except Exception as exc: logger.debug("[TableDet] pymupdf table err: %s", exc) except Exception as exc: @@ -153,11 +159,17 @@ def _detect_pdfplumber(self) -> list[TableCandidate]: conf = self._plumber_confidence(tdata) if conf < self.CONFIDENCE_THRESHOLD: continue - cands.append(TableCandidate( - strategy=TableDetectionStrategy.PDFPLUMBER_LINES, - confidence=conf, bbox=bbox, grid=None, - cells=[], data=tdata, raw_table=None, - )) + cands.append( + TableCandidate( + strategy=TableDetectionStrategy.PDFPLUMBER_LINES, + confidence=conf, + bbox=bbox, + grid=None, + cells=[], + data=tdata, + raw_table=None, + ) + ) except ImportError: logger.debug("[TableDet] pdfplumber not installed — skipping") except Exception as exc: @@ -185,11 +197,17 @@ def _detect_lines(self) -> list[TableCandidate]: conf = self._line_confidence(grid, data) if conf < self.CONFIDENCE_THRESHOLD: return [] - return [TableCandidate( - strategy=TableDetectionStrategy.HYBRID_ANALYSIS, - confidence=conf, bbox=grid.bbox, grid=grid, - cells=cells, data=data, raw_table=None, - )] + return [ + TableCandidate( + strategy=TableDetectionStrategy.HYBRID_ANALYSIS, + confidence=conf, + bbox=grid.bbox, + grid=grid, + cells=cells, + data=data, + raw_table=None, + ) + ] # ================================================================== # Header–data merging @@ -235,8 +253,10 @@ def _can_merge_hd(self, hd: TableCandidate, dt: TableCandidate) -> bool: def _do_merge_hd(self, hd: TableCandidate, dt: TableCandidate) -> TableCandidate: bbox = ( - min(hd.bbox[0], dt.bbox[0]), hd.bbox[1], - max(hd.bbox[2], dt.bbox[2]), dt.bbox[3], + min(hd.bbox[0], dt.bbox[0]), + hd.bbox[1], + max(hd.bbox[2], dt.bbox[2]), + dt.bbox[3], ) hcols = max((len(r) for r in hd.data), default=0) dcols = max((len(r) for r in dt.data), default=0) @@ -255,18 +275,30 @@ def _do_merge_hd(self, hd: TableCandidate, dt: TableCandidate) -> TableCandidate merged_cells.extend(hd.cells) # Data rows for ri, row in enumerate(dt.data): - adj = list(row) + ([""] * (ncols - len(row))) if len(row) < ncols else list(row) + adj = ( + list(row) + ([""] * (ncols - len(row))) + if len(row) < ncols + else list(row) + ) merged_data.append(adj) for c in dt.cells: - merged_cells.append(CellInfo( - row=c.row + hrc, col=c.col, - rowspan=c.rowspan, colspan=c.colspan, bbox=c.bbox, - )) + merged_cells.append( + CellInfo( + row=c.row + hrc, + col=c.col, + rowspan=c.rowspan, + colspan=c.colspan, + bbox=c.bbox, + ) + ) return TableCandidate( strategy=hd.strategy, confidence=max(hd.confidence, dt.confidence), - bbox=bbox, grid=hd.grid or dt.grid, - cells=merged_cells, data=merged_data, raw_table=None, + bbox=bbox, + grid=hd.grid or dt.grid, + cells=merged_cells, + data=merged_data, + raw_table=None, ) # ================================================================== @@ -324,17 +356,23 @@ def _col_groups(widths: list[float], min_w: float) -> list[list[int]]: groups.append(cur) cur = [] if cur: - (groups[-1] if groups else groups.append(cur) or groups).extend([] if not groups else []) + (groups[-1] if groups else groups.append(cur) or groups).extend( + [] if not groups else [] + ) if cur and groups and cur is not groups[-1]: groups[-1].extend(cur) return groups - def _merge_cols_by_content(self, data: list[list]) -> Tuple[list[list[str]], dict[int, int]]: + def _merge_cols_by_content( + self, data: list[list] + ) -> Tuple[list[list[str]], dict[int, int]]: ncols = max(len(r) for r in data) nrows = len(data) ratios = [] for ci in range(ncols): - empty = sum(1 for r in data if ci >= len(r) or not r[ci] or not str(r[ci]).strip()) + empty = sum( + 1 for r in data if ci >= len(r) or not r[ci] or not str(r[ci]).strip() + ) ratios.append(empty / nrows if nrows else 1.0) groups: list[list[int]] = [] cur: list[int] = [] @@ -365,7 +403,9 @@ def _merge_cols_by_content(self, data: list[list]) -> Tuple[list[list[str]], dic # PyMuPDF cell extraction # ================================================================== - def _cells_from_pymupdf(self, table: Any, col_map: dict[int, int]) -> list[CellInfo]: + def _cells_from_pymupdf( + self, table: Any, col_map: dict[int, int] + ) -> list[CellInfo]: raw = self._extract_raw_cells(table) if not col_map or not raw: return raw @@ -382,7 +422,15 @@ def _cells_from_pymupdf(self, table: Any, col_map: dict[int, int]) -> list[CellI if mc != nc: ncs = max(ncs, mc - nc + 1) ncs = min(ncs, new_nc - nc) - out.append(CellInfo(row=c.row, col=nc, rowspan=c.rowspan, colspan=max(1, ncs), bbox=c.bbox)) + out.append( + CellInfo( + row=c.row, + col=nc, + rowspan=c.rowspan, + colspan=max(1, ncs), + bbox=c.bbox, + ) + ) for ri in range(c.row, c.row + c.rowspan): for ci in range(nc, nc + max(1, ncs)): seen.add((ri, ci)) @@ -416,11 +464,15 @@ def idx(val: float, coords: list[float], tol: float = 3.0) -> int: if (rs, cs) in seen: continue seen.add((rs, cs)) - cells.append(CellInfo( - row=rs, col=cs, - rowspan=max(1, re - rs), colspan=max(1, ce - cs), - bbox=cb, - )) + cells.append( + CellInfo( + row=rs, + col=cs, + rowspan=max(1, re - rs), + colspan=max(1, ce - cs), + bbox=cb, + ) + ) for r in range(rs, rs + max(1, re - rs)): for c in range(cs, cs + max(1, ce - cs)): if (r, c) != (rs, cs): @@ -528,24 +580,41 @@ def _validate(self, cands: list[TableCandidate]) -> list[TableCandidate]: for c in cands: skip_gfx = c.strategy == TableDetectionStrategy.PYMUPDF_NATIVE ok, conf, reason = self.quality_validator.validate( - data=c.data, bbox=c.bbox, + data=c.data, + bbox=c.bbox, cells_info=[ - {"row": ci.row, "col": ci.col, "rowspan": ci.rowspan, - "colspan": ci.colspan, "bbox": ci.bbox} + { + "row": ci.row, + "col": ci.col, + "rowspan": ci.rowspan, + "colspan": ci.colspan, + "bbox": ci.bbox, + } for ci in c.cells - ] if c.cells else None, + ] + if c.cells + else None, skip_graphic_check=skip_gfx, ) if ok: - out.append(TableCandidate( - strategy=c.strategy, - confidence=min(c.confidence, conf), - bbox=c.bbox, grid=c.grid, - cells=c.cells, data=c.data, raw_table=c.raw_table, - )) + out.append( + TableCandidate( + strategy=c.strategy, + confidence=min(c.confidence, conf), + bbox=c.bbox, + grid=c.grid, + cells=c.cells, + data=c.data, + raw_table=c.raw_table, + ) + ) else: - logger.debug("[TableDet] filtered: page=%d bbox=%s reason=%s", - self.page_num + 1, c.bbox, reason) + logger.debug( + "[TableDet] filtered: page=%d bbox=%s reason=%s", + self.page_num + 1, + c.bbox, + reason, + ) return out def _select_best(self, cands: list[TableCandidate]) -> list[TableCandidate]: @@ -606,13 +675,16 @@ def _text_in_bbox(pd: dict, bbox: tuple) -> str: continue for ln in blk.get("lines", []): lb = ln.get("bbox", (0, 0, 0, 0)) - if not (max(lb[0], bbox[0]) < min(lb[2], bbox[2]) - and max(lb[1], bbox[1]) < min(lb[3], bbox[3])): + if not ( + max(lb[0], bbox[0]) < min(lb[2], bbox[2]) + and max(lb[1], bbox[1]) < min(lb[3], bbox[3]) + ): continue for sp in ln.get("spans", []): sb = sp.get("bbox", (0, 0, 0, 0)) - if max(sb[0], bbox[0]) < min(sb[2], bbox[2]) \ - and max(sb[1], bbox[1]) < min(sb[3], bbox[3]): + if max(sb[0], bbox[0]) < min(sb[2], bbox[2]) and max( + sb[1], bbox[1] + ) < min(sb[3], bbox[3]): t = sp.get("text", "").strip() if t: texts.append(t) @@ -621,9 +693,18 @@ def _text_in_bbox(pd: dict, bbox: tuple) -> str: @staticmethod def _cells_from_grid(grid: GridInfo) -> list[CellInfo]: return [ - CellInfo(row=ri, col=ci, rowspan=1, colspan=1, - bbox=(grid.v_lines[ci], grid.h_lines[ri], - grid.v_lines[ci + 1], grid.h_lines[ri + 1])) + CellInfo( + row=ri, + col=ci, + rowspan=1, + colspan=1, + bbox=( + grid.v_lines[ci], + grid.h_lines[ri], + grid.v_lines[ci + 1], + grid.h_lines[ri + 1], + ), + ) for ri in range(grid.row_count) for ci in range(grid.col_count) ] @@ -638,8 +719,14 @@ def _estimate_bbox_plumber(pp: Any, data: list[list]) -> Optional[tuple]: words = pp.extract_words() if not words: return None - table_texts = {str(c).strip()[:20] for r in data for c in r if c and str(c).strip()} - hits = [w for w in words if any(w["text"] in t or t in w["text"] for t in table_texts)] + table_texts = { + str(c).strip()[:20] for r in data for c in r if c and str(c).strip() + } + hits = [ + w + for w in words + if any(w["text"] in t or t in w["text"] for t in table_texts) + ] if not hits: return None m = 5 diff --git a/contextifier/handlers/pdf_plus/_table_processor.py b/contextifier/handlers/pdf_plus/_table_processor.py index d69b47f..e4cc2f1 100644 --- a/contextifier/handlers/pdf_plus/_table_processor.py +++ b/contextifier/handlers/pdf_plus/_table_processor.py @@ -17,11 +17,10 @@ import logging from dataclasses import dataclass, field -from typing import Any, Dict, List, Optional, Tuple +from typing import Any, Dict, List, Tuple from contextifier.handlers.pdf_plus._types import ( AnnotationInfo, - CellInfo, PdfPlusConfig, TableCandidate, ) @@ -38,9 +37,11 @@ # Data containers for processed tables # ────────────────────────────────────────────────────────────────────── + @dataclass class ProcessedTable: """One fully-processed table.""" + page_num: int bbox: Tuple[float, float, float, float] data: List[List[str]] @@ -55,6 +56,7 @@ class ProcessedTable: # TableProcessor # ────────────────────────────────────────────────────────────────────── + class TableProcessor: """ High-level table-processing orchestrator for a *single page*. @@ -92,15 +94,17 @@ def process(self) -> List[ProcessedTable]: cells = self._analyse_cells(cand) annotations = self._detect_annotations(cand) html = self._to_html(cand.data, cells, annotations) - results.append(ProcessedTable( - page_num=self.page_num, - bbox=cand.bbox, - data=cand.data, - cells=cells, - html=html, - annotations=annotations, - confidence=cand.confidence, - )) + results.append( + ProcessedTable( + page_num=self.page_num, + bbox=cand.bbox, + data=cand.data, + cells=cells, + html=html, + annotations=annotations, + confidence=cand.confidence, + ) + ) return results def check_continuity( @@ -130,7 +134,9 @@ def check_continuity( ct.is_continuation = True logger.debug( "[TableProc] page %d table is continuation (prev cols=%d, cur cols=%d)", - ct.page_num + 1, pcols, ccols, + ct.page_num + 1, + pcols, + ccols, ) # ------------------------------------------------------------------ @@ -172,11 +178,13 @@ def _detect_annotations(self, cand: TableCandidate) -> List[AnnotationInfo]: if not text: continue atype = self._classify_annotation(text) - annotations.append(AnnotationInfo( - text=text, - bbox=lb, - annotation_type=atype, - )) + annotations.append( + AnnotationInfo( + text=text, + bbox=lb, + annotation_type=atype, + ) + ) return annotations @staticmethod diff --git a/contextifier/handlers/pdf_plus/_table_quality_analyzer.py b/contextifier/handlers/pdf_plus/_table_quality_analyzer.py index 88461b4..b7eca9d 100644 --- a/contextifier/handlers/pdf_plus/_table_quality_analyzer.py +++ b/contextifier/handlers/pdf_plus/_table_quality_analyzer.py @@ -22,7 +22,7 @@ from __future__ import annotations import logging -from typing import Any, Dict, List, Optional, Tuple +from typing import Any, Dict, Optional, Tuple from contextifier.handlers.pdf_plus._types import ( PdfPlusConfig, @@ -87,7 +87,13 @@ def analyze_table( logger.debug( "[TQAnalyzer] %s quality=%s total=%.2f border=%.2f grid=%.2f cell=%.2f simple=%.2f", - bbox, quality.name, total, border_sc, grid_sc, cell_sc, simple_sc, + bbox, + quality.name, + total, + border_sc, + grid_sc, + cell_sc, + simple_sc, ) return TableQualityResult( bbox=bbox, @@ -111,13 +117,15 @@ def analyze_page_tables(self) -> Dict[str, Any]: results: list[dict] = [] for bbox in table_regions: qr = self.analyze_table(bbox) - results.append({ - "bbox": bbox, - "quality": qr.quality, - "score": qr.score, - "is_processable": qr.text_extractable, - "issues": qr.issues, - }) + results.append( + { + "bbox": bbox, + "quality": qr.quality, + "score": qr.score, + "is_processable": qr.text_extractable, + "issues": qr.issues, + } + ) has_ok = any(r["is_processable"] for r in results) return { "table_candidates": results, @@ -142,7 +150,12 @@ def _border_completeness( return 0.0, ["no_border_lines"] tol = CFG.QUALITY_BORDER_TOLERANCE x0, y0, x1, y1 = bbox - sides: dict[str, bool] = {"top": False, "bottom": False, "left": False, "right": False} + sides: dict[str, bool] = { + "top": False, + "bottom": False, + "left": False, + "right": False, + } for ln in lines: if ln["is_horizontal"]: if abs(ln["y1"] - y0) <= tol: @@ -167,12 +180,12 @@ def _grid_regularity( lines = self._extract_lines(drawings) if not lines: return 0.0, ["no_grid_lines"] - orth_cnt = sum(1 for l in lines if l["is_horizontal"] or l["is_vertical"]) + orth_cnt = sum(1 for ln in lines if ln["is_horizontal"] or ln["is_vertical"]) orth_ratio = orth_cnt / len(lines) if orth_ratio < CFG.QUALITY_MIN_ORTHOGONAL_RATIO: issues.append(f"non_orth({(1 - orth_ratio) * 100:.0f}%)") - h_align = self._alignment([l["y1"] for l in lines if l["is_horizontal"]]) - v_align = self._alignment([l["x1"] for l in lines if l["is_vertical"]]) + h_align = self._alignment([ln["y1"] for ln in lines if ln["is_horizontal"]]) + v_align = self._alignment([ln["x1"] for ln in lines if ln["is_vertical"]]) alignment = (h_align + v_align) / 2 if alignment < 0.8: issues.append("misaligned_grid") @@ -183,18 +196,36 @@ def _cell_structure( ) -> Tuple[float, list[str]]: issues: list[str] = [] lines = self._extract_lines(drawings) - h_sorted = sorted((l for l in lines if l["is_horizontal"]), key=lambda l: l["y1"]) - v_sorted = sorted((l for l in lines if l["is_vertical"]), key=lambda l: l["x1"]) + h_sorted = sorted( + (ln for ln in lines if ln["is_horizontal"]), key=lambda ln: ln["y1"] + ) + v_sorted = sorted( + (ln for ln in lines if ln["is_vertical"]), key=lambda ln: ln["x1"] + ) if len(h_sorted) < 2 or len(v_sorted) < 2: return 0.5, ["insufficient_lines"] - heights = [h_sorted[i + 1]["y1"] - h_sorted[i]["y1"] for i in range(len(h_sorted) - 1) if h_sorted[i + 1]["y1"] > h_sorted[i]["y1"]] - widths = [v_sorted[i + 1]["x1"] - v_sorted[i]["x1"] for i in range(len(v_sorted) - 1) if v_sorted[i + 1]["x1"] > v_sorted[i]["x1"]] - tiny = sum(1 for h in heights if h < CFG.QUALITY_MIN_CELL_SIZE) + sum(1 for w in widths if w < CFG.QUALITY_MIN_CELL_SIZE) + heights = [ + h_sorted[i + 1]["y1"] - h_sorted[i]["y1"] + for i in range(len(h_sorted) - 1) + if h_sorted[i + 1]["y1"] > h_sorted[i]["y1"] + ] + widths = [ + v_sorted[i + 1]["x1"] - v_sorted[i]["x1"] + for i in range(len(v_sorted) - 1) + if v_sorted[i + 1]["x1"] > v_sorted[i]["x1"] + ] + tiny = sum(1 for h in heights if h < CFG.QUALITY_MIN_CELL_SIZE) + sum( + 1 for w in widths if w < CFG.QUALITY_MIN_CELL_SIZE + ) total_dims = len(heights) + len(widths) extreme_ratio = 0 for h in heights: for w in widths: - if h > 0 and w > 0 and max(h / w, w / h) > CFG.QUALITY_MAX_CELL_ASPECT_RATIO: + if ( + h > 0 + and w > 0 + and max(h / w, w / h) > CFG.QUALITY_MAX_CELL_ASPECT_RATIO + ): extreme_ratio += 1 sc = 1.0 if tiny and total_dims and tiny / total_dims > 0.1: @@ -218,7 +249,10 @@ def _element_simplicity(self, drawings: list[dict]) -> Tuple[float, list[str]]: curve_c += 1 elif i[0] == "l": p1, p2 = i[1], i[2] - if abs(p2.x - p1.x) > CFG.QUALITY_LINE_ANGLE_TOLERANCE and abs(p2.y - p1.y) > CFG.QUALITY_LINE_ANGLE_TOLERANCE: + if ( + abs(p2.x - p1.x) > CFG.QUALITY_LINE_ANGLE_TOLERANCE + and abs(p2.y - p1.y) > CFG.QUALITY_LINE_ANGLE_TOLERANCE + ): diag_c += 1 if d.get("fill"): fill_c += 1 @@ -253,7 +287,8 @@ def _get_region_drawings(self, bbox: tuple) -> list[dict]: if self._drawings is None: self._drawings = self.page.get_drawings() return [ - d for d in self._drawings + d + for d in self._drawings if d.get("rect") and self._overlaps(bbox, tuple(d["rect"])) ] @@ -264,8 +299,8 @@ def _find_table_regions(self) -> list[tuple]: if not lines: return [] # quick bbox of all lines - xs = [l["x1"] for l in lines] + [l["x2"] for l in lines] - ys = [l["y1"] for l in lines] + [l["y2"] for l in lines] + xs = [ln["x1"] for ln in lines] + [ln["x2"] for ln in lines] + ys = [ln["y1"] for ln in lines] + [ln["y2"] for ln in lines] bbox = (min(xs), min(ys), max(xs), max(ys)) w, h = bbox[2] - bbox[0], bbox[3] - bbox[1] return [bbox] if w > 100 and h > 50 else [] @@ -278,13 +313,17 @@ def _extract_lines(drawings: list[dict]) -> list[dict]: if item[0] == "l": p1, p2 = item[1], item[2] x1, y1, x2, y2 = p1.x, p1.y, p2.x, p2.y - result.append({ - "x1": min(x1, x2), "y1": min(y1, y2), - "x2": max(x1, x2), "y2": max(y1, y2), - "is_horizontal": abs(y2 - y1) <= 2.0, - "is_vertical": abs(x2 - x1) <= 2.0, - "length": ((x2 - x1) ** 2 + (y2 - y1) ** 2) ** 0.5, - }) + result.append( + { + "x1": min(x1, x2), + "y1": min(y1, y2), + "x2": max(x1, x2), + "y2": max(y1, y2), + "is_horizontal": abs(y2 - y1) <= 2.0, + "is_vertical": abs(x2 - x1) <= 2.0, + "length": ((x2 - x1) ** 2 + (y2 - y1) ** 2) ** 0.5, + } + ) elif item[0] == "re": r = item[1] x0, y0, x1, y1 = r.x0, r.y0, r.x1, r.y1 @@ -294,11 +333,17 @@ def _extract_lines(drawings: list[dict]) -> list[dict]: (x0, y0, x0, y1, False, True), (x1, y0, x1, y1, False, True), ]: - result.append({ - "x1": a, "y1": b, "x2": c, "y2": dd, - "is_horizontal": ih, "is_vertical": iv, - "length": abs(c - a) + abs(dd - b), - }) + result.append( + { + "x1": a, + "y1": b, + "x2": c, + "y2": dd, + "is_horizontal": ih, + "is_vertical": iv, + "length": abs(c - a) + abs(dd - b), + } + ) return result @staticmethod diff --git a/contextifier/handlers/pdf_plus/_table_validator.py b/contextifier/handlers/pdf_plus/_table_validator.py index 18ced45..0e04d49 100644 --- a/contextifier/handlers/pdf_plus/_table_validator.py +++ b/contextifier/handlers/pdf_plus/_table_validator.py @@ -117,9 +117,7 @@ def validate( confidence -= 0.15 # --- 5. valid rows --- - valid_rows = sum( - 1 for row in data if any(c and str(c).strip() for c in row) - ) + valid_rows = sum(1 for row in data if any(c and str(c).strip() for c in row)) if valid_rows < CFG.TABLE_MIN_VALID_ROWS: penalties.append(f"few_valid_rows({valid_rows})") confidence -= 0.15 @@ -145,7 +143,10 @@ def validate( long_cnt, extreme_cnt = self._cell_lengths(data) if extreme_cnt > 0: return False, 0.0, f"extreme_long_cell({extreme_cnt})" - if filled_cells > 0 and long_cnt / filled_cells > CFG.TABLE_MAX_LONG_CELLS_RATIO: + if ( + filled_cells > 0 + and long_cnt / filled_cells > CFG.TABLE_MAX_LONG_CELLS_RATIO + ): penalties.append(f"long_cells({long_cnt / filled_cells:.2f})") confidence -= 0.2 @@ -178,7 +179,9 @@ def validate( if not is_valid: logger.debug( "[TableValidator] Rejected: %s reason=%s conf=%.2f", - bbox, reason, confidence, + bbox, + reason, + confidence, ) return is_valid, confidence, reason @@ -186,13 +189,35 @@ def validate( # Helpers # ------------------------------------------------------------------ - _SIMPLE_SYMS = frozenset(["", "-", "–", "—", ".", ":", ";", "|", "/", - "\\", "*", "#", "@", "!", "?", ",", " "]) + _SIMPLE_SYMS = frozenset( + [ + "", + "-", + "–", + "—", + ".", + ":", + ";", + "|", + "/", + "\\", + "*", + "#", + "@", + "!", + "?", + ",", + " ", + ] + ) def _count_meaningful(self, data: list) -> int: return sum( - 1 for row in data for cell in row - if cell and len(str(cell).strip()) >= 2 + 1 + for row in data + for cell in row + if cell + and len(str(cell).strip()) >= 2 and str(cell).strip() not in self._SIMPLE_SYMS ) @@ -253,11 +278,17 @@ def _validate_two_col(self, data: list, bbox: tuple) -> Tuple[bool, str]: c1_short += 1 if len(c2) > 80: c2_long += 1 - if any(p in c2 for p in [".", "。", ",", "、"]) and len(c2.split()) >= 5: + if ( + any(p in c2 for p in [".", "。", ",", "、"]) + and len(c2.split()) >= 5 + ): c2_para += 1 if nr > 0: if c1_empty / nr >= 0.6 and c2_long / nr >= 0.3: - return False, f"empty_col1({c1_empty / nr:.0%})_long_col2({c2_long / nr:.0%})" + return ( + False, + f"empty_col1({c1_empty / nr:.0%})_long_col2({c2_long / nr:.0%})", + ) if nr > 5 and c2_para >= 2: return False, f"col2_paragraphs({c2_para})" if nr > 10 and (c1_empty + c1_short) / nr >= 0.8 and c2_long >= 5: diff --git a/contextifier/handlers/pdf_plus/_text_extractor.py b/contextifier/handlers/pdf_plus/_text_extractor.py index f5c859a..9b439d1 100644 --- a/contextifier/handlers/pdf_plus/_text_extractor.py +++ b/contextifier/handlers/pdf_plus/_text_extractor.py @@ -18,7 +18,6 @@ PdfPlusConfig, ) from contextifier.handlers.pdf_plus._utils import ( - escape_html, is_inside_any_bbox, ) from contextifier.handlers.pdf_plus._text_quality_analyzer import ( @@ -71,24 +70,28 @@ def extract_text_blocks( if not text.strip(): continue text_found = True - elements.append(PageElement( - element_type=ElementType.TEXT, - content=text, - bbox=bb, - page_num=page_num, - )) + elements.append( + PageElement( + element_type=ElementType.TEXT, + content=text, + bbox=bb, + page_num=page_num, + ) + ) # 2. Fallback: quality-aware extraction if structured gave nothing if not text_found: qa = QualityAwareTextExtractor(page, page_num) result = qa.extract() if result.text.strip(): - elements.append(PageElement( - element_type=ElementType.TEXT, - content=result.text, - bbox=(0, 0, page.rect.width, page.rect.height), - page_num=page_num, - )) + elements.append( + PageElement( + element_type=ElementType.TEXT, + content=result.text, + bbox=(0, 0, page.rect.width, page.rect.height), + page_num=page_num, + ) + ) if result.used_ocr: logger.debug( "[TextExtractor] page %d: used OCR fallback (quality=%.2f)", diff --git a/contextifier/handlers/pdf_plus/_text_quality_analyzer.py b/contextifier/handlers/pdf_plus/_text_quality_analyzer.py index 93fb4ca..5b56b9f 100644 --- a/contextifier/handlers/pdf_plus/_text_quality_analyzer.py +++ b/contextifier/handlers/pdf_plus/_text_quality_analyzer.py @@ -20,7 +20,7 @@ from __future__ import annotations import logging -from typing import Any, List, Optional, Tuple +from typing import Any, Optional from contextifier.handlers.pdf_plus._types import ( PdfPlusConfig, @@ -37,6 +37,7 @@ # Text Quality Analyzer # ====================================================================== + class TextQualityAnalyzer: """Measure extractable-text quality on a single PDF page.""" @@ -95,6 +96,7 @@ def _count_pua(self, text: str) -> int: # Page OCR Fallback # ====================================================================== + class PageOCRFallbackEngine: """Render page → image → pytesseract OCR.""" @@ -120,7 +122,9 @@ def ocr(self) -> str: logger.debug("[OCRFallback] pytesseract / Pillow not installed") return "" except Exception as exc: - logger.warning("[OCRFallback] page %d OCR failed: %s", self.page_num + 1, exc) + logger.warning( + "[OCRFallback] page %d OCR failed: %s", self.page_num + 1, exc + ) return "" @@ -128,6 +132,7 @@ def ocr(self) -> str: # Quality-Aware Text Extractor # ====================================================================== + class QualityAwareTextExtractor: """ Extract page text, automatically switching to OCR when quality is diff --git a/contextifier/handlers/pdf_plus/_types.py b/contextifier/handlers/pdf_plus/_types.py index daecf28..4cc0a01 100644 --- a/contextifier/handlers/pdf_plus/_types.py +++ b/contextifier/handlers/pdf_plus/_types.py @@ -29,21 +29,24 @@ class LineThickness(Enum): """Classification of a PDF drawing-line thickness.""" - THIN = auto() # < 0.5 pt — inner grid / hairline - NORMAL = auto() # 0.5–1.5 pt — regular border - THICK = auto() # > 1.5 pt — emphasis / header divider + + THIN = auto() # < 0.5 pt — inner grid / hairline + NORMAL = auto() # 0.5–1.5 pt — regular border + THICK = auto() # > 1.5 pt — emphasis / header divider class TableDetectionStrategy(Enum): """Which strategy found this table candidate.""" - PYMUPDF_NATIVE = auto() # fitz find_tables() - PDFPLUMBER_LINES = auto() # pdfplumber line-based - HYBRID_ANALYSIS = auto() # custom line-analysis grid + + PYMUPDF_NATIVE = auto() # fitz find_tables() + PDFPLUMBER_LINES = auto() # pdfplumber line-based + HYBRID_ANALYSIS = auto() # custom line-analysis grid BORDERLESS_HEURISTIC = auto() # text-clustering heuristic class ElementType(Enum): """Type of page element for the position-aware merger.""" + TEXT = "text" TABLE = "table" IMAGE = "image" @@ -52,31 +55,35 @@ class ElementType(Enum): class ComplexityLevel(Enum): """Page-level complexity classification.""" - SIMPLE = auto() # score < 0.35 → text extraction - MODERATE = auto() # 0.35–0.65 → hybrid - COMPLEX = auto() # 0.65–0.90 → hybrid (not full OCR) - EXTREME = auto() # ≥ 0.90 → full-page OCR (with caveats) + + SIMPLE = auto() # score < 0.35 → text extraction + MODERATE = auto() # 0.35–0.65 → hybrid + COMPLEX = auto() # 0.65–0.90 → hybrid (not full OCR) + EXTREME = auto() # ≥ 0.90 → full-page OCR (with caveats) class ProcessingStrategy(Enum): """Per-page processing strategy chosen by the complexity analyzer.""" - TEXT_EXTRACTION = auto() # standard text + table + image - HYBRID = auto() # text + block-image for complex regions - BLOCK_IMAGE_OCR = auto() # block-image dominant + remaining text - FULL_PAGE_OCR = auto() # smart blocks / grid / full-page image + + TEXT_EXTRACTION = auto() # standard text + table + image + HYBRID = auto() # text + block-image for complex regions + BLOCK_IMAGE_OCR = auto() # block-image dominant + remaining text + FULL_PAGE_OCR = auto() # smart blocks / grid / full-page image class TableQuality(Enum): """Processability grade for a table region.""" - EXCELLENT = auto() # ≥ 0.95 → must text-extract - GOOD = auto() # ≥ 0.85 → recommended text - MODERATE = auto() # ≥ 0.65 → attempt text - POOR = auto() # ≥ 0.40 → image recommended + + EXCELLENT = auto() # ≥ 0.95 → must text-extract + GOOD = auto() # ≥ 0.85 → recommended text + MODERATE = auto() # ≥ 0.65 → attempt text + POOR = auto() # ≥ 0.40 → image recommended UNPROCESSABLE = auto() # < 0.40 → must image class LayoutBlockType(Enum): """Semantic block type detected by the layout detector.""" + ARTICLE = auto() IMAGE_WITH_CAPTION = auto() STANDALONE_IMAGE = auto() @@ -91,9 +98,10 @@ class LayoutBlockType(Enum): class BlockProcessingStrategy(Enum): """Strategy used by the block image engine.""" - SEMANTIC_BLOCKS = auto() # layout detector → per-block images - GRID_BLOCKS = auto() # NxM grid division - FULL_PAGE = auto() # single full-page image + + SEMANTIC_BLOCKS = auto() # layout detector → per-block images + GRID_BLOCKS = auto() # NxM grid division + FULL_PAGE = auto() # single full-page image # ═════════════════════════════════════════════════════════════════════════════ @@ -111,7 +119,7 @@ class PdfPlusConfig: # ── Line analysis ──────────────────────────────────────────────────── THIN_LINE_MAX: float = 0.5 NORMAL_LINE_MAX: float = 1.5 - DOUBLE_LINE_TOLERANCE: float = 3.0 # merge lines within 5 pt + DOUBLE_LINE_TOLERANCE: float = 3.0 # merge lines within 5 pt DOUBLE_LINE_OVERLAP_RATIO: float = 0.5 # positional overlap threshold # ── Table detection ────────────────────────────────────────────────── @@ -157,7 +165,15 @@ class PdfPlusConfig: MERGE_X_OVERLAP: float = 0.80 ANNOTATION_Y_MARGIN: float = 30.0 ANNOTATION_PATTERNS: tuple = ( - "주)", "주 )", "※", "*", "†", "‡", "¹", "²", "³", + "주)", + "주 )", + "※", + "*", + "†", + "‡", + "¹", + "²", + "³", ) CONTINUITY_TOP_RATIO: float = 0.30 CONTINUITY_BOTTOM_RATIO: float = 0.70 @@ -245,7 +261,7 @@ class PdfPlusConfig: # ── OCR common ─────────────────────────────────────────────────────── OCR_LANGUAGE: str = "kor+eng" - BLOCK_IMAGE_DPI: int = 300 # alias for BLOCK_DPI + BLOCK_IMAGE_DPI: int = 300 # alias for BLOCK_DPI # ═════════════════════════════════════════════════════════════════════ # Aliases — used by modules that reference slightly different names. @@ -253,71 +269,71 @@ class PdfPlusConfig: # ═════════════════════════════════════════════════════════════════════ # _line_analysis.py - LINE_MERGE_TOLERANCE: float = 5.0 # ≡ TABLE_MERGE_TOLERANCE - BORDER_EXTENSION_MARGIN: float = 2.0 # page-border extension buffer - THIN_LINE_THRESHOLD: float = 0.5 # ≡ THIN_LINE_MAX - THICK_LINE_THRESHOLD: float = 1.5 # ≡ NORMAL_LINE_MAX - DOUBLE_LINE_GAP: float = 3.0 # ≡ DOUBLE_LINE_TOLERANCE + LINE_MERGE_TOLERANCE: float = 5.0 # ≡ TABLE_MERGE_TOLERANCE + BORDER_EXTENSION_MARGIN: float = 2.0 # page-border extension buffer + THIN_LINE_THRESHOLD: float = 0.5 # ≡ THIN_LINE_MAX + THICK_LINE_THRESHOLD: float = 1.5 # ≡ NORMAL_LINE_MAX + DOUBLE_LINE_GAP: float = 3.0 # ≡ DOUBLE_LINE_TOLERANCE # _graphic_detector.py - GRAPHIC_CURVE_RATIO_THRESHOLD: float = 0.3 # ≡ GRAPHIC_CURVE_RATIO - GRAPHIC_MIN_CURVE_COUNT: int = 10 # ≡ GRAPHIC_MIN_CURVES - GRAPHIC_FILL_RATIO_THRESHOLD: float = 0.2 # ≡ GRAPHIC_FILL_RATIO - GRAPHIC_COLOR_VARIETY_THRESHOLD: int = 3 # ≡ GRAPHIC_COLOR_VARIETY + GRAPHIC_CURVE_RATIO_THRESHOLD: float = 0.3 # ≡ GRAPHIC_CURVE_RATIO + GRAPHIC_MIN_CURVE_COUNT: int = 10 # ≡ GRAPHIC_MIN_CURVES + GRAPHIC_FILL_RATIO_THRESHOLD: float = 0.2 # ≡ GRAPHIC_FILL_RATIO + GRAPHIC_COLOR_VARIETY_THRESHOLD: int = 3 # ≡ GRAPHIC_COLOR_VARIETY # _table_detection.py - TABLE_CONFIDENCE_THRESHOLD: float = 0.35 # minimum to accept a table - PYMUPDF_SNAP_TOLERANCE: float = 7.0 # ≡ PYMUPDF_SNAP - PYMUPDF_JOIN_TOLERANCE: float = 7.0 # ≡ PYMUPDF_JOIN - PYMUPDF_EDGE_MIN_LENGTH: float = 10.0 # ≡ PYMUPDF_EDGE_MIN - PYMUPDF_INTERSECTION_TOLERANCE: float = 7.0 # ≡ PYMUPDF_INTERSECTION - PYMUPDF_CONFIDENCE_BASE: float = 0.6 # ≡ PYMUPDF_BASE_CONFIDENCE - PDFPLUMBER_CONFIDENCE_BASE: float = 0.4 # ≡ PDFPLUMBER_BASE_CONFIDENCE - LINE_CONFIDENCE_BASE: float = 0.3 # ≡ LINE_ANALYSIS_BASE_CONFIDENCE + TABLE_CONFIDENCE_THRESHOLD: float = 0.35 # minimum to accept a table + PYMUPDF_SNAP_TOLERANCE: float = 7.0 # ≡ PYMUPDF_SNAP + PYMUPDF_JOIN_TOLERANCE: float = 7.0 # ≡ PYMUPDF_JOIN + PYMUPDF_EDGE_MIN_LENGTH: float = 10.0 # ≡ PYMUPDF_EDGE_MIN + PYMUPDF_INTERSECTION_TOLERANCE: float = 7.0 # ≡ PYMUPDF_INTERSECTION + PYMUPDF_CONFIDENCE_BASE: float = 0.6 # ≡ PYMUPDF_BASE_CONFIDENCE + PDFPLUMBER_CONFIDENCE_BASE: float = 0.4 # ≡ PDFPLUMBER_BASE_CONFIDENCE + LINE_CONFIDENCE_BASE: float = 0.3 # ≡ LINE_ANALYSIS_BASE_CONFIDENCE # _table_validator.py - TABLE_MIN_FILLED_CELL_RATIO: float = 0.15 # ≡ MIN_FILLED_CELL_RATIO - TABLE_MAX_EMPTY_ROW_RATIO: float = 0.7 # ≡ MAX_EMPTY_ROW_RATIO - TABLE_MIN_MEANINGFUL_CELLS: int = 2 # ≡ MIN_MEANINGFUL_CELLS - TABLE_MIN_VALID_ROWS: int = 2 # ≡ MIN_VALID_ROWS - TABLE_MIN_TEXT_DENSITY: float = 0.005 # ≡ MIN_TEXT_DENSITY - TABLE_MAX_LONG_CELLS_RATIO: float = 0.4 # ≡ MAX_LONG_CELLS_RATIO - TABLE_EXTREME_CELL_LENGTH: int = 800 # ≡ EXTREME_CELL_LENGTH - TABLE_MAX_CELL_TEXT_LENGTH: int = 300 # ≡ MAX_CELL_TEXT_LENGTH + TABLE_MIN_FILLED_CELL_RATIO: float = 0.15 # ≡ MIN_FILLED_CELL_RATIO + TABLE_MAX_EMPTY_ROW_RATIO: float = 0.7 # ≡ MAX_EMPTY_ROW_RATIO + TABLE_MIN_MEANINGFUL_CELLS: int = 2 # ≡ MIN_MEANINGFUL_CELLS + TABLE_MIN_VALID_ROWS: int = 2 # ≡ MIN_VALID_ROWS + TABLE_MIN_TEXT_DENSITY: float = 0.005 # ≡ MIN_TEXT_DENSITY + TABLE_MAX_LONG_CELLS_RATIO: float = 0.4 # ≡ MAX_LONG_CELLS_RATIO + TABLE_EXTREME_CELL_LENGTH: int = 800 # ≡ EXTREME_CELL_LENGTH + TABLE_MAX_CELL_TEXT_LENGTH: int = 300 # ≡ MAX_CELL_TEXT_LENGTH # _table_quality_analyzer.py - QUALITY_WEIGHT_SIMPLE: float = 0.20 # ≡ QUALITY_WEIGHT_ELEMENT - QUALITY_BORDER_TOLERANCE: float = 3.0 # snap tolerance for border scoring - QUALITY_MIN_ORTHOGONAL_RATIO: float = 0.7 # ratio of orthogonal lines - QUALITY_MIN_CELL_SIZE: float = 8.0 # minimum cell dimension - QUALITY_MAX_CELL_ASPECT_RATIO: float = 20.0 # reject excessively narrow cells - QUALITY_LINE_ANGLE_TOLERANCE: float = 5.0 # degrees from hor/vert - QUALITY_MAX_CURVE_RATIO: float = 0.3 # max curve-to-total-line ratio - QUALITY_MAX_DIAGONAL_RATIO: float = 0.2 # max diagonal-to-total ratio - QUALITY_EXCELLENT_THRESHOLD: float = 0.95 # ≡ QUALITY_EXCELLENT - QUALITY_GOOD_THRESHOLD: float = 0.85 # ≡ QUALITY_GOOD - QUALITY_MODERATE_THRESHOLD: float = 0.65 # ≡ QUALITY_MODERATE - QUALITY_POOR_THRESHOLD: float = 0.40 # ≡ QUALITY_POOR + QUALITY_WEIGHT_SIMPLE: float = 0.20 # ≡ QUALITY_WEIGHT_ELEMENT + QUALITY_BORDER_TOLERANCE: float = 3.0 # snap tolerance for border scoring + QUALITY_MIN_ORTHOGONAL_RATIO: float = 0.7 # ratio of orthogonal lines + QUALITY_MIN_CELL_SIZE: float = 8.0 # minimum cell dimension + QUALITY_MAX_CELL_ASPECT_RATIO: float = 20.0 # reject excessively narrow cells + QUALITY_LINE_ANGLE_TOLERANCE: float = 5.0 # degrees from hor/vert + QUALITY_MAX_CURVE_RATIO: float = 0.3 # max curve-to-total-line ratio + QUALITY_MAX_DIAGONAL_RATIO: float = 0.2 # max diagonal-to-total ratio + QUALITY_EXCELLENT_THRESHOLD: float = 0.95 # ≡ QUALITY_EXCELLENT + QUALITY_GOOD_THRESHOLD: float = 0.85 # ≡ QUALITY_GOOD + QUALITY_MODERATE_THRESHOLD: float = 0.65 # ≡ QUALITY_MODERATE + QUALITY_POOR_THRESHOLD: float = 0.40 # ≡ QUALITY_POOR # _cell_analysis.py - CELL_GRID_TOLERANCE: float = 3.0 # snap cells to grid lines - CELL_OVERLAP_THRESHOLD: float = 0.5 # overlap ratio for cell merging + CELL_GRID_TOLERANCE: float = 3.0 # snap cells to grid lines + CELL_OVERLAP_THRESHOLD: float = 0.5 # overlap ratio for cell merging # _table_processor.py - TABLE_CONTINUITY_BOTTOM_RATIO: float = 0.70 # ≡ CONTINUITY_BOTTOM_RATIO - TABLE_CONTINUITY_TOP_RATIO: float = 0.30 # ≡ CONTINUITY_TOP_RATIO - TABLE_ANNOTATION_GAP: float = 30.0 # ≡ ANNOTATION_Y_MARGIN + TABLE_CONTINUITY_BOTTOM_RATIO: float = 0.70 # ≡ CONTINUITY_BOTTOM_RATIO + TABLE_CONTINUITY_TOP_RATIO: float = 0.30 # ≡ CONTINUITY_TOP_RATIO + TABLE_ANNOTATION_GAP: float = 30.0 # ≡ ANNOTATION_Y_MARGIN # _text_quality_analyzer.py - PUA_RANGES: list = [ # for tuple-based iteration + PUA_RANGES: list = [ # for tuple-based iteration (0xE000, 0xF8FF), (0xF0000, 0xFFFFD), ] - OCR_QUALITY_THRESHOLD: float = 0.7 # ≡ QUALITY_OCR_THRESHOLD + OCR_QUALITY_THRESHOLD: float = 0.7 # ≡ QUALITY_OCR_THRESHOLD # _vector_text_ocr.py VECTOR_TEXT_MAX_GLYPH_AREA: float = 2500.0 # max area for a single glyph path - VECTOR_TEXT_MIN_GLYPH_CLUSTER: int = 5 # ≡ VECTOR_MIN_ITEMS + VECTOR_TEXT_MIN_GLYPH_CLUSTER: int = 5 # ≡ VECTOR_MIN_ITEMS # ═════════════════════════════════════════════════════════════════════════════ @@ -328,6 +344,7 @@ class PdfPlusConfig: @dataclass class LineInfo: """A single detected line from page drawings.""" + x0: float y0: float x1: float @@ -349,8 +366,9 @@ def midpoint(self) -> Tuple[float, float]: @dataclass class GridInfo: """Grid constructed from horizontal / vertical lines.""" - h_lines: List[float] = field(default_factory=list) # Y positions - v_lines: List[float] = field(default_factory=list) # X positions + + h_lines: List[float] = field(default_factory=list) # Y positions + v_lines: List[float] = field(default_factory=list) # X positions cells: List["CellInfo"] = field(default_factory=list) bbox: Tuple[float, float, float, float] = (0, 0, 0, 0) is_complete: bool = False @@ -368,6 +386,7 @@ def col_count(self) -> int: @dataclass class CellInfo: """A single table cell with grid position and span info.""" + row: int col: int bbox: Tuple[float, float, float, float] @@ -381,6 +400,7 @@ class CellInfo: @dataclass class AnnotationInfo: """Annotation (footnote / endnote) associated with a table.""" + text: str = "" bbox: Tuple[float, float, float, float] = (0.0, 0.0, 0.0, 0.0) annotation_type: str = "annotation" @@ -390,6 +410,7 @@ class AnnotationInfo: @dataclass class VectorTextRegion: """Region containing outlined / vector-path text.""" + bbox: Tuple[float, float, float, float] drawing_count: int curve_count: int @@ -402,6 +423,7 @@ class VectorTextRegion: @dataclass class GraphicRegionInfo: """Detected graphic region (chart, diagram, icon).""" + bbox: Tuple[float, float, float, float] curve_count: int = 0 line_count: int = 0 @@ -416,6 +438,7 @@ class GraphicRegionInfo: @dataclass class TableCandidate: """A candidate table found by one of the detection strategies.""" + strategy: TableDetectionStrategy confidence: float bbox: Tuple[float, float, float, float] @@ -440,6 +463,7 @@ class PageElement: Sortable by ``(page_num, y0, x0)`` for reading-order assembly. """ + element_type: ElementType content: str bbox: Tuple[float, float, float, float] @@ -452,18 +476,24 @@ class PageElement: def __lt__(self, other: "PageElement") -> bool: return (self.page_num, self.bbox[1], self.bbox[0]) < ( - other.page_num, other.bbox[1], other.bbox[0] + other.page_num, + other.bbox[1], + other.bbox[0], ) @dataclass class PageBorderInfo: """Whether the page has a decorative border frame.""" + has_border: bool = False border_bbox: Optional[Tuple[float, float, float, float]] = None border_lines: Dict[str, bool] = field( default_factory=lambda: { - "top": False, "bottom": False, "left": False, "right": False, + "top": False, + "bottom": False, + "left": False, + "right": False, } ) @@ -474,6 +504,7 @@ class PageBorderInfo: @dataclass class RegionComplexity: """Complexity analysis for a single page region (grid cell).""" + bbox: Tuple[float, float, float, float] complexity_level: ComplexityLevel complexity_score: float @@ -488,6 +519,7 @@ class RegionComplexity: @dataclass class PageComplexity: """Full-page complexity analysis result.""" + page_num: int page_size: Tuple[float, float] overall_complexity: ComplexityLevel @@ -509,6 +541,7 @@ class PageComplexity: @dataclass class BlockResult: """Result of rendering one block as an image.""" + success: bool image_tag: Optional[str] = None bbox: Tuple[float, float, float, float] = (0, 0, 0, 0) @@ -519,6 +552,7 @@ class BlockResult: @dataclass class MultiBlockResult: """Result of rendering multiple blocks (smart processing).""" + success: bool strategy_used: BlockProcessingStrategy = BlockProcessingStrategy.FULL_PAGE block_results: List[BlockResult] = field(default_factory=list) @@ -532,6 +566,7 @@ class MultiBlockResult: @dataclass class TableQualityResult: """Assessment of a table region's processability.""" + quality: TableQuality score: float border_score: float = 0.0 @@ -547,6 +582,7 @@ class TableQualityResult: @dataclass class LayoutBlock: """Semantic layout block detected by the layout detector.""" + block_type: LayoutBlockType bbox: Tuple[float, float, float, float] elements: List[Dict] = field(default_factory=list) @@ -557,6 +593,7 @@ class LayoutBlock: @dataclass class LayoutAnalysisResult: """Full layout analysis for a page.""" + blocks: List[LayoutBlock] = field(default_factory=list) columns: List[Tuple[float, float]] = field(default_factory=list) has_header: bool = False @@ -569,7 +606,8 @@ class LayoutAnalysisResult: @dataclass class TextQualityResult: """Quality analysis for a text fragment.""" - quality_score: float # 0.0 – 1.0 + + quality_score: float # 0.0 – 1.0 total_chars: int = 0 pua_chars: int = 0 garbled_ratio: float = 0.0 @@ -580,6 +618,7 @@ class TextQualityResult: @dataclass class PageTextAnalysis: """Quality-aware text extraction result for a single page.""" + text: str = "" quality: Optional[TextQualityResult] = None used_ocr: bool = False diff --git a/contextifier/handlers/pdf_plus/_vector_text_ocr.py b/contextifier/handlers/pdf_plus/_vector_text_ocr.py index 55d6496..e6a93e7 100644 --- a/contextifier/handlers/pdf_plus/_vector_text_ocr.py +++ b/contextifier/handlers/pdf_plus/_vector_text_ocr.py @@ -11,7 +11,7 @@ from __future__ import annotations import logging -from typing import Any, Dict, List, Optional, Tuple +from typing import Any, List, Tuple from contextifier.handlers.pdf_plus._types import ( PdfPlusConfig, @@ -100,10 +100,12 @@ def _detect_regions(self) -> List[VectorTextRegion]: if text.strip(): continue # already has real text - regions.append(VectorTextRegion( - bbox=region_bbox, - glyph_count=len(cluster), - )) + regions.append( + VectorTextRegion( + bbox=region_bbox, + glyph_count=len(cluster), + ) + ) return regions @@ -135,10 +137,7 @@ def _cluster_bboxes( @staticmethod def _nearby(a: tuple, b: tuple, m: float) -> bool: return not ( - a[2] + m < b[0] - or b[2] + m < a[0] - or a[3] + m < b[1] - or b[3] + m < a[1] + a[2] + m < b[0] or b[2] + m < a[0] or a[3] + m < b[1] or b[3] + m < a[1] ) # ------------------------------------------------------------------ diff --git a/contextifier/handlers/pdf_plus/content_extractor.py b/contextifier/handlers/pdf_plus/content_extractor.py index d096099..1d3027e 100644 --- a/contextifier/handlers/pdf_plus/content_extractor.py +++ b/contextifier/handlers/pdf_plus/content_extractor.py @@ -121,7 +121,11 @@ def extract_text( for page_num in range(page_count): page = doc.load_page(page_num) page_text = self._process_page( - doc, page, page_num, processed_images, file_data, + doc, + page, + page_num, + processed_images, + file_data, ) if page_text.strip(): tag = self._make_page_tag(page_num + 1) @@ -189,30 +193,49 @@ def _process_page( # 2. Dispatch to strategy if strategy == ProcessingStrategy.TEXT_EXTRACTION: return self._strategy_text( - doc, page, page_num, processed_images, file_data, + doc, + page, + page_num, + processed_images, + file_data, ) elif strategy == ProcessingStrategy.HYBRID: return self._strategy_hybrid( - doc, page, page_num, processed_images, file_data, + doc, + page, + page_num, + processed_images, + file_data, complexity, ) elif strategy == ProcessingStrategy.BLOCK_IMAGE_OCR: return self._strategy_block_ocr( - doc, page, page_num, processed_images, file_data, + doc, + page, + page_num, + processed_images, + file_data, complexity, ) else: # FULL_PAGE_OCR return self._strategy_full_ocr( - page, page_num, processed_images, + page, + page_num, + processed_images, ) except Exception as exc: logger.error( "[PdfPlus] Error on page %d, falling back to text: %s", - page_num + 1, exc, + page_num + 1, + exc, ) return self._strategy_text( - doc, page, page_num, processed_images, file_data, + doc, + page, + page_num, + processed_images, + file_data, ) # ───────────────────────────────────────────────────────────────────── @@ -234,33 +257,36 @@ def _strategy_text( tp = TableProcessor(page, page_num, file_data) for pt in tp.process(): if pt.html: - elements.append(PageElement( - element_type=ElementType.TABLE, - content=pt.html, - bbox=pt.bbox, - page_num=page_num, - table_data=pt.data, - confidence=pt.confidence, - )) + elements.append( + PageElement( + element_type=ElementType.TABLE, + content=pt.html, + bbox=pt.bbox, + page_num=page_num, + table_data=pt.data, + confidence=pt.confidence, + ) + ) table_bboxes.append(pt.bbox) if pt.annotations: for ann in pt.annotations: - elements.append(PageElement( - element_type=ElementType.ANNOTATION, - content=ann.text, - bbox=ann.bbox, - page_num=page_num, - )) + elements.append( + PageElement( + element_type=ElementType.ANNOTATION, + content=ann.text, + bbox=ann.bbox, + page_num=page_num, + ) + ) # text blocks - elements.extend( - extract_text_blocks(page, page_num, table_bboxes) - ) + elements.extend(extract_text_blocks(page, page_num, table_bboxes)) # images elements.extend( extract_images( - page, page_num, + page, + page_num, table_bboxes=table_bboxes, image_service=self._image_service, ) @@ -288,33 +314,36 @@ def _strategy_hybrid( tp = TableProcessor(page, page_num, file_data) for pt in tp.process(): if pt.html: - elements.append(PageElement( - element_type=ElementType.TABLE, - content=pt.html, - bbox=pt.bbox, - page_num=page_num, - table_data=pt.data, - confidence=pt.confidence, - )) + elements.append( + PageElement( + element_type=ElementType.TABLE, + content=pt.html, + bbox=pt.bbox, + page_num=page_num, + table_data=pt.data, + confidence=pt.confidence, + ) + ) table_bboxes.append(pt.bbox) if pt.annotations: for ann in pt.annotations: - elements.append(PageElement( - element_type=ElementType.ANNOTATION, - content=ann.text, - bbox=ann.bbox, - page_num=page_num, - )) + elements.append( + PageElement( + element_type=ElementType.ANNOTATION, + content=ann.text, + bbox=ann.bbox, + page_num=page_num, + ) + ) # text blocks - elements.extend( - extract_text_blocks(page, page_num, table_bboxes) - ) + elements.extend(extract_text_blocks(page, page_num, table_bboxes)) # images elements.extend( extract_images( - page, page_num, + page, + page_num, table_bboxes=table_bboxes, image_service=self._image_service, ) @@ -325,17 +354,20 @@ def _strategy_hybrid( engine = VectorTextOCREngine(page, page_num) for region in engine.detect_and_ocr(): if region.ocr_text.strip(): - elements.append(PageElement( - element_type=ElementType.TEXT, - content=region.ocr_text, - bbox=region.bbox, - page_num=page_num, - confidence=region.confidence, - )) + elements.append( + PageElement( + element_type=ElementType.TEXT, + content=region.ocr_text, + bbox=region.bbox, + page_num=page_num, + confidence=region.confidence, + ) + ) except Exception as exc: logger.debug( "[PdfPlus] Vector OCR skipped on page %d: %s", - page_num + 1, exc, + page_num + 1, + exc, ) # block-image complex regions (if significant complex area) @@ -343,20 +375,25 @@ def _strategy_hybrid( complex_regions = complexity.complex_regions if complex_regions and self._image_service is not None: engine = BlockImageEngine( - page, page_num, image_service=self._image_service, + page, + page_num, + image_service=self._image_service, ) for bbox in complex_regions: br = engine.process_region(bbox, region_type="complex") if br.success and br.image_tag: - elements.append(PageElement( - element_type=ElementType.IMAGE, - content=br.image_tag, - bbox=br.bbox, - page_num=page_num, - )) + elements.append( + PageElement( + element_type=ElementType.IMAGE, + content=br.image_tag, + bbox=br.bbox, + page_num=page_num, + ) + ) except Exception as exc: logger.debug( - "[PdfPlus] Block-image for complex regions skipped: %s", exc, + "[PdfPlus] Block-image for complex regions skipped: %s", + exc, ) return merge_page_elements(elements) @@ -389,31 +426,37 @@ def _strategy_block_ocr( tp = TableProcessor(page, page_num, file_data) for pt in tp.process(): if pt.html: - elements.append(PageElement( - element_type=ElementType.TABLE, - content=pt.html, - bbox=pt.bbox, - page_num=page_num, - table_data=pt.data, - )) + elements.append( + PageElement( + element_type=ElementType.TABLE, + content=pt.html, + bbox=pt.bbox, + page_num=page_num, + table_data=pt.data, + ) + ) table_bboxes.append(pt.bbox) break # processor handles all tables at once # remaining content → block-image if self._image_service is not None: engine = BlockImageEngine( - page, page_num, image_service=self._image_service, + page, + page_num, + image_service=self._image_service, ) result = engine.process_page_smart() if result.success: combined = combine_block_output(result) if combined.strip(): - elements.append(PageElement( - element_type=ElementType.IMAGE, - content=combined, - bbox=(0, 0, page.rect.width, page.rect.height), - page_num=page_num, - )) + elements.append( + PageElement( + element_type=ElementType.IMAGE, + content=combined, + bbox=(0, 0, page.rect.width, page.rect.height), + page_num=page_num, + ) + ) return merge_page_elements(elements) @@ -433,7 +476,9 @@ def _strategy_full_ocr( return text engine = BlockImageEngine( - page, page_num, image_service=self._image_service, + page, + page_num, + image_service=self._image_service, ) result = engine.process_page_smart() if result.success: diff --git a/contextifier/handlers/ppt/_constants.py b/contextifier/handlers/ppt/_constants.py index bb3e745..ecb27f3 100644 --- a/contextifier/handlers/ppt/_constants.py +++ b/contextifier/handlers/ppt/_constants.py @@ -30,12 +30,14 @@ DOC_SUMMARY_STREAM = "\x05DocumentSummaryInformation" # Known OLE streams in PPT files -PPT_KNOWN_STREAMS = frozenset({ - PP_DOCUMENT_STREAM, - CURRENT_USER_STREAM, - SUMMARY_INFO_STREAM, - DOC_SUMMARY_STREAM, -}) +PPT_KNOWN_STREAMS = frozenset( + { + PP_DOCUMENT_STREAM, + CURRENT_USER_STREAM, + SUMMARY_INFO_STREAM, + DOC_SUMMARY_STREAM, + } +) __all__ = [ diff --git a/contextifier/handlers/ppt/content_extractor.py b/contextifier/handlers/ppt/content_extractor.py index 84d56d2..eb6ba4a 100644 --- a/contextifier/handlers/ppt/content_extractor.py +++ b/contextifier/handlers/ppt/content_extractor.py @@ -29,7 +29,7 @@ import logging import re -from typing import Any, Dict, List, Optional, Set, Tuple +from typing import Any, List, Optional, Set, Tuple from contextifier.pipeline.content_extractor import BaseContentExtractor from contextifier.types import ( @@ -42,11 +42,11 @@ logger = logging.getLogger(__name__) # PPT binary record types containing text -_TEXT_BYTES_ATOM = 0x0FA0 # TextBytesAtom: single-byte text (ANSI) -_TEXT_CHARS_ATOM = 0x0FA8 # TextCharsAtom: double-byte text (Unicode UTF-16LE) -_CSTRING_ATOM = 0x0FBA # CString: Unicode string (used in headers etc.) -_SLIDE_ATOM = 0x03EF # SlideAtom -_NOTES_CONTAINER = 0x0408 # NotesContainer +_TEXT_BYTES_ATOM = 0x0FA0 # TextBytesAtom: single-byte text (ANSI) +_TEXT_CHARS_ATOM = 0x0FA8 # TextCharsAtom: double-byte text (Unicode UTF-16LE) +_CSTRING_ATOM = 0x0FBA # CString: Unicode string (used in headers etc.) +_SLIDE_ATOM = 0x03EF # SlideAtom +_NOTES_CONTAINER = 0x0408 # NotesContainer class PptContentExtractor(BaseContentExtractor): @@ -124,6 +124,7 @@ def extract_images( continue import hashlib + content_hash = hashlib.md5(img_data).hexdigest()[:16] if content_hash in processed: continue @@ -207,6 +208,7 @@ def _make_slide_tag(self, slide_number: int) -> Optional[str]: # Binary stream parsing # ═══════════════════════════════════════════════════════════════════════════════ + def _parse_text_records(data: bytes) -> List[Tuple[int, str]]: """ Parse text records from the PowerPoint Document stream. @@ -330,6 +332,7 @@ def _group_into_slides(records: List[Tuple[int, str]]) -> List[List[str]]: # Table detection from text patterns # ═══════════════════════════════════════════════════════════════════════════════ + def _detect_tabular_text( records: List[Tuple[int, str]], ) -> List[TableData]: @@ -369,11 +372,13 @@ def _rows_to_table_data(rows: List[List[str]]) -> TableData: for row_idx, cells in enumerate(rows): row: List[TableCell] = [] for col_idx, text in enumerate(cells): - row.append(TableCell( - content=text, - row_index=row_idx, - col_index=col_idx, - )) + row.append( + TableCell( + content=text, + row_index=row_idx, + col_index=col_idx, + ) + ) table_rows.append(row) return TableData( @@ -422,9 +427,9 @@ def _detect_ole_charts(ole: Any) -> List[ChartData]: parent = entry[0] # Try to read the OLE class ID of the parent storage try: - clsid = ole.get_rootentry_name() if len(entry) == 1 else None + ole.get_rootentry_name() if len(entry) == 1 else None except Exception: - clsid = None + pass # We can't reliably determine chart type without full parsing, # but we note the embedded object if parent not in chart_storages: @@ -433,10 +438,12 @@ def _detect_ole_charts(ole: Any) -> List[ChartData]: chart_storages.add(parent) for storage_name in chart_storages: - charts.append(ChartData( - chart_type="embedded_ole", - title=f"Embedded Chart ({storage_name})", - raw_content=f"OLE embedded chart object in storage: {storage_name}", - )) + charts.append( + ChartData( + chart_type="embedded_ole", + title=f"Embedded Chart ({storage_name})", + raw_content=f"OLE embedded chart object in storage: {storage_name}", + ) + ) return charts diff --git a/contextifier/handlers/ppt/converter.py b/contextifier/handlers/ppt/converter.py index b6ef754..56a7319 100644 --- a/contextifier/handlers/ppt/converter.py +++ b/contextifier/handlers/ppt/converter.py @@ -29,8 +29,9 @@ class PptConvertedData(NamedTuple): """Result of the PPT conversion stage.""" - ole: olefile.OleFileIO # Opened OLE2 compound file - file_extension: str # Original extension (always "ppt") + + ole: olefile.OleFileIO # Opened OLE2 compound file + file_extension: str # Original extension (always "ppt") class PptConverter(BaseConverter): diff --git a/contextifier/handlers/ppt/handler.py b/contextifier/handlers/ppt/handler.py index dee199e..da0a9dd 100644 --- a/contextifier/handlers/ppt/handler.py +++ b/contextifier/handlers/ppt/handler.py @@ -38,7 +38,7 @@ from contextifier.pipeline.content_extractor import BaseContentExtractor from contextifier.pipeline.postprocessor import BasePostprocessor, DefaultPostprocessor -from contextifier.handlers.ppt._constants import ZIP_MAGIC, OLE2_MAGIC +from contextifier.handlers.ppt._constants import ZIP_MAGIC from contextifier.handlers.ppt.converter import PptConverter from contextifier.handlers.ppt.preprocessor import PptPreprocessor from contextifier.handlers.ppt.metadata_extractor import PptMetadataExtractor @@ -84,7 +84,8 @@ def _check_delegation( if data[:4] == ZIP_MAGIC: self._logger.info("PPT file is actually PPTX (ZIP magic detected)") return self._delegate_to( - "pptx", file_context, + "pptx", + file_context, include_metadata=kwargs.get("include_metadata", True), **{k: v for k, v in kwargs.items() if k != "include_metadata"}, ) diff --git a/contextifier/handlers/ppt/metadata_extractor.py b/contextifier/handlers/ppt/metadata_extractor.py index 59791ca..50dfb28 100644 --- a/contextifier/handlers/ppt/metadata_extractor.py +++ b/contextifier/handlers/ppt/metadata_extractor.py @@ -53,7 +53,9 @@ def extract(self, source: Any) -> DocumentMetadata: create_time=self._to_datetime(meta.create_time), last_saved_time=self._to_datetime(meta.last_saved_time), revision=str(meta.revision_number) if meta.revision_number else None, - category=self._decode(meta.category) if hasattr(meta, "category") else None, + category=self._decode(meta.category) + if hasattr(meta, "category") + else None, ) except Exception as exc: self._logger.warning("Failed to extract PPT metadata: %s", exc) @@ -92,7 +94,9 @@ def _is_ole(obj: Any) -> bool: if content is not None and hasattr(content, "get_metadata"): return content # Direct OLE-like object with get_metadata (e.g. from preprocessed.content) - if hasattr(source, "get_metadata") and callable(getattr(source, "get_metadata", None)): + if hasattr(source, "get_metadata") and callable( + getattr(source, "get_metadata", None) + ): return source # Has .ole attribute if hasattr(source, "ole"): diff --git a/contextifier/handlers/ppt/preprocessor.py b/contextifier/handlers/ppt/preprocessor.py index cd79b56..a72263c 100644 --- a/contextifier/handlers/ppt/preprocessor.py +++ b/contextifier/handlers/ppt/preprocessor.py @@ -9,7 +9,7 @@ from __future__ import annotations import logging -from typing import Any, Dict, List, Optional +from typing import Any, List, Optional from contextifier.pipeline.preprocessor import BasePreprocessor from contextifier.types import PreprocessedData diff --git a/contextifier/handlers/pptx/_bullet.py b/contextifier/handlers/pptx/_bullet.py index 3e6a1d3..a22a4f5 100644 --- a/contextifier/handlers/pptx/_bullet.py +++ b/contextifier/handlers/pptx/_bullet.py @@ -31,6 +31,7 @@ # Public API # ═══════════════════════════════════════════════════════════════════════════════ + def extract_text_with_bullets(text_frame: Any) -> str: """ Extract text from a python-pptx TextFrame, preserving bullet @@ -92,6 +93,7 @@ def extract_text_with_bullets(text_frame: Any) -> str: # Bullet detection # ═══════════════════════════════════════════════════════════════════════════════ + def _extract_bullet_info(paragraph: Any) -> Dict[str, Any]: """ Parse bullet/numbering XML from a paragraph element. @@ -160,6 +162,7 @@ def _extract_bullet_info(paragraph: Any) -> Dict[str, Any]: # Special font conversion # ═══════════════════════════════════════════════════════════════════════════════ + def _convert_special_font_char(char: str, font_typeface: str) -> str: """Convert a Wingdings/Symbol character to a Unicode equivalent.""" if not char: @@ -195,6 +198,7 @@ def _convert_special_font_char(char: str, font_typeface: str) -> str: # Numbering helpers # ═══════════════════════════════════════════════════════════════════════════════ + def _get_or_increment_number( numbering_state: Dict[int, int], level: int, @@ -244,9 +248,19 @@ def _format_number(num: int, format_type: str) -> str: def _to_roman(num: int) -> str: """Convert integer to Roman numeral string.""" val_map = [ - (1000, "M"), (900, "CM"), (500, "D"), (400, "CD"), - (100, "C"), (90, "XC"), (50, "L"), (40, "XL"), - (10, "X"), (9, "IX"), (5, "V"), (4, "IV"), (1, "I"), + (1000, "M"), + (900, "CM"), + (500, "D"), + (400, "CD"), + (100, "C"), + (90, "XC"), + (50, "L"), + (40, "XL"), + (10, "X"), + (9, "IX"), + (5, "V"), + (4, "IV"), + (1, "I"), ] parts: list[str] = [] for value, letter in val_map: diff --git a/contextifier/handlers/pptx/_constants.py b/contextifier/handlers/pptx/_constants.py index e35bd01..f9f552c 100644 --- a/contextifier/handlers/pptx/_constants.py +++ b/contextifier/handlers/pptx/_constants.py @@ -26,9 +26,11 @@ # Slide element types # ═══════════════════════════════════════════════════════════════════════════════ + @unique class ElementType(str, Enum): """Classification of content elements found on a slide.""" + TEXT = "text" IMAGE = "image" TABLE = "table" @@ -43,6 +45,7 @@ class SlideElement: Slide content is collected per-shape, then sorted by position (top first, then left) to reconstruct visual reading order. """ + element_type: ElementType content: str position: Tuple[int, int, int, int] # (left, top, width, height) EMU @@ -127,11 +130,11 @@ def sort_key(self) -> Tuple[int, int]: # Symbol font code-point → Unicode mapping SYMBOL_MAPPING: Dict[int, str] = { - 0xB7: "•", # Bullet - 0xD7: "×", # Multiplication - 0xF7: "÷", # Division - 0xA5: "∞", # Infinity - 0xB1: "±", # Plus-minus + 0xB7: "•", # Bullet + 0xD7: "×", # Multiplication + 0xF7: "÷", # Division + 0xA5: "∞", # Infinity + 0xB1: "±", # Plus-minus } diff --git a/contextifier/handlers/pptx/_table.py b/contextifier/handlers/pptx/_table.py index fd9ca57..902132c 100644 --- a/contextifier/handlers/pptx/_table.py +++ b/contextifier/handlers/pptx/_table.py @@ -13,7 +13,7 @@ from __future__ import annotations import logging -from typing import Any, Dict, List, Optional +from typing import Any, Dict, List from contextifier.types import TableCell, TableData @@ -24,6 +24,7 @@ # Public helpers # ═══════════════════════════════════════════════════════════════════════════════ + def is_simple_table(table: Any) -> bool: """ Decide whether a table is a "simple" layout table that should @@ -144,6 +145,7 @@ def extract_table(table: Any) -> TableData: # Merge detection # ═══════════════════════════════════════════════════════════════════════════════ + def _get_merge_info( cell: Any, table: Any, diff --git a/contextifier/handlers/pptx/content_extractor.py b/contextifier/handlers/pptx/content_extractor.py index 2178ace..6424893 100644 --- a/contextifier/handlers/pptx/content_extractor.py +++ b/contextifier/handlers/pptx/content_extractor.py @@ -55,7 +55,9 @@ def __init__(self, **kwargs: Any) -> None: super().__init__(**kwargs) if self._config is not None: depth = self._config.get_format_option( - "pptx", "max_group_depth", self._MAX_GROUP_DEPTH, + "pptx", + "max_group_depth", + self._MAX_GROUP_DEPTH, ) self._MAX_GROUP_DEPTH = int(depth) @@ -201,25 +203,33 @@ def _process_shape( if shape.has_table: content = self._format_table_shape(shape.table) if content: - etype = ElementType.TEXT if is_simple_table(shape.table) else ElementType.TABLE - elements.append(SlideElement( - element_type=etype, - content=content, - position=position, - shape_id=shape_id, - )) + etype = ( + ElementType.TEXT + if is_simple_table(shape.table) + else ElementType.TABLE + ) + elements.append( + SlideElement( + element_type=etype, + content=content, + position=position, + shape_id=shape_id, + ) + ) return elements, chart_ptr # Picture (image) if _is_picture(shape): tag = self._save_image_shape(shape, slide_idx, processed_images) if tag: - elements.append(SlideElement( - element_type=ElementType.IMAGE, - content=tag, - position=position, - shape_id=shape_id, - )) + elements.append( + SlideElement( + element_type=ElementType.IMAGE, + content=tag, + position=position, + shape_id=shape_id, + ) + ) return elements, chart_ptr # Chart @@ -227,30 +237,38 @@ def _process_shape( content = self._format_chart(chart_queue, chart_ptr) chart_ptr += 1 if content: - elements.append(SlideElement( - element_type=ElementType.CHART, - content=content, - position=position, - shape_id=shape_id, - )) + elements.append( + SlideElement( + element_type=ElementType.CHART, + content=content, + position=position, + shape_id=shape_id, + ) + ) return elements, chart_ptr # Text frame (with bullets) if hasattr(shape, "text_frame") and shape.text_frame: text = extract_text_with_bullets(shape.text_frame) if text.strip(): - elements.append(SlideElement( - element_type=ElementType.TEXT, - content=text, - position=position, - shape_id=shape_id, - )) + elements.append( + SlideElement( + element_type=ElementType.TEXT, + content=text, + position=position, + shape_id=shape_id, + ) + ) return elements, chart_ptr # Group shape (recursive) if hasattr(shape, "shapes"): group_elems, chart_ptr = self._process_group( - shape, slide_idx, chart_queue, chart_ptr, processed_images, + shape, + slide_idx, + chart_queue, + chart_ptr, + processed_images, depth=depth, ) elements.extend(group_elems) @@ -258,12 +276,14 @@ def _process_shape( # Fallback: plain text if hasattr(shape, "text") and shape.text.strip(): - elements.append(SlideElement( - element_type=ElementType.TEXT, - content=shape.text.strip(), - position=position, - shape_id=shape_id, - )) + elements.append( + SlideElement( + element_type=ElementType.TEXT, + content=shape.text.strip(), + position=position, + shape_id=shape_id, + ) + ) return elements, chart_ptr @@ -281,7 +301,8 @@ def _process_group( if depth >= self._MAX_GROUP_DEPTH: logger.warning( "Group shape nesting depth %d exceeds limit %d, skipping", - depth, self._MAX_GROUP_DEPTH, + depth, + self._MAX_GROUP_DEPTH, ) return [], chart_ptr @@ -289,8 +310,12 @@ def _process_group( try: for sub_shape in group_shape.shapes: sub_elems, chart_ptr = self._process_shape( - sub_shape, slide_idx, chart_queue, chart_ptr, - processed_images, depth=depth + 1, + sub_shape, + slide_idx, + chart_queue, + chart_ptr, + processed_images, + depth=depth + 1, ) elements.extend(sub_elems) except Exception as exc: @@ -300,7 +325,10 @@ def _process_group( # ── Recursive collection helpers (for extract_tables/images/charts) ── def _collect_tables( - self, shape: Any, tables: List[TableData], depth: int = 0, + self, + shape: Any, + tables: List[TableData], + depth: int = 0, ) -> None: """Recursively collect tables from shape and nested groups.""" if hasattr(shape, "has_table") and shape.has_table: @@ -328,7 +356,10 @@ def _collect_images( self._collect_images(sub, slide_idx, processed, tags, depth + 1) def _collect_charts( - self, shape: Any, charts: List[ChartData], depth: int = 0, + self, + shape: Any, + charts: List[ChartData], + depth: int = 0, ) -> None: """Recursively collect charts from shape and nested groups.""" cd = self._try_extract_chart_data(shape) @@ -358,8 +389,7 @@ def _format_table_shape(self, table: Any) -> str: # Single-column collapse if table_data.num_cols == 1: items = [ - row[0].content for row in table_data.rows - if row and row[0].content + row[0].content for row in table_data.rows if row and row[0].content ] return "\n\n".join(items) if items else "" @@ -445,7 +475,11 @@ def _try_extract_chart_data(self, shape: Any) -> Optional[ChartData]: chart = shape.chart title = None try: - if chart.has_title and chart.chart_title and chart.chart_title.has_text_frame: + if ( + chart.has_title + and chart.chart_title + and chart.chart_title.has_text_frame + ): title = chart.chart_title.text_frame.text.strip() or None except (AttributeError, TypeError) as exc: logger.debug("Chart title extraction failed: %s", exc) @@ -524,6 +558,7 @@ def _get_presentation(preprocessed: PreprocessedData) -> Any: # Module-level helpers # ═══════════════════════════════════════════════════════════════════════════════ + def _get_position(shape: Any) -> Tuple[int, int, int, int]: """Get (left, top, width, height) in EMU from a shape.""" try: @@ -540,6 +575,7 @@ def _is_picture(shape: Any) -> bool: """Check whether a shape is an image/picture.""" try: from pptx.enum.shapes import MSO_SHAPE_TYPE + if hasattr(shape, "shape_type") and shape.shape_type == MSO_SHAPE_TYPE.PICTURE: return True except Exception: diff --git a/contextifier/handlers/pptx/metadata_extractor.py b/contextifier/handlers/pptx/metadata_extractor.py index 72f1e6a..f0d3891 100644 --- a/contextifier/handlers/pptx/metadata_extractor.py +++ b/contextifier/handlers/pptx/metadata_extractor.py @@ -48,8 +48,12 @@ def extract(self, source: Any) -> DocumentMetadata: create_time=props.created, last_saved_time=props.modified, page_count=len(prs.slides) if hasattr(prs, "slides") else None, - category=self._get(props.category) if hasattr(props, "category") else None, - revision=str(props.revision) if hasattr(props, "revision") and props.revision else None, + category=self._get(props.category) + if hasattr(props, "category") + else None, + revision=str(props.revision) + if hasattr(props, "revision") and props.revision + else None, ) except Exception as exc: self._logger.warning("Failed to extract PPTX metadata: %s", exc) diff --git a/contextifier/handlers/registry.py b/contextifier/handlers/registry.py index a94bd6d..f73d262 100644 --- a/contextifier/handlers/registry.py +++ b/contextifier/handlers/registry.py @@ -25,7 +25,6 @@ import logging from typing import ( - Callable, Dict, FrozenSet, List, @@ -100,10 +99,10 @@ def register(self, handler_class: Type[BaseHandler]) -> None: Raises: TypeError: If handler_class is not a BaseHandler subclass. """ - if not (isinstance(handler_class, type) and issubclass(handler_class, BaseHandler)): - raise TypeError( - f"Expected BaseHandler subclass, got {handler_class}" - ) + if not ( + isinstance(handler_class, type) and issubclass(handler_class, BaseHandler) + ): + raise TypeError(f"Expected BaseHandler subclass, got {handler_class}") try: handler = handler_class( @@ -115,9 +114,7 @@ def register(self, handler_class: Type[BaseHandler]) -> None: metadata_service=self._services.get("metadata_service"), ) except Exception as e: - logger.warning( - f"Failed to instantiate {handler_class.__name__}: {e}" - ) + logger.warning(f"Failed to instantiate {handler_class.__name__}: {e}") return # Inject registry reference so handler can use delegation @@ -176,13 +173,12 @@ def register_defaults(self) -> None: for module_path, class_name in default_handlers: try: import importlib + module = importlib.import_module(module_path) handler_class = getattr(module, class_name) self.register(handler_class) except (ImportError, AttributeError) as e: - logger.warning( - f"Handler {class_name} not available: {e}" - ) + logger.warning(f"Handler {class_name} not available: {e}") # Discover third-party handler plugins self._discover_plugins() @@ -265,12 +261,19 @@ def _discover_plugins(self) -> None: for ep in group: try: handler_class = ep.load() - if isinstance(handler_class, type) and issubclass(handler_class, BaseHandler): + if isinstance(handler_class, type) and issubclass( + handler_class, BaseHandler + ): self.register(handler_class) - logger.info("Loaded plugin handler: %s (%s)", ep.name, handler_class.__name__) + logger.info( + "Loaded plugin handler: %s (%s)", + ep.name, + handler_class.__name__, + ) else: logger.warning( - "Plugin '%s' does not provide a BaseHandler subclass", ep.name, + "Plugin '%s' does not provide a BaseHandler subclass", + ep.name, ) except Exception as exc: logger.warning("Failed to load plugin '%s': %s", ep.name, exc) diff --git a/contextifier/handlers/rtf/_cleaner.py b/contextifier/handlers/rtf/_cleaner.py index 1902fc6..2986c95 100644 --- a/contextifier/handlers/rtf/_cleaner.py +++ b/contextifier/handlers/rtf/_cleaner.py @@ -30,19 +30,17 @@ SKIP_DESTINATIONS, IMAGE_DESTINATIONS, ) -from contextifier.handlers.rtf._decoder import decode_bytes # Precompile shape property name pattern for performance -_SHAPE_NAME_PATTERN = re.compile( - r"\b(" + "|".join(SHAPE_PROPERTY_NAMES) + r")\b" -) +_SHAPE_NAME_PATTERN = re.compile(r"\b(" + "|".join(SHAPE_PROPERTY_NAMES) + r")\b") # ═══════════════════════════════════════════════════════════════════════════ # Main Text Cleaner # ═══════════════════════════════════════════════════════════════════════════ + def clean_rtf_text(text: str, encoding: str = "cp949") -> str: """ Remove RTF control codes and extract pure text. @@ -82,9 +80,7 @@ def save_image_tag(m: re.Match) -> str: # Remove shape property inline patterns text = re.sub(r"\{\\sp\{\\sn\s*\w+\}\{\\sv\s*[^}]*\}\}", "", text) - text = re.sub( - r"shapeType\d+[a-zA-Z0-9]+(?:posrelh\d+posrelv\d+)?", "", text - ) + text = re.sub(r"shapeType\d+[a-zA-Z0-9]+(?:posrelh\d+posrelv\d+)?", "", text) text = re.sub( r"\\shp(?:inst|txt|left|right|top|bottom|bx\w+|by\w+|wr\d+|fblwtxt\d+|z\d+|lid\d+)\b\d*", "", @@ -129,11 +125,11 @@ def save_image_tag(m: re.Match) -> str: i += 2 continue elif next_ch == "~": - result.append("\u00A0") # non-breaking space + result.append("\u00a0") # non-breaking space i += 2 continue elif next_ch == "-": - result.append("\u00AD") # soft hyphen + result.append("\u00ad") # soft hyphen i += 2 continue elif next_ch == "_": @@ -244,9 +240,7 @@ def _remove_hex_outside_image_tags(text: str) -> str: last_end = 0 for start, end in protected_ranges: before = text[last_end:start] - before = re.sub( - r"(? str: # Destination Group Removal # ═══════════════════════════════════════════════════════════════════════════ + def remove_destination_groups(content: str) -> str: """ Remove RTF destination groups. @@ -290,9 +285,7 @@ def remove_destination_groups(content: str) -> str: if content[i] == "{" and i + 1 < n and content[i + 1] == "\\": # Check for {\*\word or {\word j = i + 2 - has_star = False if j < n and content[j] == "*": - has_star = True j += 1 # Skip whitespace and the next backslash while j < n and content[j] in " \t\r\n": @@ -466,6 +459,7 @@ def remove_shprslt_blocks(content: str) -> str: # Region Finder (merged from rtf_region_finder.py) # ═══════════════════════════════════════════════════════════════════════════ + def find_excluded_regions(content: str) -> List[Tuple[int, int]]: """ Find document regions to exclude from content extraction. @@ -483,12 +477,12 @@ def find_excluded_regions(content: str) -> List[Tuple[int, int]]: # Header/footer/footnote start patterns start_patterns = [ - r"\\header[lrf]?\b", # Headers (left/right/first) - r"\\footer[lrf]?\b", # Footers - r"\\footnote\b", # Footnotes - r"\\annotation\b", # Annotations/comments - r"\{\\headerf", # First-page header - r"\{\\footerf", # First-page footer + r"\\header[lrf]?\b", # Headers (left/right/first) + r"\\footer[lrf]?\b", # Footers + r"\\footnote\b", # Footnotes + r"\\annotation\b", # Annotations/comments + r"\{\\headerf", # First-page header + r"\{\\footerf", # First-page footer ] for pattern in start_patterns: diff --git a/contextifier/handlers/rtf/_constants.py b/contextifier/handlers/rtf/_constants.py index 9ee66cb..b72111e 100644 --- a/contextifier/handlers/rtf/_constants.py +++ b/contextifier/handlers/rtf/_constants.py @@ -33,26 +33,30 @@ 866: "cp866", 869: "cp869", 874: "cp874", - 932: "cp932", # Japanese Shift-JIS - 936: "gb2312", # Simplified Chinese - 949: "cp949", # Korean - 950: "big5", # Traditional Chinese - 1250: "cp1250", # Central European - 1251: "cp1251", # Cyrillic - 1252: "cp1252", # Western European - 1253: "cp1253", # Greek - 1254: "cp1254", # Turkish - 1255: "cp1255", # Hebrew - 1256: "cp1256", # Arabic - 1257: "cp1257", # Baltic - 1258: "cp1258", # Vietnamese + 932: "cp932", # Japanese Shift-JIS + 936: "gb2312", # Simplified Chinese + 949: "cp949", # Korean + 950: "big5", # Traditional Chinese + 1250: "cp1250", # Central European + 1251: "cp1251", # Cyrillic + 1252: "cp1252", # Western European + 1253: "cp1253", # Greek + 1254: "cp1254", # Turkish + 1255: "cp1255", # Hebrew + 1256: "cp1256", # Arabic + 1257: "cp1257", # Baltic + 1258: "cp1258", # Vietnamese 10000: "mac_roman", 65001: "utf-8", } # Default encoding fallback list DEFAULT_ENCODINGS: List[str] = [ - "utf-8", "cp949", "euc-kr", "cp1252", "latin-1", + "utf-8", + "cp949", + "euc-kr", + "cp1252", + "latin-1", ] @@ -60,18 +64,41 @@ # RTF Destination Groups (to skip during text extraction) # ═══════════════════════════════════════════════════════════════════════════ -SKIP_DESTINATIONS: FrozenSet[str] = frozenset({ - "fonttbl", "colortbl", "stylesheet", "listtable", - "listoverridetable", "revtbl", "rsidtbl", "generator", - "xmlnstbl", "mmathPr", "themedata", "colorschememapping", - "datastore", "latentstyles", "pgptbl", "protusertbl", - "bookmarkstart", "bookmarkend", "bkmkstart", "bkmkend", - "fldinst", "fldrslt", -}) - -IMAGE_DESTINATIONS: FrozenSet[str] = frozenset({ - "pict", "shppict", "nonshppict", "blipuid", -}) +SKIP_DESTINATIONS: FrozenSet[str] = frozenset( + { + "fonttbl", + "colortbl", + "stylesheet", + "listtable", + "listoverridetable", + "revtbl", + "rsidtbl", + "generator", + "xmlnstbl", + "mmathPr", + "themedata", + "colorschememapping", + "datastore", + "latentstyles", + "pgptbl", + "protusertbl", + "bookmarkstart", + "bookmarkend", + "bkmkstart", + "bkmkend", + "fldinst", + "fldrslt", + } +) + +IMAGE_DESTINATIONS: FrozenSet[str] = frozenset( + { + "pict", + "shppict", + "nonshppict", + "blipuid", + } +) # ═══════════════════════════════════════════════════════════════════════════ @@ -79,23 +106,61 @@ # ═══════════════════════════════════════════════════════════════════════════ SHAPE_PROPERTY_NAMES: List[str] = [ - "shapeType", "fFlipH", "fFlipV", "rotation", - "posh", "posrelh", "posv", "posrelv", - "fLayoutInCell", "fAllowOverlap", "fBehindDocument", - "fPseudoInline", "fLockAnchor", "fLockPosition", - "fLockAspectRatio", "fLockRotation", "fLockAgainstSelect", - "fLockCropping", "fLockVerticies", "fLockText", - "fLockAdjustHandles", "fLockAgainstGrouping", - "geoLeft", "geoTop", "geoRight", "geoBottom", - "shapePath", "pWrapPolygonVertices", "dxWrapDistLeft", - "dyWrapDistTop", "dxWrapDistRight", "dyWrapDistBottom", - "fLine", "fFilled", "fillType", "fillColor", - "fillOpacity", "fillBackColor", "fillBackOpacity", - "lineColor", "lineOpacity", "lineWidth", "lineStyle", - "lineDashing", "lineStartArrowhead", "lineStartArrowWidth", - "lineStartArrowLength", "lineEndArrowhead", "lineEndArrowWidth", - "lineEndArrowLength", "shadowType", "shadowColor", - "shadowOpacity", "shadowOffsetX", "shadowOffsetY", + "shapeType", + "fFlipH", + "fFlipV", + "rotation", + "posh", + "posrelh", + "posv", + "posrelv", + "fLayoutInCell", + "fAllowOverlap", + "fBehindDocument", + "fPseudoInline", + "fLockAnchor", + "fLockPosition", + "fLockAspectRatio", + "fLockRotation", + "fLockAgainstSelect", + "fLockCropping", + "fLockVerticies", + "fLockText", + "fLockAdjustHandles", + "fLockAgainstGrouping", + "geoLeft", + "geoTop", + "geoRight", + "geoBottom", + "shapePath", + "pWrapPolygonVertices", + "dxWrapDistLeft", + "dyWrapDistTop", + "dxWrapDistRight", + "dyWrapDistBottom", + "fLine", + "fFilled", + "fillType", + "fillColor", + "fillOpacity", + "fillBackColor", + "fillBackOpacity", + "lineColor", + "lineOpacity", + "lineWidth", + "lineStyle", + "lineDashing", + "lineStartArrowhead", + "lineStartArrowWidth", + "lineStartArrowLength", + "lineEndArrowhead", + "lineEndArrowWidth", + "lineEndArrowLength", + "shadowType", + "shadowColor", + "shadowOpacity", + "shadowOffsetX", + "shadowOffsetY", ] @@ -123,6 +188,11 @@ "wbitmap": "bmp", } -SUPPORTED_IMAGE_FORMATS: FrozenSet[str] = frozenset({ - "jpeg", "png", "gif", "bmp", -}) +SUPPORTED_IMAGE_FORMATS: FrozenSet[str] = frozenset( + { + "jpeg", + "png", + "gif", + "bmp", + } +) diff --git a/contextifier/handlers/rtf/_table_parser.py b/contextifier/handlers/rtf/_table_parser.py index 43b92d0..364e7b5 100644 --- a/contextifier/handlers/rtf/_table_parser.py +++ b/contextifier/handlers/rtf/_table_parser.py @@ -26,7 +26,7 @@ import logging import re -from typing import Dict, List, NamedTuple, Optional, Tuple +from typing import List, NamedTuple, Optional, Tuple from contextifier.types import TableCell, TableData from contextifier.handlers.rtf._decoder import decode_hex_escapes @@ -47,17 +47,20 @@ # Internal Data Structures # ═══════════════════════════════════════════════════════════════════════════ + class _CellDef(NamedTuple): """Cell definition parsed from row header (before \\cell content).""" - h_merge_first: bool # \clmgf — start horizontal merge - h_merge_cont: bool # \clmrg — continue horizontal merge - v_merge_first: bool # \clvmgf — start vertical merge - v_merge_cont: bool # \clvmrg — continue vertical merge - right_boundary: int # \cellxN — right boundary in twips + + h_merge_first: bool # \clmgf — start horizontal merge + h_merge_cont: bool # \clmrg — continue horizontal merge + v_merge_first: bool # \clvmgf — start vertical merge + v_merge_cont: bool # \clvmrg — continue vertical merge + right_boundary: int # \cellxN — right boundary in twips class _ParsedCell(NamedTuple): """Cell with content + merge info, prior to TableCell conversion.""" + text: str h_merge_first: bool h_merge_cont: bool @@ -69,6 +72,7 @@ class _ParsedCell(NamedTuple): # Public API # ═══════════════════════════════════════════════════════════════════════════ + def extract_tables( content: str, encoding: str = "cp949", @@ -109,10 +113,7 @@ def extract_tables_with_positions( excluded_regions = find_excluded_regions(content) # Find all \row end positions - row_ends: List[int] = [ - m.end() - for m in re.finditer(r"\\row(?![a-z])", content) - ] + row_ends: List[int] = [m.end() for m in re.finditer(r"\\row(?![a-z])", content)] if not row_ends: return [], [] @@ -130,7 +131,8 @@ def extract_tables_with_positions( if is_in_excluded_region(row_start, excluded_regions): _logger.debug( - "Skipping table row at %d (header/footer/footnote)", row_start, + "Skipping table row at %d (header/footer/footnote)", + row_start, ) continue @@ -162,6 +164,7 @@ def extract_tables_with_positions( # Row Grouping # ═══════════════════════════════════════════════════════════════════════════ + def _group_rows_into_tables( raw_rows: List[Tuple[int, int, str]], ) -> List[Tuple[int, int, List[str]]]: @@ -204,6 +207,7 @@ def _group_rows_into_tables( # Table Parsing (Rows → TableData) # ═══════════════════════════════════════════════════════════════════════════ + def _parse_table( row_texts: List[str], encoding: str, @@ -246,7 +250,8 @@ def _is_real_table(parsed_rows: List[List[_ParsedCell]]) -> bool: effective_counts: List[int] = [] for row in parsed_rows: non_empty = [ - i for i, cell in enumerate(row) + i + for i, cell in enumerate(row) if not cell.h_merge_cont and (cell.text.strip() or cell.v_merge_first) ] if non_empty: @@ -301,7 +306,10 @@ def _build_table_data( if cell.v_merge_first: rowspan = 1 for nr in range(r_idx + 1, num_rows): - if c_idx < len(parsed_rows[nr]) and parsed_rows[nr][c_idx].v_merge_cont: + if ( + c_idx < len(parsed_rows[nr]) + and parsed_rows[nr][c_idx].v_merge_cont + ): rowspan += 1 merge[nr][c_idx] = (0, 0) # consumed else: @@ -330,14 +338,16 @@ def _build_table_data( cs, rs = 1, 1 text = re.sub(r"\s+", " ", cell.text).strip() - table_cells.append(TableCell( - content=text, - row_span=rs, - col_span=cs, - is_header=False, - row_index=r_idx, - col_index=c_idx, - )) + table_cells.append( + TableCell( + content=text, + row_span=rs, + col_span=cs, + is_header=False, + row_index=r_idx, + col_index=c_idx, + ) + ) table_rows.append(table_cells) return TableData( @@ -352,6 +362,7 @@ def _build_table_data( # Cell Extraction (single row) # ═══════════════════════════════════════════════════════════════════════════ + def _extract_cells_with_merge( row_text: str, encoding: str, @@ -377,21 +388,25 @@ def _extract_cells_with_merge( for i, text in enumerate(cell_texts): if i < len(cell_defs): d = cell_defs[i] - cells.append(_ParsedCell( - text=text, - h_merge_first=d.h_merge_first, - h_merge_cont=d.h_merge_cont, - v_merge_first=d.v_merge_first, - v_merge_cont=d.v_merge_cont, - )) + cells.append( + _ParsedCell( + text=text, + h_merge_first=d.h_merge_first, + h_merge_cont=d.h_merge_cont, + v_merge_first=d.v_merge_first, + v_merge_cont=d.v_merge_cont, + ) + ) else: - cells.append(_ParsedCell( - text=text, - h_merge_first=False, - h_merge_cont=False, - v_merge_first=False, - v_merge_cont=False, - )) + cells.append( + _ParsedCell( + text=text, + h_merge_first=False, + h_merge_cont=False, + v_merge_first=False, + v_merge_cont=False, + ) + ) return cells @@ -445,13 +460,15 @@ def _parse_cell_definitions(def_part: str) -> List[_CellDef]: v_cont = True elif token.startswith("\\cellx"): boundary = int(match.group(1)) if match.group(1) else 0 - defs.append(_CellDef( - h_merge_first=h_first, - h_merge_cont=h_cont, - v_merge_first=v_first, - v_merge_cont=v_cont, - right_boundary=boundary, - )) + defs.append( + _CellDef( + h_merge_first=h_first, + h_merge_cont=h_cont, + v_merge_first=v_first, + v_merge_cont=v_cont, + right_boundary=boundary, + ) + ) # Reset for next cell definition h_first = h_cont = v_first = v_cont = False @@ -510,6 +527,7 @@ def _extract_cell_texts( # Utility: Convert single-column "table" to text # ═══════════════════════════════════════════════════════════════════════════ + def single_column_to_text( row_texts: List[str], encoding: str, diff --git a/contextifier/handlers/rtf/content_extractor.py b/contextifier/handlers/rtf/content_extractor.py index b62b02d..e926438 100644 --- a/contextifier/handlers/rtf/content_extractor.py +++ b/contextifier/handlers/rtf/content_extractor.py @@ -32,12 +32,11 @@ import logging import re -from typing import Any, Dict, List, Optional, Tuple, TYPE_CHECKING +from typing import Any, List, Optional, Tuple, TYPE_CHECKING from contextifier.pipeline.content_extractor import BaseContentExtractor from contextifier.services.table_service import TableService from contextifier.types import ( - ExtractionResult, PreprocessedData, TableData, ) @@ -53,13 +52,9 @@ ) from contextifier.handlers.rtf._table_parser import ( extract_tables_with_positions, - single_column_to_text, ) if TYPE_CHECKING: - from contextifier.services.image_service import ImageService - from contextifier.services.tag_service import TagService - from contextifier.services.chart_service import ChartService from contextifier.services.table_service import TableService _logger = logging.getLogger("contextifier.rtf.content") @@ -136,7 +131,8 @@ def extract_tables( return [] _, table_regions = extract_tables_with_positions( - parsed.text, encoding, + parsed.text, + encoding, ) return [table for _, _, table in table_regions] @@ -176,15 +172,16 @@ def _build_inline_content( if not table_regions: clean = self._clean_segment( - content[header_end:], header_end, encoding, excluded, + content[header_end:], + header_end, + encoding, + excluded, ) return clean.strip() # Adjust table regions to start after header adjusted = [ - (max(s, header_end), e, t) - for s, e, t in table_regions - if e > header_end + (max(s, header_end), e, t) for s, e, t in table_regions if e > header_end ] parts: List[str] = [] @@ -195,7 +192,10 @@ def _build_inline_content( if start_pos > last_end: segment = content[last_end:start_pos] clean = self._clean_segment( - segment, last_end, encoding, excluded, + segment, + last_end, + encoding, + excluded, ) if clean.strip(): parts.append(clean) @@ -211,7 +211,10 @@ def _build_inline_content( if last_end < len(content): segment = content[last_end:] clean = self._clean_segment( - segment, last_end, encoding, excluded, + segment, + last_end, + encoding, + excluded, ) if clean.strip(): parts.append(clean) @@ -307,9 +310,7 @@ def _save_images(self, preprocessed: PreprocessedData) -> List[str]: return [] images: List[RtfImageData] = ( - preprocessed.resources.get("images", []) - if preprocessed.resources - else [] + preprocessed.resources.get("images", []) if preprocessed.resources else [] ) if not images: return [] @@ -340,6 +341,7 @@ def _fallback_striprtf(parsed: RtfParsedData) -> str: """ try: from striprtf.striprtf import rtf_to_text # type: ignore + text = rtf_to_text(parsed.text) if text: _logger.info("Using striprtf fallback (%d chars)", len(text)) diff --git a/contextifier/handlers/rtf/converter.py b/contextifier/handlers/rtf/converter.py index abf50c2..69a3c6e 100644 --- a/contextifier/handlers/rtf/converter.py +++ b/contextifier/handlers/rtf/converter.py @@ -39,6 +39,7 @@ class RtfConvertedData(NamedTuple): encoding: Detected encoding from \\ansicpg (hint for preprocessor). file_extension: From FileContext ("rtf"). """ + raw_bytes: bytes encoding: str file_extension: str @@ -85,7 +86,7 @@ def convert( stripped = file_data.lstrip() if not stripped.startswith(_RTF_MAGIC): raise ConversionError( - f"Not a valid RTF file (missing {{\\rtf header)", + "Not a valid RTF file (missing {\\rtf header)", stage="convert", handler="RtfConverter", ) diff --git a/contextifier/handlers/rtf/metadata_extractor.py b/contextifier/handlers/rtf/metadata_extractor.py index 9260aaf..07238d9 100644 --- a/contextifier/handlers/rtf/metadata_extractor.py +++ b/contextifier/handlers/rtf/metadata_extractor.py @@ -35,11 +35,11 @@ # ── Field extraction patterns (within \info group) ────────────────────── # RTF structure: {\title content} — the { precedes the control word _METADATA_FIELDS: Dict[str, str] = { - "title": r"\{\\title\s+([^}]*)\}", - "subject": r"\{\\subject\s+([^}]*)\}", - "author": r"\{\\author\s+([^}]*)\}", - "keywords": r"\{\\keywords\s+([^}]*)\}", - "comments": r"\{\\doccomm\s+([^}]*)\}", + "title": r"\{\\title\s+([^}]*)\}", + "subject": r"\{\\subject\s+([^}]*)\}", + "author": r"\{\\author\s+([^}]*)\}", + "keywords": r"\{\\keywords\s+([^}]*)\}", + "comments": r"\{\\doccomm\s+([^}]*)\}", "last_saved_by": r"\{\\operator\s+([^}]*)\}", } @@ -94,7 +94,9 @@ def extract(self, source: Any) -> DocumentMetadata: fields: Dict[str, Optional[str]] = {} for field_name, pattern in _METADATA_FIELDS.items(): fields[field_name] = self._extract_field( - info_content, pattern, encoding, + info_content, + pattern, + encoding, ) # Extract date fields (search full content, not just \info) diff --git a/contextifier/handlers/rtf/preprocessor.py b/contextifier/handlers/rtf/preprocessor.py index 828250c..e5fe038 100644 --- a/contextifier/handlers/rtf/preprocessor.py +++ b/contextifier/handlers/rtf/preprocessor.py @@ -35,7 +35,7 @@ import hashlib import logging import re -from typing import Any, Dict, List, NamedTuple, Optional, Set, Tuple +from typing import Any, List, NamedTuple, Set, Tuple from contextifier.pipeline.preprocessor import BasePreprocessor from contextifier.types import PreprocessedData @@ -57,6 +57,7 @@ # Data Structures # ═══════════════════════════════════════════════════════════════════════════ + class RtfImageData(NamedTuple): """ Extracted image from RTF binary content. @@ -64,10 +65,11 @@ class RtfImageData(NamedTuple): Stored in PreprocessedData.resources["images"]. ContentExtractor saves these via ImageService. """ - image_format: str # "jpeg", "png", "gif", "bmp", etc. - image_bytes: bytes # Raw image binary data - position: int # Position in original RTF stream - content_hash: str # MD5 hash for deduplication + + image_format: str # "jpeg", "png", "gif", "bmp", etc. + image_bytes: bytes # Raw image binary data + position: int # Position in original RTF stream + content_hash: str # MD5 hash for deduplication class RtfParsedData(NamedTuple): @@ -77,16 +79,18 @@ class RtfParsedData(NamedTuple): Stored in PreprocessedData.content. Consumed by RtfMetadataExtractor and RtfContentExtractor. """ - text: str # Decoded RTF string (binary data removed) - encoding: str # Detected encoding - image_count: int # Number of images extracted + + text: str # Decoded RTF string (binary data removed) + encoding: str # Detected encoding + image_count: int # Number of images extracted class _BinaryRegion(NamedTuple): """Internal: describes a binary data region in RTF bytes.""" + start_pos: int end_pos: int - bin_type: str # "bin" or "pict" + bin_type: str # "bin" or "pict" image_format: str image_data: bytes @@ -95,6 +99,7 @@ class _BinaryRegion(NamedTuple): # Preprocessor # ═══════════════════════════════════════════════════════════════════════════ + class RtfPreprocessor(BasePreprocessor): """ RTF-specific preprocessor. @@ -217,12 +222,14 @@ def _process_binary( continue seen_hashes.add(content_hash) - images.append(RtfImageData( - image_format=region.image_format, - image_bytes=region.image_data, - position=region.start_pos, - content_hash=content_hash, - )) + images.append( + RtfImageData( + image_format=region.image_format, + image_bytes=region.image_data, + position=region.start_pos, + content_hash=content_hash, + ) + ) # Remove binary regions from content clean_bytes = _remove_binary_regions(content, all_regions) @@ -234,6 +241,7 @@ def _process_binary( # Binary Region Discovery (module-level for testability) # ═══════════════════════════════════════════════════════════════════════════ + def _find_bin_regions(content: bytes) -> List[_BinaryRegion]: """ Find ``\\binN`` tags that mark raw binary data embedded in RTF. @@ -257,7 +265,10 @@ def _find_bin_regions(content: bytes) -> List[_BinaryRegion]: data_start = bin_tag_end # Skip optional space after \binN - if data_start < len(content) and content[data_start : data_start + 1] == b" ": + if ( + data_start < len(content) + and content[data_start : data_start + 1] == b" " + ): data_start += 1 data_end = data_start + bin_size @@ -293,13 +304,15 @@ def _find_bin_regions(content: bytes) -> List[_BinaryRegion]: j += 1 group_end = j - regions.append(_BinaryRegion( - start_pos=group_start, - end_pos=group_end, - bin_type="bin", - image_format=image_format, - image_data=binary_data, - )) + regions.append( + _BinaryRegion( + start_pos=group_start, + end_pos=group_end, + bin_type="bin", + image_format=image_format, + image_data=binary_data, + ) + ) except (ValueError, IndexError): continue @@ -387,13 +400,15 @@ def has_bin_nearby(pict_pos: int) -> bool: image_format = _detect_image_format(image_data) if image_format: - regions.append(_BinaryRegion( - start_pos=start_pos, - end_pos=i, - bin_type="pict", - image_format=image_format, - image_data=image_data, - )) + regions.append( + _BinaryRegion( + start_pos=start_pos, + end_pos=i, + bin_type="pict", + image_format=image_format, + image_data=image_data, + ) + ) except ValueError: continue except Exception as e: diff --git a/contextifier/handlers/text/content_extractor.py b/contextifier/handlers/text/content_extractor.py index 7542d93..512c88e 100644 --- a/contextifier/handlers/text/content_extractor.py +++ b/contextifier/handlers/text/content_extractor.py @@ -38,12 +38,14 @@ # File categories where content should be treated as code # (preserve indentation, replace tabs, don't collapse blank lines) -_CODE_CATEGORIES: FrozenSet[str] = frozenset({ - "code", # .py, .js, .ts, .java, .cpp, ... - "config", # .json, .yaml, .xml, .toml, .ini, ... - "script", # .sh, .bat, .ps1, ... - "web", # .htm, .xhtml -}) +_CODE_CATEGORIES: FrozenSet[str] = frozenset( + { + "code", # .py, .js, .ts, .java, .cpp, ... + "config", # .json, .yaml, .xml, .toml, .ini, ... + "script", # .sh, .bat, .ps1, ... + "web", # .htm, .xhtml + } +) class TextContentExtractor(BaseContentExtractor): diff --git a/contextifier/handlers/text/converter.py b/contextifier/handlers/text/converter.py index 9efe774..ead9123 100644 --- a/contextifier/handlers/text/converter.py +++ b/contextifier/handlers/text/converter.py @@ -23,12 +23,10 @@ from __future__ import annotations -import logging from typing import Any, List, NamedTuple, Optional from contextifier.pipeline.converter import BaseConverter from contextifier.types import FileContext -from contextifier.errors import ConversionError class TextConvertedData(NamedTuple): @@ -38,6 +36,7 @@ class TextConvertedData(NamedTuple): Carries decoded text along with metadata that downstream pipeline stages need for format-aware processing. """ + text: str encoding: str file_extension: str diff --git a/contextifier/handlers/text/handler.py b/contextifier/handlers/text/handler.py index 92b9a07..dbb2433 100644 --- a/contextifier/handlers/text/handler.py +++ b/contextifier/handlers/text/handler.py @@ -47,27 +47,75 @@ # All text-based extensions supported by this handler. # Covers: plain text, markup, source code, config, scripts, stylesheets. -_TEXT_EXTENSIONS = frozenset({ - # Plain text & markup - "txt", "md", "markdown", "rst", "log", - # Config & data formats - "cfg", "ini", "conf", "yaml", "yml", "toml", "json", "xml", "svg", - "properties", "env", - # Source code - "py", "js", "ts", "jsx", "tsx", "java", "cpp", "c", "h", "hpp", - "cs", "go", "rs", "php", "rb", "swift", "kt", "scala", "dart", - "r", "m", "lua", "pl", "pm", - # Scripts - "sh", "bash", "zsh", "bat", "ps1", "cmd", "fish", - # SQL - "sql", - # Web - "css", "scss", "less", "sass", - # Vue / Svelte (single-file components) - "vue", "svelte", - # Dotfiles - "gitignore", "dockerignore", "editorconfig", -}) +_TEXT_EXTENSIONS = frozenset( + { + # Plain text & markup + "txt", + "md", + "markdown", + "rst", + "log", + # Config & data formats + "cfg", + "ini", + "conf", + "yaml", + "yml", + "toml", + "json", + "xml", + "svg", + "properties", + "env", + # Source code + "py", + "js", + "ts", + "jsx", + "tsx", + "java", + "cpp", + "c", + "h", + "hpp", + "cs", + "go", + "rs", + "php", + "rb", + "swift", + "kt", + "scala", + "dart", + "r", + "m", + "lua", + "pl", + "pm", + # Scripts + "sh", + "bash", + "zsh", + "bat", + "ps1", + "cmd", + "fish", + # SQL + "sql", + # Web + "css", + "scss", + "less", + "sass", + # Vue / Svelte (single-file components) + "vue", + "svelte", + # Dotfiles + "gitignore", + "dockerignore", + "editorconfig", + } +) class TextHandler(BaseHandler): diff --git a/contextifier/handlers/text/preprocessor.py b/contextifier/handlers/text/preprocessor.py index acbcbca..f1c4bc0 100644 --- a/contextifier/handlers/text/preprocessor.py +++ b/contextifier/handlers/text/preprocessor.py @@ -20,8 +20,7 @@ from __future__ import annotations -import logging -from typing import Any, Union +from typing import Any from contextifier.pipeline.preprocessor import BasePreprocessor from contextifier.types import PreprocessedData @@ -129,7 +128,12 @@ def _unpack( ) # Last resort - return (str(converted_data) if converted_data is not None else "", "utf-8", "", "") + return ( + str(converted_data) if converted_data is not None else "", + "utf-8", + "", + "", + ) __all__ = ["TextPreprocessor"] diff --git a/contextifier/handlers/xls/_layout.py b/contextifier/handlers/xls/_layout.py index 1843948..26ab0d9 100644 --- a/contextifier/handlers/xls/_layout.py +++ b/contextifier/handlers/xls/_layout.py @@ -40,7 +40,9 @@ def cell_count(self) -> int: return self.rows * self.cols def contains(self, row: int, col: int) -> bool: - return self.min_row <= row <= self.max_row and self.min_col <= col <= self.max_col + return ( + self.min_row <= row <= self.max_row and self.min_col <= col <= self.max_col + ) def overlaps(self, other: LayoutRange) -> bool: return not ( @@ -123,7 +125,9 @@ def layout_detect_range(sheet: object) -> Optional[LayoutRange]: # ── Object detection ───────────────────────────────────────────────────────── -def object_detect(sheet: object, book: object = None, layout: Optional[LayoutRange] = None) -> List[LayoutRange]: +def object_detect( + sheet: object, book: object = None, layout: Optional[LayoutRange] = None +) -> List[LayoutRange]: """ Detect individual data regions in an xlrd sheet. @@ -140,7 +144,9 @@ def object_detect(sheet: object, book: object = None, layout: Optional[LayoutRan if layout is None: return [] - bordered = _detect_bordered_regions(sheet, book, layout) if book is not None else [] + bordered = ( + _detect_bordered_regions(sheet, book, layout) if book is not None else [] + ) value_regions = _detect_value_regions(sheet, layout, bordered) all_regions = bordered + value_regions @@ -162,7 +168,12 @@ def _has_border(sheet: object, book: object, row0: int, col0: int) -> bool: try: xf_idx = sheet.cell_xf_index(row0, col0) # type: ignore[attr-defined] xf = book.xf_list[xf_idx] # type: ignore[attr-defined] - for attr in ("top_line_style", "bottom_line_style", "left_line_style", "right_line_style"): + for attr in ( + "top_line_style", + "bottom_line_style", + "left_line_style", + "right_line_style", + ): val = getattr(xf.border, attr, 0) if val and val > 0: return True @@ -171,7 +182,9 @@ def _has_border(sheet: object, book: object, row0: int, col0: int) -> bool: return False -def _detect_bordered_regions(sheet: object, book: object, layout: LayoutRange) -> List[LayoutRange]: +def _detect_bordered_regions( + sheet: object, book: object, layout: LayoutRange +) -> List[LayoutRange]: bordered: Set[Tuple[int, int]] = set() for r1 in range(layout.min_row, layout.max_row + 1): for c1 in range(layout.min_col, layout.max_col + 1): @@ -185,7 +198,9 @@ def _detect_bordered_regions(sheet: object, book: object, layout: LayoutRange) - # ── Value detection ────────────────────────────────────────────────────────── -def _detect_value_regions(sheet: object, layout: LayoutRange, exclude: List[LayoutRange]) -> List[LayoutRange]: +def _detect_value_regions( + sheet: object, layout: LayoutRange, exclude: List[LayoutRange] +) -> List[LayoutRange]: def _excluded(r: int, c: int) -> bool: return any(reg.contains(r, c) for reg in exclude) @@ -232,7 +247,9 @@ def _bfs_group(cells: Set[Tuple[int, int]]) -> List[LayoutRange]: max_r = max(r for r, _ in group) min_c = min(c for _, c in group) max_c = max(c for _, c in group) - regions.append(LayoutRange(min_row=min_r, max_row=max_r, min_col=min_c, max_col=max_c)) + regions.append( + LayoutRange(min_row=min_r, max_row=max_r, min_col=min_c, max_col=max_c) + ) return regions diff --git a/contextifier/handlers/xls/_table.py b/contextifier/handlers/xls/_table.py index 2e245cc..50d2e6e 100644 --- a/contextifier/handlers/xls/_table.py +++ b/contextifier/handlers/xls/_table.py @@ -11,7 +11,7 @@ from __future__ import annotations import logging -from typing import Any, List, Optional, Set, Tuple +from typing import List, Optional, Set, Tuple from contextifier.types import TableData, TableCell @@ -184,7 +184,12 @@ def _has_merged_in_region(sheet: object, region: LayoutRange) -> bool: for rlo, rhi, clo, chi in getattr(sheet, "merged_cells", []): mr_min, mr_max = rlo + 1, rhi mc_min, mc_max = clo + 1, chi - if mr_min <= region.max_row and mr_max >= region.min_row and mc_min <= region.max_col and mc_max >= region.min_col: + if ( + mr_min <= region.max_row + and mr_max >= region.min_row + and mc_min <= region.max_col + and mc_max >= region.min_col + ): return True return False @@ -198,7 +203,12 @@ def _get_merged_in_region( for rlo, rhi, clo, chi in getattr(sheet, "merged_cells", []): mr_min, mr_max = rlo + 1, rhi mc_min, mc_max = clo + 1, chi - if mr_min <= region.max_row and mr_max >= region.min_row and mc_min <= region.max_col and mc_max >= region.min_col: + if ( + mr_min <= region.max_row + and mr_max >= region.min_row + and mc_min <= region.max_col + and mc_max >= region.min_col + ): result.append((rlo, rhi, clo, chi)) return result diff --git a/contextifier/handlers/xls/content_extractor.py b/contextifier/handlers/xls/content_extractor.py index a1d40ab..96fdab5 100644 --- a/contextifier/handlers/xls/content_extractor.py +++ b/contextifier/handlers/xls/content_extractor.py @@ -22,14 +22,11 @@ from contextifier.pipeline.content_extractor import BaseContentExtractor from contextifier.types import ( ChartData, - ExtractionResult, PreprocessedData, TableData, ) from contextifier.handlers.xls._layout import ( - LayoutRange, - layout_detect_range, object_detect, ) from contextifier.handlers.xls._table import ( @@ -41,21 +38,28 @@ # Image signatures for OLE stream scanning (same as DOC handler) _IMAGE_SIGNATURES: dict[str, tuple[bytes, int]] = { - "png": (b"\x89PNG\r\n\x1a\n", 8), - "jpeg": (b"\xff\xd8", 2), - "gif87": (b"GIF87a", 6), - "gif89": (b"GIF89a", 6), - "bmp": (b"BM", 2), - "tiff_le": (b"II\x2a\x00", 4), - "tiff_be": (b"MM\x00\x2a", 4), - "emf": (b"\x01\x00\x00\x00", 4), + "png": (b"\x89PNG\r\n\x1a\n", 8), + "jpeg": (b"\xff\xd8", 2), + "gif87": (b"GIF87a", 6), + "gif89": (b"GIF89a", 6), + "bmp": (b"BM", 2), + "tiff_le": (b"II\x2a\x00", 4), + "tiff_be": (b"MM\x00\x2a", 4), + "emf": (b"\x01\x00\x00\x00", 4), } # OLE stream keywords that may contain images -_IMAGE_STREAM_KEYWORDS: frozenset[str] = frozenset({ - "pictures", "data", "object", "oleobject", "objectpool", - "mbd", "workbook", # Some images are embedded in Drawing records -}) +_IMAGE_STREAM_KEYWORDS: frozenset[str] = frozenset( + { + "pictures", + "data", + "object", + "oleobject", + "objectpool", + "mbd", + "workbook", # Some images are embedded in Drawing records + } +) # BIFF sheet type for charts _XL_CHART_SHEET = 2 @@ -109,7 +113,9 @@ def extract_text(self, preprocessed: PreprocessedData, **kw: Any) -> str: # ── tables ─────────────────────────────────────────────────────────── - def extract_tables(self, preprocessed: PreprocessedData, **kw: Any) -> List[TableData]: + def extract_tables( + self, preprocessed: PreprocessedData, **kw: Any + ) -> List[TableData]: book = self._get_book(preprocessed) if book is None: return [] @@ -154,7 +160,7 @@ def extract_images(self, preprocessed: PreprocessedData, **kw: Any) -> List[str] try: for entry in ole.listdir(): - entry_path = "/".join(entry) + "/".join(entry) if not any( kw in part.lower() for part in entry @@ -195,7 +201,9 @@ def extract_images(self, preprocessed: PreprocessedData, **kw: Any) -> List[str] # ── charts (BIFF sheet type detection) ─────────────────────────────── - def extract_charts(self, preprocessed: PreprocessedData, **kw: Any) -> List[ChartData]: + def extract_charts( + self, preprocessed: PreprocessedData, **kw: Any + ) -> List[ChartData]: """ Detect chart sheets in the XLS workbook. @@ -216,16 +224,20 @@ def extract_charts(self, preprocessed: PreprocessedData, **kw: Any) -> List[Char sheet = book.sheet_by_index(idx) # xlrd Sheet objects have a `sheet_type` attribute in # the internal book.sheet_types list - sheet_type = book.sheet_type(idx) if hasattr(book, "sheet_type") else None + sheet_type = ( + book.sheet_type(idx) if hasattr(book, "sheet_type") else None + ) except Exception: continue if sheet_type == _XL_CHART_SHEET: - charts.append(ChartData( - chart_type="biff_chart_sheet", - title=sheet.name, - raw_content=f"BIFF chart sheet: {sheet.name}", - )) + charts.append( + ChartData( + chart_type="biff_chart_sheet", + title=sheet.name, + raw_content=f"BIFF chart sheet: {sheet.name}", + ) + ) return charts @@ -234,7 +246,11 @@ def extract_charts(self, preprocessed: PreprocessedData, **kw: Any) -> List[Char def _get_book(self, preprocessed: PreprocessedData) -> Any: if preprocessed is None: return None - content = preprocessed.content if isinstance(preprocessed, PreprocessedData) else preprocessed + content = ( + preprocessed.content + if isinstance(preprocessed, PreprocessedData) + else preprocessed + ) if content is None: return None if hasattr(content, "nsheets"): @@ -252,6 +268,7 @@ def _make_sheet_tag(self, name: str) -> str: # ── Module-level helpers ───────────────────────────────────────────────────── + def _detect_image_format(data: bytes) -> Optional[str]: """Detect image format from binary data using header signatures.""" if not data or len(data) < 2: diff --git a/contextifier/handlers/xls/converter.py b/contextifier/handlers/xls/converter.py index 0936e89..12385e8 100644 --- a/contextifier/handlers/xls/converter.py +++ b/contextifier/handlers/xls/converter.py @@ -8,7 +8,6 @@ from __future__ import annotations -import io import logging from typing import Any, NamedTuple @@ -25,8 +24,9 @@ class XlsConvertedData(NamedTuple): """Wrapper returned by XlsConverter.convert().""" - book: Any # xlrd.Book - file_data: bytes # original bytes (for OLE metadata) + + book: Any # xlrd.Book + file_data: bytes # original bytes (for OLE metadata) class XlsConverter(BaseConverter): diff --git a/contextifier/handlers/xls/metadata_extractor.py b/contextifier/handlers/xls/metadata_extractor.py index 4335236..78ba309 100644 --- a/contextifier/handlers/xls/metadata_extractor.py +++ b/contextifier/handlers/xls/metadata_extractor.py @@ -93,7 +93,9 @@ def _extract_from_ole(self, file_data: bytes) -> DocumentMetadata: comments = self._safe_str(getattr(meta_obj, "comments", None)) last_saved_by = self._safe_str(getattr(meta_obj, "last_saved_by", None)) create_time = self._safe_datetime(getattr(meta_obj, "create_time", None)) - last_saved_time = self._safe_datetime(getattr(meta_obj, "last_saved_time", None)) + last_saved_time = self._safe_datetime( + getattr(meta_obj, "last_saved_time", None) + ) category = self._safe_str(getattr(meta_obj, "category", None)) revision = self._safe_str(getattr(meta_obj, "revision_number", None)) diff --git a/contextifier/handlers/xlsx/_constants.py b/contextifier/handlers/xlsx/_constants.py index be0aa0f..56f15bf 100644 --- a/contextifier/handlers/xlsx/_constants.py +++ b/contextifier/handlers/xlsx/_constants.py @@ -25,7 +25,9 @@ NS_DRAWING_MAIN = "http://schemas.openxmlformats.org/drawingml/2006/main" NS_SPREADSHEET = "http://schemas.openxmlformats.org/spreadsheetml/2006/main" NS_RELATIONSHIPS = "http://schemas.openxmlformats.org/officeDocument/2006/relationships" -NS_SPREADSHEET_DRAWING = "http://schemas.openxmlformats.org/drawingml/2006/spreadsheetDrawing" +NS_SPREADSHEET_DRAWING = ( + "http://schemas.openxmlformats.org/drawingml/2006/spreadsheetDrawing" +) NS_PACKAGE = "http://schemas.openxmlformats.org/package/2006/relationships" OOXML_NS = { @@ -64,7 +66,9 @@ # Image file extensions # ═══════════════════════════════════════════════════════════════════════════════ -SUPPORTED_IMAGE_EXTENSIONS = frozenset({".png", ".jpg", ".jpeg", ".gif", ".bmp", ".tiff"}) +SUPPORTED_IMAGE_EXTENSIONS = frozenset( + {".png", ".jpg", ".jpeg", ".gif", ".bmp", ".tiff"} +) UNSUPPORTED_IMAGE_EXTENSIONS = frozenset({".emf", ".wmf"}) # ═══════════════════════════════════════════════════════════════════════════════ diff --git a/contextifier/handlers/xlsx/_layout.py b/contextifier/handlers/xlsx/_layout.py index fecfef8..50249f9 100644 --- a/contextifier/handlers/xlsx/_layout.py +++ b/contextifier/handlers/xlsx/_layout.py @@ -15,7 +15,7 @@ import logging from collections import deque from dataclasses import dataclass -from typing import Dict, List, Optional, Set, Tuple +from typing import List, Optional, Set, Tuple from contextifier.handlers.xlsx._constants import MAX_SCAN_ROWS, MAX_SCAN_COLS @@ -45,8 +45,12 @@ def cell_count(self) -> int: def is_adjacent(self, other: "LayoutRange", *, tolerance: int = 1) -> bool: """Check if two ranges are adjacent (within tolerance rows/cols).""" - row_gap = max(0, max(self.min_row, other.min_row) - min(self.max_row, other.max_row) - 1) - col_gap = max(0, max(self.min_col, other.min_col) - min(self.max_col, other.max_col) - 1) + row_gap = max( + 0, max(self.min_row, other.min_row) - min(self.max_row, other.max_row) - 1 + ) + col_gap = max( + 0, max(self.min_col, other.min_col) - min(self.max_col, other.max_col) - 1 + ) return row_gap <= tolerance and col_gap <= tolerance def merge_with(self, other: "LayoutRange") -> "LayoutRange": @@ -70,8 +74,7 @@ def overlaps(self, other: "LayoutRange") -> bool: def contains(self, row: int, col: int) -> bool: """Check if a cell is within this range.""" return ( - self.min_row <= row <= self.max_row - and self.min_col <= col <= self.max_col + self.min_row <= row <= self.max_row and self.min_col <= col <= self.max_col ) diff --git a/contextifier/handlers/xlsx/_table.py b/contextifier/handlers/xlsx/_table.py index 2506d22..60e3cc8 100644 --- a/contextifier/handlers/xlsx/_table.py +++ b/contextifier/handlers/xlsx/_table.py @@ -64,16 +64,18 @@ def convert_region_to_table( continue skip_cells.add((row_idx + dr, col_idx + dc)) - is_header = (row_idx == region.min_row) - - row_cells.append(TableCell( - content=value, - row_span=row_span, - col_span=col_span, - is_header=is_header, - row_index=row_idx - region.min_row, - col_index=col_idx - region.min_col, - )) + is_header = row_idx == region.min_row + + row_cells.append( + TableCell( + content=value, + row_span=row_span, + col_span=col_span, + is_header=is_header, + row_index=row_idx - region.min_row, + col_index=col_idx - region.min_col, + ) + ) if row_cells: rows.append(row_cells) @@ -82,11 +84,7 @@ def convert_region_to_table( return None # Filter out completely empty tables - has_content = any( - cell.content.strip() - for row in rows - for cell in row - ) + has_content = any(cell.content.strip() for row in rows for cell in row) if not has_content: return None @@ -157,7 +155,7 @@ def convert_region_to_html( for row_idx in range(region.min_row, region.max_row + 1): cells_html: List[str] = [] - is_header = (row_idx == region.min_row) + is_header = row_idx == region.min_row for col_idx in range(region.min_col, region.max_col + 1): if (row_idx, col_idx) in skip_cells: @@ -215,6 +213,7 @@ def convert_sheet_to_text( # Internal helpers # ═══════════════════════════════════════════════════════════════════════════════ + def _get_merged_cells_in_region( ws: object, region: LayoutRange, diff --git a/contextifier/handlers/xlsx/content_extractor.py b/contextifier/handlers/xlsx/content_extractor.py index 65e54eb..e9c30dd 100644 --- a/contextifier/handlers/xlsx/content_extractor.py +++ b/contextifier/handlers/xlsx/content_extractor.py @@ -25,7 +25,6 @@ TableData, ) from contextifier.handlers.xlsx._layout import ( - LayoutRange, layout_detect_range, object_detect, ) @@ -96,7 +95,9 @@ def extract_text( if not self._include_hidden_sheets: state = getattr(ws, "sheet_state", "visible") if state != "visible": - logger.debug("Skipping hidden sheet '%s' (state=%s)", sheet_name, state) + logger.debug( + "Skipping hidden sheet '%s' (state=%s)", sheet_name, state + ) continue sheet_parts: List[str] = [] @@ -134,7 +135,9 @@ def extract_text( sheet_parts.append(chart_text) chart_index += 1 except Exception as exc: - logger.debug("Error processing charts for sheet %s: %s", sheet_name, exc) + logger.debug( + "Error processing charts for sheet %s: %s", sheet_name, exc + ) # Process per-sheet images try: @@ -142,7 +145,9 @@ def extract_text( if sheet_image_tags: sheet_parts.extend(sheet_image_tags) except Exception as exc: - logger.debug("Error processing images for sheet %s: %s", sheet_name, exc) + logger.debug( + "Error processing images for sheet %s: %s", sheet_name, exc + ) # Textboxes sheet_textboxes = textboxes.get(sheet_name, []) @@ -248,12 +253,14 @@ def extract_charts( ) for s in chart_dict.get("series", []) ] - result.append(ChartData( - chart_type=chart_dict.get("chart_type", "Unknown"), - title=chart_dict.get("title", ""), - categories=chart_dict.get("categories", []), - series=series, - )) + result.append( + ChartData( + chart_type=chart_dict.get("chart_type", "Unknown"), + title=chart_dict.get("title", ""), + categories=chart_dict.get("categories", []), + series=series, + ) + ) except (KeyError, TypeError, ValueError) as exc: logger.debug("Failed to parse chart data: %s", exc) diff --git a/contextifier/handlers/xlsx/converter.py b/contextifier/handlers/xlsx/converter.py index 8c7f4bc..d64ab5f 100644 --- a/contextifier/handlers/xlsx/converter.py +++ b/contextifier/handlers/xlsx/converter.py @@ -30,8 +30,9 @@ class XlsxConvertedData(NamedTuple): """Result of the XLSX conversion stage.""" - workbook: openpyxl.Workbook # Opened openpyxl Workbook - file_data: bytes # Original file bytes (for ZIP re-reading) + + workbook: openpyxl.Workbook # Opened openpyxl Workbook + file_data: bytes # Original file bytes (for ZIP re-reading) class XlsxConverter(BaseConverter): diff --git a/contextifier/handlers/xlsx/handler.py b/contextifier/handlers/xlsx/handler.py index 7381205..1222abe 100644 --- a/contextifier/handlers/xlsx/handler.py +++ b/contextifier/handlers/xlsx/handler.py @@ -17,7 +17,7 @@ from __future__ import annotations -from typing import Any, FrozenSet +from typing import FrozenSet from contextifier.handlers.base import BaseHandler from contextifier.pipeline.converter import BaseConverter diff --git a/contextifier/handlers/xlsx/metadata_extractor.py b/contextifier/handlers/xlsx/metadata_extractor.py index 6346e46..af348c2 100644 --- a/contextifier/handlers/xlsx/metadata_extractor.py +++ b/contextifier/handlers/xlsx/metadata_extractor.py @@ -56,7 +56,9 @@ def extract(self, source: Any) -> DocumentMetadata: create_time=self._safe_datetime(props.created), last_saved_time=self._safe_datetime(props.modified), category=self._safe_str(props.category), - revision=self._safe_str(props.revision) if hasattr(props, "revision") else None, + revision=self._safe_str(props.revision) + if hasattr(props, "revision") + else None, page_count=len(wb.sheetnames), ) except Exception as exc: diff --git a/contextifier/handlers/xlsx/preprocessor.py b/contextifier/handlers/xlsx/preprocessor.py index 0a6b693..3674dd4 100644 --- a/contextifier/handlers/xlsx/preprocessor.py +++ b/contextifier/handlers/xlsx/preprocessor.py @@ -127,6 +127,7 @@ def get_format_name(self) -> str: # ZIP-level extraction helpers # ═══════════════════════════════════════════════════════════════════════════════ + def _extract_charts_from_zip(file_data: bytes) -> List[dict]: """ Extract all chart data from ``xl/charts/chart*.xml`` in the ZIP. @@ -139,7 +140,8 @@ def _extract_charts_from_zip(file_data: bytes) -> List[dict]: try: with zipfile.ZipFile(io.BytesIO(file_data)) as zf: chart_files = sorted( - n for n in zf.namelist() + n + for n in zf.namelist() if n.startswith("xl/charts/chart") and n.endswith(".xml") ) for chart_file in chart_files: @@ -228,9 +230,7 @@ def _extract_chart_title(chart_el: ET.Element, ns_c: str, ns_a: str) -> Optional return " ".join(parts).strip() # Try string reference: c:title/c:tx/c:strRef/c:strCache/c:pt/c:v - str_cache = title_el.find( - f"{{{ns_c}}}tx/{{{ns_c}}}strRef/{{{ns_c}}}strCache" - ) + str_cache = title_el.find(f"{{{ns_c}}}tx/{{{ns_c}}}strRef/{{{ns_c}}}strCache") if str_cache is not None: pt = str_cache.find(f"{{{ns_c}}}pt/{{{ns_c}}}v") if pt is not None and pt.text: @@ -362,7 +362,7 @@ def _extract_textboxes_from_zip(file_data: bytes) -> Dict[str, List[str]]: """ ns_xdr = OOXML_NS["xdr"] ns_a = OOXML_NS["a"] - ns_r = OOXML_NS["r"] + OOXML_NS["r"] textboxes: Dict[str, List[str]] = {} @@ -380,7 +380,9 @@ def _extract_textboxes_from_zip(file_data: bytes) -> Dict[str, List[str]]: if texts: textboxes[sheet_name] = texts except Exception as exc: - logger.debug("Failed to parse textboxes from %s: %s", drawing_path, exc) + logger.debug( + "Failed to parse textboxes from %s: %s", drawing_path, exc + ) except zipfile.BadZipFile: pass @@ -428,7 +430,9 @@ def _build_sheet_drawing_map(zf: zipfile.ZipFile) -> Dict[str, str]: continue # Resolve relative path: target is relative to xl/ - sheet_path = f"xl/{target}" if not target.startswith("/") else target.lstrip("/") + sheet_path = ( + f"xl/{target}" if not target.startswith("/") else target.lstrip("/") + ) rels_path = sheet_path.replace( os.path.basename(sheet_path), f"_rels/{os.path.basename(sheet_path)}.rels", @@ -443,7 +447,11 @@ def _build_sheet_drawing_map(zf: zipfile.ZipFile) -> Dict[str, str]: rel_target = rel.get("Target", "") rel_type = rel.get("Type", "") if "drawing" in rel_type.lower() and rel_target: - drawing_path = f"xl/{rel_target}" if not rel_target.startswith("/") else rel_target.lstrip("/") + drawing_path = ( + f"xl/{rel_target}" + if not rel_target.startswith("/") + else rel_target.lstrip("/") + ) # Normalize path (remove ../) drawing_path = os.path.normpath(drawing_path).replace("\\", "/") mapping[sheet_name] = drawing_path diff --git a/contextifier/integrations/langchain_loader.py b/contextifier/integrations/langchain_loader.py index 4fdbca4..fa2497a 100644 --- a/contextifier/integrations/langchain_loader.py +++ b/contextifier/integrations/langchain_loader.py @@ -30,7 +30,7 @@ import os from pathlib import Path -from typing import Any, Iterator, List, Optional, Union +from typing import Any, Iterator, Optional, Union from langchain_core.document_loaders import BaseLoader from langchain_core.documents import Document diff --git a/contextifier/ocr/base.py b/contextifier/ocr/base.py index 5665054..672c60a 100644 --- a/contextifier/ocr/base.py +++ b/contextifier/ocr/base.py @@ -104,11 +104,13 @@ def get_ocr_prompt(language: str = "ko") -> str: """ return _OCR_PROMPT_TEMPLATE.get(language, _OCR_PROMPT_TEMPLATE["en"]) + SIMPLE_OCR_PROMPT: str = "Describe the contents of this image." # ── Abstract Base ───────────────────────────────────────────────────────── + class BaseOCREngine(ABC): """ Abstract base class for OCR engine implementations. @@ -205,7 +207,9 @@ def convert_image_to_text(self, image_path: str) -> Optional[str]: response = self._llm_client.invoke([message]) result = response.content.strip() - logger.info(f"[{self.provider.upper()}] OCR completed: {os.path.basename(image_path)}") + logger.info( + f"[{self.provider.upper()}] OCR completed: {os.path.basename(image_path)}" + ) return f"[Figure:{result}]" except Exception as e: diff --git a/contextifier/ocr/processor.py b/contextifier/ocr/processor.py index 7009e15..284ead5 100644 --- a/contextifier/ocr/processor.py +++ b/contextifier/ocr/processor.py @@ -20,7 +20,7 @@ import os import re from dataclasses import dataclass -from typing import Any, Callable, Dict, List, Optional, Pattern, Protocol, Union +from typing import Any, Callable, List, Optional, Pattern, Protocol from contextifier.config import ProcessingConfig from contextifier.ocr.base import BaseOCREngine @@ -30,27 +30,28 @@ # ── Progress Protocol ───────────────────────────────────────────────────── + class OCRProgressCallback(Protocol): """Protocol for OCR progress reporting.""" - def __call__(self, event: OCRProgressEvent) -> None: - ... + def __call__(self, event: OCRProgressEvent) -> None: ... @dataclass(frozen=True) class OCRProgressEvent: """Structured progress event (replaces raw dict from old code).""" - event_type: str # 'tag_processing' | 'tag_processed' | 'completed' - current_index: int # 0-based index of current image - total_count: int # Total number of images - image_path: str = "" # Path of image being processed - status: str = "" # 'success' | 'failed' | '' - error: str = "" # Error message if failed + event_type: str # 'tag_processing' | 'tag_processed' | 'completed' + current_index: int # 0-based index of current image + total_count: int # Total number of images + image_path: str = "" # Path of image being processed + status: str = "" # 'success' | 'failed' | '' + error: str = "" # Error message if failed # ── OCRProcessor ────────────────────────────────────────────────────────── + class OCRProcessor: """ Orchestrates OCR replacement of image tags in text. @@ -140,12 +141,14 @@ def process( # Final notification if progress_callback: - progress_callback(OCRProgressEvent( - event_type="completed", - current_index=total, - total_count=total, - status=f"{success_count}/{total} succeeded", - )) + progress_callback( + OCRProgressEvent( + event_type="completed", + current_index=total, + total_count=total, + status=f"{success_count}/{total} succeeded", + ) + ) return result @@ -176,7 +179,11 @@ def _run_ocr_batch( ) as executor: future_to_idx = { executor.submit( - self._ocr_single, idx, img_path, total, progress_callback, + self._ocr_single, + idx, + img_path, + total, + progress_callback, ): idx for idx, img_path in enumerate(image_paths) } @@ -202,37 +209,46 @@ def _ocr_single( """ # Notify: processing started if progress_callback: - progress_callback(OCRProgressEvent( - event_type="tag_processing", - current_index=idx, - total_count=total, - image_path=img_path, - )) + progress_callback( + OCRProgressEvent( + event_type="tag_processing", + current_index=idx, + total_count=total, + image_path=img_path, + ) + ) # Resolve and convert local_path = self._resolve_image_path(img_path) if local_path is None: logger.warning(f"Image not found, keeping original tag: {img_path}") - self._notify_failed(progress_callback, idx, total, img_path, "File not found") + self._notify_failed( + progress_callback, idx, total, img_path, "File not found" + ) return (img_path, None) ocr_text = self._engine.convert_image_to_text(local_path) if ocr_text is None or ocr_text.startswith("[Image conversion error:"): logger.warning(f"OCR failed, keeping original tag: {img_path}") self._notify_failed( - progress_callback, idx, total, img_path, + progress_callback, + idx, + total, + img_path, ocr_text or "OCR returned None", ) return (img_path, None) if progress_callback: - progress_callback(OCRProgressEvent( - event_type="tag_processed", - current_index=idx, - total_count=total, - image_path=img_path, - status="success", - )) + progress_callback( + OCRProgressEvent( + event_type="tag_processed", + current_index=idx, + total_count=total, + image_path=img_path, + status="success", + ) + ) return (img_path, ocr_text) @@ -261,7 +277,10 @@ def _replace_tag(self, text: str, img_path: str, replacement: str) -> str: # Replace the capture group with the escaped literal path. # Use a lambda to avoid backslash interpretation in replacement strings. tag_pattern_str = re.sub( - r"\([^)]+\)", lambda _: escaped, pattern_str, count=1, + r"\([^)]+\)", + lambda _: escaped, + pattern_str, + count=1, ) tag_re = re.compile(tag_pattern_str) return tag_re.sub(lambda _: replacement, text) @@ -275,14 +294,16 @@ def _notify_failed( error: str, ) -> None: if callback: - callback(OCRProgressEvent( - event_type="tag_processed", - current_index=idx, - total_count=total, - image_path=path, - status="failed", - error=error, - )) + callback( + OCRProgressEvent( + event_type="tag_processed", + current_index=idx, + total_count=total, + image_path=path, + status="failed", + error=error, + ) + ) __all__ = [ diff --git a/contextifier/pipeline/__init__.py b/contextifier/pipeline/__init__.py index 9196f6d..64dc646 100644 --- a/contextifier/pipeline/__init__.py +++ b/contextifier/pipeline/__init__.py @@ -31,8 +31,14 @@ from contextifier.pipeline.converter import BaseConverter, NullConverter from contextifier.pipeline.preprocessor import BasePreprocessor, NullPreprocessor -from contextifier.pipeline.metadata_extractor import BaseMetadataExtractor, NullMetadataExtractor -from contextifier.pipeline.content_extractor import BaseContentExtractor, NullContentExtractor +from contextifier.pipeline.metadata_extractor import ( + BaseMetadataExtractor, + NullMetadataExtractor, +) +from contextifier.pipeline.content_extractor import ( + BaseContentExtractor, + NullContentExtractor, +) from contextifier.pipeline.postprocessor import BasePostprocessor, NullPostprocessor __all__ = [ diff --git a/contextifier/pipeline/content_extractor.py b/contextifier/pipeline/content_extractor.py index f62a1e1..b4a29f6 100644 --- a/contextifier/pipeline/content_extractor.py +++ b/contextifier/pipeline/content_extractor.py @@ -34,7 +34,7 @@ import logging from abc import ABC, abstractmethod -from typing import TYPE_CHECKING, Any, Dict, List, Optional +from typing import TYPE_CHECKING, Any, List, Optional from contextifier.types import ( ChartData, diff --git a/contextifier/pipeline/converter.py b/contextifier/pipeline/converter.py index de12d4a..74e067d 100644 --- a/contextifier/pipeline/converter.py +++ b/contextifier/pipeline/converter.py @@ -25,7 +25,7 @@ import logging import zipfile from abc import ABC, abstractmethod -from typing import Any, BinaryIO, Optional +from typing import Any from contextifier.types import FileContext from contextifier.errors import ConversionError @@ -159,4 +159,9 @@ def get_format_name(self) -> str: return "raw" -__all__ = ["BaseConverter", "NullConverter", "check_zip_bomb", "MAX_ZIP_DECOMPRESSED_BYTES"] +__all__ = [ + "BaseConverter", + "NullConverter", + "check_zip_bomb", + "MAX_ZIP_DECOMPRESSED_BYTES", +] diff --git a/contextifier/pipeline/metadata_extractor.py b/contextifier/pipeline/metadata_extractor.py index 6d54673..f3729f0 100644 --- a/contextifier/pipeline/metadata_extractor.py +++ b/contextifier/pipeline/metadata_extractor.py @@ -24,10 +24,9 @@ import logging from abc import ABC, abstractmethod -from typing import Any, Optional +from typing import Any from contextifier.types import DocumentMetadata -from contextifier.errors import ExtractionError class BaseMetadataExtractor(ABC): diff --git a/contextifier/pipeline/postprocessor.py b/contextifier/pipeline/postprocessor.py index 96a4e12..93cb476 100644 --- a/contextifier/pipeline/postprocessor.py +++ b/contextifier/pipeline/postprocessor.py @@ -43,8 +43,7 @@ from abc import ABC, abstractmethod from typing import TYPE_CHECKING, Any, Optional -from contextifier.types import DocumentMetadata, ExtractionResult -from contextifier.errors import PostprocessingError +from contextifier.types import ExtractionResult if TYPE_CHECKING: from contextifier.config import ProcessingConfig @@ -143,9 +142,7 @@ def postprocess( # 3. Append extraction warnings as HTML comments if result.warnings: - warning_lines = "\n".join( - f"" for w in result.warnings - ) + warning_lines = "\n".join(f"" for w in result.warnings) text = text + "\n\n" + warning_lines return text diff --git a/contextifier/pipeline/preprocessor.py b/contextifier/pipeline/preprocessor.py index 6be0fb0..0d11513 100644 --- a/contextifier/pipeline/preprocessor.py +++ b/contextifier/pipeline/preprocessor.py @@ -28,7 +28,6 @@ from typing import Any from contextifier.types import PreprocessedData -from contextifier.errors import PreprocessingError class BasePreprocessor(ABC): diff --git a/contextifier/raw/__init__.py b/contextifier/raw/__init__.py new file mode 100644 index 0000000..c8f21fb --- /dev/null +++ b/contextifier/raw/__init__.py @@ -0,0 +1,93 @@ +# contextifier/raw +""" +Raw document access — the lossless twin of the extraction pipeline. + +Contextifier has two ways to look at a document: + +* ``DocumentProcessor.extract_text()`` / ``.process()`` — the existing + pipeline that renders an **AI-friendly** view (clean text, normalized + tables/charts) and throws the rest away. +* ``open_raw()`` (this package) — a **lossless, addressable, writable** + view of the same file. Nothing is discarded: every OPC part stays + available, XML is parsed lazily, and edits are *surgical* — when you + save, untouched parts are written back **byte-identical** (the + byte-preservation contract), so charts, pivot tables, sparklines, + custom XML, styles and anything else the higher-level libraries can't + model all survive. + +Usage:: + + from contextifier import open_raw + + raw = open_raw("report.xlsx") # XlsxRawDocument + raw.sheets["Sales"].set_cell("B3", 142) + raw.charts[0].set_data(categories=["Q1", "Q2"], series=[("Sales", [1, 2])]) + raw.save("report-edited.xlsx") # or raw.to_bytes() + + raw = open_raw("deck.pptx") # PptxRawDocument + raw = open_raw("paper.docx") # DocxRawDocument + +Every format model also exposes ``.package`` (the raw +:class:`~contextifier.raw.opc.OpcPackage`) for part-level work, so the +"easy" interface never locks you out of the full container. + +Supported today: the OOXML trio (.xlsx / .docx / .pptx). Other handlers +raise :class:`RawUnsupportedError`. +""" + +from __future__ import annotations + +from contextifier.raw.opc import OpcPackage, OpcPart, RawUnsupportedError +from contextifier.raw.xmlpart import NS, XmlPart, qn + +__all__ = [ + "OpcPackage", + "OpcPart", + "RawUnsupportedError", + "XmlPart", + "NS", + "qn", + "open_raw", +] + + +def open_raw(source, *, extension: str | None = None): + """Open a document for lossless, writable access. + + Args: + source: path (str/Path), bytes, or a binary file object. + extension: override the format sniff (e.g. ``"xlsx"``); by default + the file extension (for paths) or the package content is used. + + Returns: + ``XlsxRawDocument`` / ``DocxRawDocument`` / ``PptxRawDocument``. + + Raises: + RawUnsupportedError: format has no raw model yet. + """ + from pathlib import Path + + ext = (extension or "").lower().lstrip(".") + if not ext and isinstance(source, (str, Path)): + ext = Path(source).suffix.lower().lstrip(".") + + package = OpcPackage.open(source) + if not ext: + ext = package.sniff_format() or "" + + if ext == "xlsx": + from contextifier.raw.xlsx import XlsxRawDocument + + return XlsxRawDocument(package) + if ext == "docx": + from contextifier.raw.docx import DocxRawDocument + + return DocxRawDocument(package) + if ext == "pptx": + from contextifier.raw.pptx import PptxRawDocument + + return PptxRawDocument(package) + raise RawUnsupportedError( + f"No raw model for {ext or 'unknown format'!r} yet " + "(supported: xlsx, docx, pptx)" + ) diff --git a/contextifier/raw/base.py b/contextifier/raw/base.py new file mode 100644 index 0000000..fdb7abd --- /dev/null +++ b/contextifier/raw/base.py @@ -0,0 +1,69 @@ +# contextifier/raw/base.py +""" +Shared base for format raw-document models (xlsx / docx / pptx). + +A format model owns a set of :class:`~contextifier.raw.xmlpart.XmlPart` +facades over the parts it understands. ``save()`` flushes every dirty +facade into the package, then serializes the package under the +byte-preservation contract. Parts the model does NOT understand are +never touched at all. +""" + +from __future__ import annotations + +from pathlib import Path +from typing import BinaryIO + +from contextifier.raw.opc import OpcPackage, OpcPart +from contextifier.raw.xmlpart import XmlPart + +__all__ = ["RawDocumentBase"] + + +class RawDocumentBase: + """Common plumbing: part registry, flush-on-save, byte export.""" + + #: subclasses set this ("xlsx" / "docx" / "pptx") + format: str = "" + + def __init__(self, package: OpcPackage): + self.package = package + self._xml_parts: dict[str, XmlPart] = {} + + # -- part facades ------------------------------------------------------------ + + def xml_part(self, name: str) -> XmlPart: + """The (cached) XmlPart facade for a package part.""" + if name not in self._xml_parts: + self._xml_parts[name] = XmlPart(self.package.get_part(name)) + return self._xml_parts[name] + + def raw_part(self, name: str) -> OpcPart: + """Direct part access — the escape hatch for anything the model + doesn't cover.""" + return self.package.get_part(name) + + # -- persistence ---------------------------------------------------------- + + def flush(self) -> None: + """Serialize every dirty XML facade into its package part.""" + for xp in self._xml_parts.values(): + xp.flush() + + def to_bytes(self) -> bytes: + self.flush() + return self.package.to_bytes() + + def save(self, target: str | Path | BinaryIO | None = None) -> bytes: + """Write the package; untouched parts stay byte-identical.""" + self.flush() + return self.package.save(target) + + def close(self) -> None: + self.package.close() + + def __enter__(self): + return self + + def __exit__(self, *exc) -> None: + self.close() diff --git a/contextifier/raw/chart.py b/contextifier/raw/chart.py new file mode 100644 index 0000000..9fb50ab --- /dev/null +++ b/contextifier/raw/chart.py @@ -0,0 +1,771 @@ +# contextifier/raw/chart.py +""" +ChartModel — read & write DrawingML charts, shared across all three +OOXML formats. + +A chart is a ``chartN.xml`` part (classic ``c:`` namespace, or ``cx:`` +chartEx for the 2016+ types) referenced from a sheet drawing (xlsx), a +document drawing (docx) or a slide graphicFrame (pptx), usually paired +with an embedded ``.xlsx`` workbook holding the source table. This +module is format-agnostic: it only needs the chart part and the +package, so the three format models can share one implementation. + +Reading:: + + chart = raw.charts[0] + chart.kind # "bar" | "line" | "pie" | ... | "chartex:" + chart.title # str | None + chart.series # [ChartSeriesData(name, categories, values), ...] + +Writing:: + + chart.set_title("Q3 Sales") + chart.set_data( + categories=["Q1", "Q2", "Q3"], + series=[("Sales", [120, 135, 150]), ("Cost", [80, 90, 95])], + ) + raw.save("out.xlsx") + +``set_data`` must rewrite BOTH the caches inside the chart XML +(``c:cat/c:strRef/c:strCache``, ``c:val/c:numRef/c:numCache`` per +series — creating/removing ``c:ser`` elements as the series count +changes) AND the embedded workbook part (if present) so that +double-click-edit in Office shows the same numbers. Formula references +(``c:f``) should be regenerated against the embedded workbook's sheet +("Sheet1!$B$2:$B$4" style). For xlsx-hosted charts whose series +reference the HOST workbook's own cells, ``set_data`` rewrites caches +and leaves the ``c:f`` references pointing at the host sheet (values +there are the caller's responsibility — typically edited through +``sheet.set_cell`` alongside). + +Embedded workbook regeneration +------------------------------ +When the chart has an embedded ``.../embeddings/*.xlsx`` workbook (the +usual case for pptx/docx-hosted charts), ``set_data`` REGENERATES that +part from scratch with openpyxl: a single sheet holding the plain data +table (header row ``[None, name1, name2, ...]`` followed by one +``[category, v1, v2, ...]`` row per category). The embedded workbook is +chart-internal source data, not user content, so replacing it wholesale +is safe — it is exactly what keeps Office's "Edit Data" view in sync +with the rewritten caches. + +chartEx (cx:) limitations in this milestone +------------------------------------------- +chartEx charts are fully **readable** (``kind`` / ``title`` / +``series``), but ``set_title`` and ``set_data`` raise +:class:`~contextifier.raw.opc.RawUnsupportedError` for them — the cx: +write path is deferred (documented limitation, see the class contract). +""" + +from __future__ import annotations + +import copy +import io +import posixpath +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Sequence + +from lxml import etree + +from contextifier.raw.opc import RawUnsupportedError +from contextifier.raw.xmlpart import NS, XmlPart, qn + +if TYPE_CHECKING: # pragma: no cover + from lxml.etree import _Element + + from contextifier.raw.opc import OpcPackage, OpcPart + +__all__ = ["ChartSeriesData", "ChartModel", "find_chart_parts", "load_chart"] + + +@dataclass +class ChartSeriesData: + """One series as read from (or written to) the chart caches.""" + + name: str | None + categories: list[str] = field(default_factory=list) + values: list[float | None] = field(default_factory=list) + + +# Classic plot-type element (local name) → kind. bar/bar3D are handled +# separately because they split on c:barDir ("bar" vs "column"). +_PLOT_KIND: dict[str, str] = { + "lineChart": "line", + "line3DChart": "line", + "pieChart": "pie", + "pie3DChart": "pie", + "areaChart": "area", + "area3DChart": "area", + "scatterChart": "scatter", + "doughnutChart": "doughnut", + "radarChart": "radar", + "bubbleChart": "bubble", + "ofPieChart": "of_pie", + "surfaceChart": "surface", + "surface3DChart": "surface", + "stockChart": "stock", +} + +# Plot types whose series carry c:xVal/c:yVal instead of c:cat/c:val. +_XY_PLOTS = {"scatterChart", "bubbleChart"} + +# c:ser children that must stay AFTER c:cat/c:val (schema tail). +_SER_TAIL_TAGS = ("c:smooth", "c:shape", "c:bubbleSize", "c:bubble3D", "c:extLst") + + +def _local(el: "_Element") -> str: + """Local (namespace-stripped) tag name; '' for comments/PIs.""" + tag = el.tag + if not isinstance(tag, str): + return "" + return tag.rsplit("}", 1)[-1] + + +def _col_letter(n: int) -> str: + """1-based column index → spreadsheet letters (1→A, 27→AA).""" + letters = "" + while n > 0: + n, rem = divmod(n - 1, 26) + letters = chr(ord("A") + rem) + letters + return letters + + +def _num_text(value: float) -> str: + """Float/int → the shortest cache text Excel reads back exactly.""" + f = float(value) + if f.is_integer(): + return str(int(f)) + return repr(f) + + +def _unquote_sheet(ref_prefix: str) -> str: + """``"'My Sheet'"`` → ``"My Sheet"`` (c:f quoting → real sheet name).""" + if ref_prefix.startswith("'") and ref_prefix.endswith("'"): + return ref_prefix[1:-1].replace("''", "'") + return ref_prefix + + +class ChartModel: + """Read/write view over one chart part. + + Contract (implemented in this module; consumed by the xlsx/docx/pptx + models — do not change signatures without updating all three): + + * ``ChartModel(xml_part, package)`` — *xml_part* is the chart part + facade; *package* is the owning :class:`OpcPackage` (used to reach + the embedded workbook through the chart part's rels). + * ``kind: str`` (property) — first plot type found: ``"bar"``, + ``"line"``, ``"pie"``, ``"area"``, ``"scatter"``, ``"doughnut"``, + ``"radar"``, ``"bubble"``, ... or ``"chartex:funnel"`` etc. for + cx: charts. + * ``title: str | None`` (property) — concatenated ``a:t`` runs of the + chart title, if any. + * ``series: list[ChartSeriesData]`` (property) — parsed from + str/num caches (classic) or ``cx:strDim``/``cx:numDim`` (chartEx). + Missing cache points yield ``None`` values. + * ``set_title(text: str) -> None`` — replace/insert the title text, + preserving existing run formatting where present. + * ``set_data(categories, series) -> None`` — see module docstring. + ``series`` accepts ``[(name, values), ...]`` or + ``[ChartSeriesData, ...]``. Raises ``ValueError`` on ragged input + (series length != len(categories)). chartEx write support may + raise ``RawUnsupportedError`` (documented limitation) in v0.4. + * ``embedded_workbook_part() -> OpcPart | None`` — the + ``.../embeddings/*.xlsx`` part referenced by this chart's rels. + + Implementation notes (behavior, not signatures): + + * Parsing is lazy — ``__init__`` stores the facade; the tree is + first parsed on the first property access. + * Mutators (``set_title`` / ``set_data``) mark the XML facade dirty + AND flush it immediately, so ``package.to_bytes()`` alone reflects + the edit; format-model ``save()`` flushing stays a no-op for it. + * chartEx: read-only in this milestone — ``set_title`` and + ``set_data`` raise :class:`RawUnsupportedError` for cx: charts. + * Charts without value caches (e.g. openpyxl-authored charts, which + write only ``c:f`` references into the host workbook) read as + series with ``name=None``, empty categories and ``values`` of + ``None`` per cached ``ptCount`` (or empty when no cache exists). + """ + + def __init__(self, xml_part: "XmlPart", package: "OpcPackage"): + self.xml = xml_part + self.package = package + + # -- detection --------------------------------------------------------------- + + @property + def _is_chartex(self) -> bool: + return self.xml.root.tag == qn("cx:chartSpace") + + def _plot_elements(self) -> list["_Element"]: + """Plot-type children of c:plotArea (barChart, lineChart, ...).""" + plot_area = self.xml.root.find("c:chart/c:plotArea", NS) + if plot_area is None: + return [] + return [el for el in plot_area if _local(el).endswith("Chart")] + + # -- reading ----------------------------------------------------------------- + + @property + def kind(self) -> str: + if self._is_chartex: + ser = self.xml.root.find( + "cx:chart/cx:plotArea/cx:plotAreaRegion/cx:series", NS + ) + if ser is None: + ser = next(self.xml.root.iter(qn("cx:series")), None) + layout = ser.get("layoutId") if ser is not None else None + return f"chartex:{layout or 'unknown'}" + for plot in self._plot_elements(): + name = _local(plot) + if name in ("barChart", "bar3DChart"): + bar_dir = plot.find("c:barDir", NS) + val = bar_dir.get("val", "col") if bar_dir is not None else "col" + return "bar" if val == "bar" else "column" + return _PLOT_KIND.get(name, name) + return "unknown" + + @property + def title(self) -> str | None: + root = self.xml.root + if self._is_chartex: + title_el = root.find("cx:chart/cx:title", NS) + else: + chart = root.find("c:chart", NS) + if chart is None: + return None + atd = chart.find("c:autoTitleDeleted", NS) + if atd is not None and atd.get("val", "1") in ("1", "true"): + return None + title_el = chart.find("c:title", NS) + if title_el is None: + return None + text = "".join(t.text for t in title_el.iter(qn("a:t")) if t.text) + return text or None + + @property + def series(self) -> list[ChartSeriesData]: + if self._is_chartex: + return self._series_chartex() + out: list[ChartSeriesData] = [] + for plot in self._plot_elements(): + uses_xy = _local(plot) in _XY_PLOTS + for ser in plot.findall("c:ser", NS): + out.append(self._read_ser(ser, uses_xy)) + return out + + def _read_ser(self, ser: "_Element", uses_xy: bool) -> ChartSeriesData: + name = None + name_el = ser.find("c:tx/c:strRef/c:strCache/c:pt/c:v", NS) + if name_el is None: + name_el = ser.find("c:tx/c:v", NS) + if name_el is not None and name_el.text: + name = name_el.text + cat_holder = ser.find("c:xVal" if uses_xy else "c:cat", NS) + val_holder = ser.find("c:yVal" if uses_xy else "c:val", NS) + return ChartSeriesData( + name=name, + categories=self._cache_texts(cat_holder), + values=self._cache_values(val_holder), + ) + + @staticmethod + def _find_cache(holder: "_Element | None") -> "_Element | None": + """The cache element inside a c:cat/c:val/c:xVal/c:yVal holder.""" + if holder is None: + return None + for path in ( + "c:strRef/c:strCache", + "c:numRef/c:numCache", + "c:strLit", + "c:numLit", + # multi-level categories: first c:lvl = leaf labels + "c:multiLvlStrRef/c:multiLvlStrCache/c:lvl", + ): + cache = holder.find(path, NS) + if cache is not None: + return cache + return None + + @classmethod + def _cache_texts(cls, holder: "_Element | None") -> list[str]: + """Cache points as strings in idx order; gaps are skipped.""" + cache = cls._find_cache(holder) + if cache is None: + return [] + pts: list[tuple[int, str]] = [] + for pt in cache.findall("c:pt", NS): + v = pt.find("c:v", NS) + if v is None: + continue + try: + idx = int(pt.get("idx", str(len(pts)))) + except ValueError: + idx = len(pts) + pts.append((idx, v.text or "")) + return [text for _, text in sorted(pts, key=lambda p: p[0])] + + @classmethod + def _cache_values(cls, holder: "_Element | None") -> list[float | None]: + """Numeric cache; length = ptCount, gaps/unparseable → None.""" + cache = cls._find_cache(holder) + if cache is None: + return [] + by_idx: dict[int, float | None] = {} + for pt in cache.findall("c:pt", NS): + v = pt.find("c:v", NS) + try: + idx = int(pt.get("idx", "0")) + except ValueError: + continue + if v is None or v.text is None: + by_idx[idx] = None + continue + try: + by_idx[idx] = float(v.text) + except ValueError: + by_idx[idx] = None + pt_count_el = cache.find("c:ptCount", NS) + try: + count = int(pt_count_el.get("val", "")) if pt_count_el is not None else -1 + except ValueError: + count = -1 + if count < 0: + count = (max(by_idx) + 1) if by_idx else 0 + return [by_idx.get(i) for i in range(count)] + + def _series_chartex(self) -> list[ChartSeriesData]: + root = self.xml.root + data_by_id = { + data.get("id"): data for data in root.findall("cx:chartData/cx:data", NS) + } + series_els = root.findall( + "cx:chart/cx:plotArea/cx:plotAreaRegion/cx:series", NS + ) or list(root.iter(qn("cx:series"))) + out: list[ChartSeriesData] = [] + if not series_els: + # No series elements: expose each data block as an unnamed series. + for data in data_by_id.values(): + cats, vals = self._chartex_dims(data) + out.append(ChartSeriesData(name=None, categories=cats, values=vals)) + return out + for ser in series_els: + name_el = ser.find("cx:tx/cx:txData/cx:v", NS) + name = name_el.text if name_el is not None and name_el.text else None + data_id_el = ser.find("cx:dataId", NS) + data = ( + data_by_id.get(data_id_el.get("val")) + if data_id_el is not None + else None + ) + cats: list[str] = [] + vals: list[float | None] = [] + if data is not None: + cats, vals = self._chartex_dims(data) + out.append(ChartSeriesData(name=name, categories=cats, values=vals)) + return out + + @staticmethod + def _chartex_dims(data: "_Element") -> tuple[list[str], list[float | None]]: + """(categories, values) from a cx:data block's dims.""" + + def pick(dims: list["_Element"], wanted_type: str) -> "_Element | None": + for dim in dims: + if dim.get("type") == wanted_type: + return dim + return dims[0] if dims else None + + cats: list[str] = [] + vals: list[float | None] = [] + str_dim = pick(data.findall("cx:strDim", NS), "cat") + if str_dim is not None: + lvl = str_dim.find("cx:lvl", NS) + if lvl is not None: + pts = [] + for pt in lvl.findall("cx:pt", NS): + try: + idx = int(pt.get("idx", str(len(pts)))) + except ValueError: + idx = len(pts) + pts.append((idx, pt.text or "")) + cats = [text for _, text in sorted(pts, key=lambda p: p[0])] + num_dim = pick(data.findall("cx:numDim", NS), "val") + if num_dim is not None: + lvl = num_dim.find("cx:lvl", NS) + if lvl is not None: + by_idx: dict[int, float | None] = {} + for pt in lvl.findall("cx:pt", NS): + try: + idx = int(pt.get("idx", "0")) + except ValueError: + continue + try: + by_idx[idx] = float(pt.text) if pt.text else None + except ValueError: + by_idx[idx] = None + try: + count = int(lvl.get("ptCount", "")) + except ValueError: + count = (max(by_idx) + 1) if by_idx else 0 + vals = [by_idx.get(i) for i in range(count)] + return cats, vals + + # -- writing: title ---------------------------------------------------------- + + def set_title(self, text: str) -> None: + """Replace (or insert) the chart title text. + + Keeps the first existing run's ``a:rPr`` formatting; extra runs + and extra paragraphs of the old title are removed so the title + reads exactly *text*. Removes ``c:autoTitleDeleted val="1"`` so + the explicit title is shown. chartEx charts are read-only in + this milestone → :class:`RawUnsupportedError`. + """ + if self._is_chartex: + raise RawUnsupportedError( + "set_title is not supported for chartEx (cx:) charts yet; " + "classic c: charts only" + ) + chart = self.xml.root.find("c:chart", NS) + if chart is None: + raise RawUnsupportedError("chart part has no c:chart element") + + atd = chart.find("c:autoTitleDeleted", NS) + if atd is not None and atd.get("val", "1") in ("1", "true"): + chart.remove(atd) + + title_el = chart.find("c:title", NS) + if title_el is None: + title_el = etree.Element(qn("c:title")) + chart.insert(0, title_el) # c:title is the first child per schema + tx = title_el.find("c:tx", NS) + if tx is None: + tx = etree.Element(qn("c:tx")) + title_el.insert(0, tx) # c:tx is the first child of c:title + rich = tx.find("c:rich", NS) + if rich is None: + # Replace whatever c:tx held (e.g. a c:strRef) with rich text. + for child in list(tx): + tx.remove(child) + rich = etree.SubElement(tx, qn("c:rich")) + etree.SubElement(rich, qn("a:bodyPr")) + etree.SubElement(rich, qn("a:lstStyle")) + + paragraphs = rich.findall("a:p", NS) + if paragraphs: + para = paragraphs[0] + for extra in paragraphs[1:]: + rich.remove(extra) + else: + para = etree.SubElement(rich, qn("a:p")) + + run_props = None + first_run = para.find("a:r", NS) + if first_run is not None: + run_props = first_run.find("a:rPr", NS) + for el in list(para): + if el.tag in (qn("a:r"), qn("a:fld"), qn("a:br")): + para.remove(el) + run = etree.Element(qn("a:r")) + if run_props is not None: + run.append(copy.deepcopy(run_props)) + t = etree.SubElement(run, qn("a:t")) + t.text = text + end_props = para.find("a:endParaRPr", NS) + if end_props is not None: + end_props.addprevious(run) + else: + para.append(run) + self._commit() + + # -- writing: data ----------------------------------------------------------- + + def set_data( + self, + categories: Sequence[object], + series: Sequence["tuple[str | None, Sequence[float | None]] | ChartSeriesData"], + ) -> None: + """Rewrite the chart's cached data (and its embedded workbook). + + *series* accepts ``[(name, values), ...]`` or + ``[ChartSeriesData, ...]`` (whose ``.categories`` are ignored — + the *categories* argument is canonical for every series). + + Raises ``ValueError`` on ragged input and + :class:`RawUnsupportedError` for chartEx (cx:) charts. + """ + if self._is_chartex: + raise RawUnsupportedError( + "set_data is not supported for chartEx (cx:) charts in this " + "version; classic c: charts only" + ) + cats = list(categories) + normalized: list[tuple[str | None, list[float | None]]] = [] + for item in series: + if isinstance(item, ChartSeriesData): + normalized.append((item.name, list(item.values))) + else: + name, values = item + normalized.append((name, list(values))) + if not normalized: + raise ValueError("set_data requires at least one series") + for name, values in normalized: + if len(values) != len(cats): + raise ValueError( + f"series {name!r} has {len(values)} values for " + f"{len(cats)} categories (lengths must match)" + ) + + sers: list[tuple["_Element", "_Element"]] = [ + (plot, ser) + for plot in self._plot_elements() + for ser in plot.findall("c:ser", NS) + ] + if not sers: + raise RawUnsupportedError( + "chart has no existing c:ser element to rewrite or clone from" + ) + + sheet_ref = self._sheet_ref() + + # Shrink: drop trailing series. + while len(sers) > len(normalized): + plot, ser = sers.pop() + plot.remove(ser) + # Grow: clone the LAST series (deep copy keeps its styling). + if len(sers) < len(normalized): + used_indices = [] + for _, ser in sers: + for tag in ("c:idx", "c:order"): + el = ser.find(tag, NS) + if el is not None: + try: + used_indices.append(int(el.get("val", ""))) + except ValueError: + pass + next_index = (max(used_indices) + 1) if used_indices else len(sers) + plot, last = sers[-1] + while len(sers) < len(normalized): + clone = copy.deepcopy(last) + self._set_ser_index(clone, next_index) + last.addnext(clone) + sers.append((plot, clone)) + last = clone + next_index += 1 + + n = len(cats) + for i, ((plot, ser), (name, values)) in enumerate(zip(sers, normalized)): + col = _col_letter(i + 2) # data columns start at B + self._rewrite_ser( + ser, + name, + cats, + values, + uses_xy=_local(plot) in _XY_PLOTS, + cat_ref=f"{sheet_ref}!$A$2:$A${n + 1}", + val_ref=f"{sheet_ref}!${col}$2:${col}${n + 1}", + name_ref=f"{sheet_ref}!${col}$1", + ) + + workbook_part = self.embedded_workbook_part() + if workbook_part is not None: + workbook_part.write( + self._workbook_bytes(_unquote_sheet(sheet_ref), cats, normalized) + ) + self._commit() + + def _sheet_ref(self) -> str: + """Sheet prefix (verbatim, quotes kept) of the first c:f found.""" + for f in self.xml.root.iter(qn("c:f")): + text = f.text or "" + if "!" in text: + return text.rsplit("!", 1)[0] + return "Sheet1" + + @staticmethod + def _set_ser_index(ser: "_Element", index: int) -> None: + for tag in ("c:idx", "c:order"): + el = ser.find(tag, NS) + if el is None: + el = etree.Element(qn(tag)) + ser.insert(0 if tag == "c:idx" else 1, el) + el.set("val", str(index)) + + def _rewrite_ser( + self, + ser: "_Element", + name: str | None, + cats: list[object], + values: list[float | None], + *, + uses_xy: bool, + cat_ref: str, + val_ref: str, + name_ref: str, + ) -> None: + if name is not None: + self._rewrite_ser_name(ser, name, name_ref) + cat_tag = "c:xVal" if uses_xy else "c:cat" + val_tag = "c:yVal" if uses_xy else "c:val" + numeric_cats = all( + isinstance(c, (int, float)) and not isinstance(c, bool) for c in cats + ) + if uses_xy and numeric_cats: + new_cat = self._build_num_holder(cat_tag, cat_ref, cats) # type: ignore[arg-type] + else: + new_cat = self._build_str_holder(cat_tag, cat_ref, cats) + new_val = self._build_num_holder(val_tag, val_ref, values) + self._replace_ser_child(ser, cat_tag, new_cat, also_before=(val_tag,)) + self._replace_ser_child(ser, val_tag, new_val) + + def _rewrite_ser_name(self, ser: "_Element", name: str, name_ref: str) -> None: + tx = ser.find("c:tx", NS) + if tx is None: + tx = etree.Element(qn("c:tx")) + anchor = ser.find("c:order", NS) + if anchor is None: + anchor = ser.find("c:idx", NS) + if anchor is not None: + anchor.addnext(tx) + else: + ser.insert(0, tx) + literal = tx.find("c:v", NS) + if literal is not None and tx.find("c:strRef", NS) is None: + literal.text = name + return + for child in list(tx): + tx.remove(child) + str_ref = etree.SubElement(tx, qn("c:strRef")) + f = etree.SubElement(str_ref, qn("c:f")) + f.text = name_ref + cache = etree.SubElement(str_ref, qn("c:strCache")) + pt_count = etree.SubElement(cache, qn("c:ptCount")) + pt_count.set("val", "1") + pt = etree.SubElement(cache, qn("c:pt")) + pt.set("idx", "0") + v = etree.SubElement(pt, qn("c:v")) + v.text = name + + @staticmethod + def _build_str_holder(tag: str, ref: str, items: list[object]) -> "_Element": + holder = etree.Element(qn(tag)) + str_ref = etree.SubElement(holder, qn("c:strRef")) + f = etree.SubElement(str_ref, qn("c:f")) + f.text = ref + cache = etree.SubElement(str_ref, qn("c:strCache")) + pt_count = etree.SubElement(cache, qn("c:ptCount")) + pt_count.set("val", str(len(items))) + for idx, item in enumerate(items): + if item is None: + continue # gap + pt = etree.SubElement(cache, qn("c:pt")) + pt.set("idx", str(idx)) + v = etree.SubElement(pt, qn("c:v")) + v.text = str(item) + return holder + + @staticmethod + def _build_num_holder(tag: str, ref: str, values: list[float | None]) -> "_Element": + holder = etree.Element(qn(tag)) + num_ref = etree.SubElement(holder, qn("c:numRef")) + f = etree.SubElement(num_ref, qn("c:f")) + f.text = ref + cache = etree.SubElement(num_ref, qn("c:numCache")) + format_code = etree.SubElement(cache, qn("c:formatCode")) + format_code.text = "General" + pt_count = etree.SubElement(cache, qn("c:ptCount")) + pt_count.set("val", str(len(values))) + for idx, value in enumerate(values): + if value is None: + continue # gap + pt = etree.SubElement(cache, qn("c:pt")) + pt.set("idx", str(idx)) + v = etree.SubElement(pt, qn("c:v")) + v.text = _num_text(value) + return holder + + @staticmethod + def _replace_ser_child( + ser: "_Element", + tag: str, + new_el: "_Element", + also_before: tuple[str, ...] = (), + ) -> None: + """Swap the *tag* child in place, or insert it at a schema-valid + position (before the c:ser tail elements).""" + existing = ser.find(tag, NS) + if existing is not None: + existing.addprevious(new_el) + ser.remove(existing) + return + stop_tags = {qn(t) for t in (*_SER_TAIL_TAGS, *also_before)} + for child in ser: + if child.tag in stop_tags: + child.addprevious(new_el) + return + ser.append(new_el) + + # -- embedded workbook --------------------------------------------------------- + + def embedded_workbook_part(self) -> "OpcPart | None": + """The ``.../embeddings/*.xlsx`` part behind this chart, if any.""" + rels = self.package.rels_for(self.xml.name) + if rels is None: + return None + for rel in rels: + if rel["mode"] == "External": + continue + rel_type = rel["type"] or "" + if not (rel_type.endswith("/package") or rel_type.endswith("/oleObject")): + continue + target = rel["target"] or "" + if posixpath.splitext(target)[1].lower() != ".xlsx": + continue + name = rels.resolve(self.xml.name, target) + if self.package.has_part(name): + return self.package.get_part(name) + return None + + @staticmethod + def _workbook_bytes( + sheet_name: str, + cats: list[object], + normalized: list[tuple[str | None, list[float | None]]], + ) -> bytes: + """A fresh data-only workbook mirroring the rewritten caches.""" + from openpyxl import Workbook + + wb = Workbook() + ws = wb.active + ws.title = sheet_name + ws.append([None, *(name for name, _ in normalized)]) + for i, cat in enumerate(cats): + ws.append([cat, *(values[i] for _, values in normalized)]) + buf = io.BytesIO() + wb.save(buf) + return buf.getvalue() + + # -- persistence ----------------------------------------------------------- + + def _commit(self) -> None: + """Mark the facade dirty and serialize into the package part.""" + self.xml.mark_dirty() + self.xml.flush() + + +def load_chart(package: "OpcPackage", chart_part_name: str) -> ChartModel: + """Convenience: chart part name → :class:`ChartModel` facade.""" + return ChartModel(XmlPart(package.get_part(chart_part_name)), package) + + +def find_chart_parts(package: "OpcPackage", from_part: str) -> list[str]: + """Chart part names referenced (directly) by *from_part*'s rels. + + Format models use this from a drawing/slide part: + ``rels.by_type("/chart")`` → resolve targets → chart part names. + """ + rels = package.rels_for(from_part) + if rels is None: + return [] + return [rels.resolve(from_part, rel["target"]) for rel in rels.by_type("/chart")] diff --git a/contextifier/raw/docx.py b/contextifier/raw/docx.py new file mode 100644 index 0000000..befb7cb --- /dev/null +++ b/contextifier/raw/docx.py @@ -0,0 +1,567 @@ +# contextifier/raw/docx.py +""" +DocxRawDocument — the raw (lossless, addressable, writable) model for +WordprocessingML documents. + +Addressing invariant +-------------------- +``paragraphs[i]`` is the *i*-th direct ``w:p`` child of ``w:body`` and +``tables[t]`` is the *t*-th direct ``w:tbl`` child — exactly the +elements python-docx exposes as ``document.paragraphs`` / +``document.tables``. Table cells use **grid addressing** with the same +semantics as python-docx ``row.cells``: a horizontally merged cell +(``w:gridSpan``) occupies every grid column it spans, and a vertically +merged continuation (``w:vMerge``) resolves to its start cell. This +keeps raw addresses interchangeable with edit2docs' existing addresses. + +Run-preserving text replacement (the P0-3 fix) +---------------------------------------------- +``set_paragraph_text`` never rebuilds a paragraph. Within the +paragraph, runs are classified: + +* **text runs** — ``w:r`` containing ``w:t`` and *no* protected content + (``w:drawing`` / ``w:pict`` / ``w:object`` / ``mc:AlternateContent``); +* **protected elements** — runs holding drawings/pictures/OLE, plus + every non-run child (bookmarks, math, comment ranges, revision + containers, ...). These are left in place, in order, untouched. + +The new text goes into the **first direct text run** (its ``w:rPr`` is +kept; ``w:t`` gets ``xml:space="preserve"`` when the text has leading +or trailing whitespace); the other direct pure-text runs are deleted. + +Hyperlink policy: ``w:hyperlink`` elements are never deleted by a text +replace. If the paragraph has no direct text run, the new text goes +into the *first hyperlink's first text run* (the hyperlink itself is +preserved). All other pure-text runs inside hyperlinks are emptied, but +the (now empty) hyperlink elements remain — callers that want them gone +use :meth:`DocxRawDocument.strip_empty_hyperlinks`, which also drops +the hyperlink relationship when nothing else references it. + +``RawCell.set_text`` applies the same rules per paragraph: the first +paragraph that owns a text run becomes the carrier, paragraphs with +protected content (inline images, ...) are left completely untouched, +and the remaining pure-text paragraphs are emptied — never removed, so +cell layout, nested tables and images survive (the edit2docs +cell-image-destruction fix). +""" + +from __future__ import annotations + +import copy +from typing import TYPE_CHECKING + +from contextifier.raw.base import RawDocumentBase +from contextifier.raw.chart import ChartModel, find_chart_parts +from contextifier.raw.xmlpart import NS, qn + +if TYPE_CHECKING: # pragma: no cover + from lxml.etree import _Element + + from contextifier.raw.opc import OpcPackage + +__all__ = ["DocxRawDocument", "RawParagraph", "RawTable", "RawCell"] + +_DOCUMENT_PART = "word/document.xml" + +#: Content that makes a run (or a paragraph) non-replaceable. +_PROTECTED_TAGS = ( + qn("w:drawing"), + qn("w:pict"), + qn("w:object"), + qn("mc:AlternateContent"), +) +_TEXTUAL_TAGS = (qn("w:t"), qn("w:tab"), qn("w:br"), qn("w:cr")) + + +# -- element-level helpers --------------------------------------------------- + + +def _contains_protected(el: "_Element") -> bool: + return next(el.iter(*_PROTECTED_TAGS), None) is not None + + +def _is_text_run(el: "_Element") -> bool: + """A pure-text run: ``w:r`` with a ``w:t`` and no protected content.""" + return ( + el.tag == qn("w:r") + and next(el.iter(qn("w:t")), None) is not None + and not _contains_protected(el) + ) + + +def _has_del_ancestor(node: "_Element", stop: "_Element") -> bool: + parent = node.getparent() + while parent is not None and parent is not stop: + if parent.tag == qn("w:del"): + return True + parent = parent.getparent() + return False + + +def _element_text(el: "_Element") -> str: + """Visible text of *el* (python-docx ``.text`` semantics). + + Concatenates ``w:t`` anywhere below *el* — including inside + ``w:hyperlink`` and ``w:ins`` — excluding deleted (``w:del``) + content; run-level ``w:tab``/``w:br``/``w:cr`` render as + ``"\\t"``/``"\\n"``. + """ + parts: list[str] = [] + for node in el.iter(*_TEXTUAL_TAGS): + if _has_del_ancestor(node, el): + continue + if node.tag == qn("w:t"): + parts.append(node.text or "") + elif node.getparent() is not None and node.getparent().tag == qn("w:r"): + parts.append("\t" if node.tag == qn("w:tab") else "\n") + return "".join(parts) + + +def _set_run_text(run: "_Element", text: str) -> None: + """Replace the run's content with a single ``w:t``, keeping ``w:rPr``.""" + from lxml import etree + + for child in list(run): + if child.tag != qn("w:rPr"): + run.remove(child) + t = etree.SubElement(run, qn("w:t")) + t.text = text + if text != text.strip(): + t.set("{http://www.w3.org/XML/1998/namespace}space", "preserve") + + +def _clear_run_text(run: "_Element") -> None: + """Empty a pure-text run (drop everything but ``w:rPr``).""" + for child in list(run): + if child.tag != qn("w:rPr"): + run.remove(child) + + +def _paragraph_text_runs(p_el: "_Element") -> tuple[list, list[tuple]]: + """(direct pure-text runs, [(hyperlink, its pure-text runs), ...]).""" + direct = [r for r in p_el.findall(qn("w:r")) if _is_text_run(r)] + in_links = [ + (h, [r for r in h.findall(qn("w:r")) if _is_text_run(r)]) + for h in p_el.findall(qn("w:hyperlink")) + ] + return direct, in_links + + +def _has_text_run(p_el: "_Element") -> bool: + direct, in_links = _paragraph_text_runs(p_el) + return bool(direct) or any(runs for _, runs in in_links) + + +def _replace_paragraph_text(p_el: "_Element", text: str) -> None: + """The run-preserving replace documented in the module docstring.""" + from lxml import etree + + direct, in_links = _paragraph_text_runs(p_el) + + carrier = direct[0] if direct else None + if carrier is None: + for _, runs in in_links: + if runs: + carrier = runs[0] + break + if carrier is None: # no text run anywhere — append a fresh one + carrier = etree.SubElement(p_el, qn("w:r")) + + _set_run_text(carrier, text) + for run in direct: + if run is not carrier: + p_el.remove(run) + for _, runs in in_links: + for run in runs: + if run is not carrier: + _clear_run_text(run) + + +def _clear_paragraph_text(p_el: "_Element") -> None: + """Empty a paragraph's text: delete direct pure-text runs, empty the + text runs inside hyperlinks (hyperlink elements stay).""" + direct, in_links = _paragraph_text_runs(p_el) + for run in direct: + p_el.remove(run) + for _, runs in in_links: + for run in runs: + _clear_run_text(run) + + +def _clear_row_copy(tr_el: "_Element") -> None: + """Blank a deep-copied row template: clear every paragraph's text and + drop the (copied) hyperlink elements. Cell count, ``w:trPr``, + ``w:tcPr`` (shading, spans, merges) all stay.""" + for p_el in tr_el.iter(qn("w:p")): + _clear_paragraph_text(p_el) + for h in list(p_el.findall(qn("w:hyperlink"))): + p_el.remove(h) + + +# -- model classes ------------------------------------------------------------- + + +class RawParagraph: + """A body-level paragraph (positional snapshot — re-fetch + ``document.paragraphs`` after structural edits).""" + + __slots__ = ("element", "index") + + def __init__(self, element: "_Element", index: int): + self.element = element + self.index = index + + @property + def text(self) -> str: + return _element_text(self.element) + + @property + def style(self) -> str: + """The ``w:pStyle`` style *id* (e.g. ``"Heading1"``), or ``"Normal"``.""" + style = self.element.find(f"{qn('w:pPr')}/{qn('w:pStyle')}") + return style.get(qn("w:val")) if style is not None else "Normal" + + def __repr__(self) -> str: # pragma: no cover - debug aid + return f"" + + +class RawCell: + """One (grid-resolved) table cell.""" + + __slots__ = ("element", "_doc") + + def __init__(self, element: "_Element", doc: "DocxRawDocument"): + self.element = element + self._doc = doc + + def _paragraph_els(self) -> list: + return self.element.findall(qn("w:p")) + + @property + def text(self) -> str: + """python-docx parity: direct paragraphs joined with newlines + (nested-table text is *not* included).""" + return "\n".join(_element_text(p) for p in self._paragraph_els()) + + @property + def paragraph_count(self) -> int: + return len(self._paragraph_els()) + + def set_text(self, text: str) -> None: + """Run- and layout-preserving text replace (see module docstring). + + The carrier is the first paragraph that owns a text run (else + the first paragraph without protected content, else a new + trailing paragraph). Paragraphs containing drawings / pictures / + OLE stay byte-for-byte untouched; other pure-text paragraphs are + emptied but never removed, so images and nested tables keep + their positions. + """ + from lxml import etree + + paras = self._paragraph_els() + target = next((p for p in paras if _has_text_run(p)), None) + if target is None: + target = next((p for p in paras if not _contains_protected(p)), None) + if target is None: + target = etree.SubElement(self.element, qn("w:p")) + _replace_paragraph_text(target, text) + for p in paras: + if p is not target and not _contains_protected(p): + _clear_paragraph_text(p) + self._doc._mark_document_dirty() + + def __repr__(self) -> str: # pragma: no cover - debug aid + return f"" + + +class RawTable: + """A table (body-level or nested) with grid-resolved cell access.""" + + __slots__ = ("element", "index", "_doc") + + def __init__(self, element: "_Element", index: int, doc: "DocxRawDocument"): + self.element = element + self.index = index + self._doc = doc + + def _row_els(self) -> list: + return self.element.findall(qn("w:tr")) + + @property + def n_rows(self) -> int: + return len(self._row_els()) + + @property + def n_cols(self) -> int: + grid = self.element.find(qn("w:tblGrid")) + return len(grid.findall(qn("w:gridCol"))) if grid is not None else 0 + + # -- grid resolution ----------------------------------------------------- + + def _grid(self) -> list[list]: + """``matrix[r][c]`` → the ``w:tc`` owning grid position (r, c). + + gridSpan repeats a cell across the columns it spans; a vMerge + continuation resolves to the start cell of the merge — the same + answers python-docx gives for ``table.rows[r].cells[c]``. + """ + rows = self._row_els() + n_cols = self.n_cols + matrix: list[list] = [[None] * n_cols for _ in rows] + for ri, tr in enumerate(rows): + ci = 0 + before = tr.find(f"{qn('w:trPr')}/{qn('w:gridBefore')}") + if before is not None: + ci = int(before.get(qn("w:val"), 0)) + for tc in tr.findall(qn("w:tc")): + tc_pr = tc.find(qn("w:tcPr")) + span, resolved = 1, tc + if tc_pr is not None: + grid_span = tc_pr.find(qn("w:gridSpan")) + if grid_span is not None: + span = int(grid_span.get(qn("w:val"), 1)) + v_merge = tc_pr.find(qn("w:vMerge")) + if ( + v_merge is not None + and v_merge.get(qn("w:val"), "continue") == "continue" + and ri > 0 + and ci < n_cols + and matrix[ri - 1][ci] is not None + ): + resolved = matrix[ri - 1][ci] + for k in range(span): + if ci + k < n_cols: + matrix[ri][ci + k] = resolved + ci += span + return matrix + + def _tc_at(self, r: int, c: int) -> "_Element": + if not (0 <= r < self.n_rows and 0 <= c < self.n_cols): + raise IndexError( + f"cell ({r}, {c}) out of range for {self.n_rows}x{self.n_cols} table" + ) + tc = self._grid()[r][c] + if tc is None: + raise IndexError(f"grid position ({r}, {c}) has no cell") + return tc + + def cell(self, r: int, c: int) -> RawCell: + return RawCell(self._tc_at(r, c), self._doc) + + def nested_tables(self, r: int, c: int) -> list["RawTable"]: + """Tables directly inside cell (r, c), addressable like any table.""" + return [ + RawTable(el, i, self._doc) + for i, el in enumerate(self._tc_at(r, c).findall(qn("w:tbl"))) + ] + + # -- row editing -------------------------------------------------------- + + def insert_row(self, idx: int) -> None: + """Insert a blank row at *idx* (0 ≤ idx ≤ n_rows), deep-copying the + row above (row 0 for idx=0) as the template: cell count and all + row/cell properties carry over, text is cleared.""" + rows = self._row_els() + if not rows: + raise ValueError("cannot insert into a table with no rows") + if not 0 <= idx <= len(rows): + raise IndexError(f"row index {idx} out of range 0..{len(rows)}") + template = rows[idx - 1] if idx > 0 else rows[0] + new_tr = copy.deepcopy(template) + _clear_row_copy(new_tr) + if idx == len(rows): + rows[-1].addnext(new_tr) + else: + rows[idx].addprevious(new_tr) + self._doc._mark_document_dirty() + + def delete_row(self, idx: int) -> None: + rows = self._row_els() + if not 0 <= idx < len(rows): + raise IndexError(f"row index {idx} out of range 0..{len(rows) - 1}") + self.element.remove(rows[idx]) + self._doc._mark_document_dirty() + + def __repr__(self) -> str: # pragma: no cover - debug aid + return f"" + + +class DocxRawDocument(RawDocumentBase): + """Raw model over ``word/document.xml`` (plus header/footer/chart + parts). Everything the model doesn't touch round-trips + byte-identically via the OPC container.""" + + format = "docx" + + def __init__(self, package: "OpcPackage"): + super().__init__(package) + if not package.has_part(_DOCUMENT_PART): + from contextifier.raw.opc import RawUnsupportedError + + raise RawUnsupportedError("Package has no word/document.xml — not a docx") + self._document = self.xml_part(_DOCUMENT_PART) + self._charts: list[ChartModel] | None = None + + # -- internals ----------------------------------------------------------- + + @property + def _body(self) -> "_Element": + return self._document.find("w:body") + + def _mark_document_dirty(self) -> None: + self._document.mark_dirty() + + def _paragraph_els(self) -> list: + return self._body.findall(qn("w:p")) + + def _paragraph_el(self, index: int) -> "_Element": + els = self._paragraph_els() + if not 0 <= index < len(els): + raise IndexError(f"paragraph index {index} out of range 0..{len(els) - 1}") + return els[index] + + # -- paragraphs ----------------------------------------------------------- + + @property + def paragraphs(self) -> list[RawParagraph]: + """Body-level paragraphs, indexed like python-docx + ``document.paragraphs``.""" + return [RawParagraph(el, i) for i, el in enumerate(self._paragraph_els())] + + def set_paragraph_text(self, index: int, text: str) -> None: + """Run-preserving text replace (see module docstring): first text + run carries the new text with its formatting, protected runs + (images, OLE, bookmarks, math, ...) stay in place, hyperlink + elements survive (possibly with empty text).""" + _replace_paragraph_text(self._paragraph_el(index), text) + self._mark_document_dirty() + + def strip_empty_hyperlinks(self, index: int) -> int: + """Remove the paragraph's ``w:hyperlink`` children whose text is + empty (e.g. after :meth:`set_paragraph_text`). Relationships no + longer referenced anywhere in the document part are dropped too. + Returns the number of hyperlinks removed.""" + p_el = self._paragraph_el(index) + removed_ids: list[str] = [] + removed = 0 + for h in list(p_el.findall(qn("w:hyperlink"))): + if _element_text(h) == "": + rid = h.get(qn("r:id")) + p_el.remove(h) + removed += 1 + if rid: + removed_ids.append(rid) + if removed: + self._mark_document_dirty() + if removed_ids: + rels = self.package.rels_for(_DOCUMENT_PART) + r_ns = "{%s}" % NS["r"] + for rid in removed_ids: + still_used = any( + value == rid + for el in self._document.root.iter() + for key, value in el.attrib.items() + if key.startswith(r_ns) + ) + if rels is not None and not still_used: + rels.remove(rid) + return removed + + def insert_paragraph_after( + self, index: int, text: str, style: str | None = None + ) -> RawParagraph: + """Insert a new paragraph after ``paragraphs[index]``; ``index=-1`` + inserts at the very start of the body. Safe with a trailing body + ``w:sectPr`` (insertion is always anchored to an existing + paragraph, never appended past it).""" + from lxml import etree + + new_p = etree.Element(qn("w:p")) + if style: + p_pr = etree.SubElement(new_p, qn("w:pPr")) + p_style = etree.SubElement(p_pr, qn("w:pStyle")) + p_style.set(qn("w:val"), style) + run = etree.SubElement(new_p, qn("w:r")) + _set_run_text(run, text) + + if index == -1: + self._body.insert(0, new_p) + new_index = 0 + else: + self._paragraph_el(index).addnext(new_p) + new_index = index + 1 + self._mark_document_dirty() + return RawParagraph(new_p, new_index) + + def delete_paragraph(self, index: int) -> None: + el = self._paragraph_el(index) + el.getparent().remove(el) + self._mark_document_dirty() + + # -- tables --------------------------------------------------------------- + + @property + def tables(self) -> list[RawTable]: + """Body-level tables, indexed like python-docx ``document.tables`` + (nested tables are reached via :meth:`RawTable.nested_tables`).""" + return [ + RawTable(el, i, self) + for i, el in enumerate(self._body.findall(qn("w:tbl"))) + ] + + # -- charts --------------------------------------------------------------- + + @property + def chart_part_names(self) -> list[str]: + """Chart parts referenced from the document part's rels.""" + return find_chart_parts(self.package, _DOCUMENT_PART) + + @property + def charts(self) -> list[ChartModel]: + """Lazy :class:`ChartModel` per chart part (shared C3 contract).""" + if self._charts is None: + self._charts = [ + ChartModel(self.xml_part(name), self.package) + for name in self.chart_part_names + ] + return self._charts + + # -- headers / footers ------------------------------------------------------ + + def _header_footer_text(self, type_suffix: str) -> dict[str, str]: + rels = self.package.rels_for(_DOCUMENT_PART) + if rels is None: + return {} + out: dict[str, str] = {} + for rel in rels.by_type(type_suffix): + name = rels.resolve(_DOCUMENT_PART, rel["target"]) + root = self.xml_part(name).root + out[name] = "\n".join(_element_text(p) for p in root.iter(qn("w:p"))) + return out + + @property + def headers(self) -> dict[str, str]: + """Read-only ``{part name: text}`` for the header parts.""" + return self._header_footer_text("/header") + + @property + def footers(self) -> dict[str, str]: + """Read-only ``{part name: text}`` for the footer parts.""" + return self._header_footer_text("/footer") + + # -- structure --------------------------------------------------------------- + + def body_order(self) -> list[tuple[str, int]]: + """Document order of body children as ``("p"|"tbl", index)`` pairs — + the skeleton outline builders walk.""" + out: list[tuple[str, int]] = [] + p_i = t_i = 0 + for child in self._body: + if child.tag == qn("w:p"): + out.append(("p", p_i)) + p_i += 1 + elif child.tag == qn("w:tbl"): + out.append(("tbl", t_i)) + t_i += 1 + return out diff --git a/contextifier/raw/opc.py b/contextifier/raw/opc.py new file mode 100644 index 0000000..223e567 --- /dev/null +++ b/contextifier/raw/opc.py @@ -0,0 +1,355 @@ +# contextifier/raw/opc.py +""" +OPC (Open Packaging Conventions) container with a byte-preservation +write contract. + +An OOXML file is a ZIP of *parts* plus ``[Content_Types].xml`` and +``_rels/*.rels`` relationship graphs. The higher-level Office libraries +(openpyxl & friends) parse the whole package into their own object model +and re-serialize everything on save — silently destroying whatever they +don't model. This container does the opposite: + +**Byte-preservation contract** — on :meth:`OpcPackage.save`, every part +that was not explicitly written through :meth:`OpcPart.write` (or +removed/added) is emitted with its original decompressed bytes, +unchanged. Only dirty parts are re-serialized. The ZIP container itself +may be re-encoded (compression is not part of the contract); part +*content* is. + +The container knows nothing about spreadsheets or slides — format +models (:mod:`contextifier.raw.xlsx` etc.) sit on top. +""" + +from __future__ import annotations + +import io +import posixpath +import zipfile +from pathlib import Path +from typing import BinaryIO, Iterator + +from contextifier.errors import ContextifierError + +__all__ = ["OpcPackage", "OpcPart", "Relationships", "RawUnsupportedError"] + +_CONTENT_TYPES = "[Content_Types].xml" + +# Mirrors the pipeline's zip-bomb guard (handlers validate the same cap). +_MAX_UNCOMPRESSED = 1 << 30 # 1 GiB + + +class RawUnsupportedError(ContextifierError): + """The requested raw capability does not exist for this format.""" + + +class OpcPart: + """One part (file entry) inside the package. + + Parts are lazy: bytes are read from the source ZIP on first access. + Writing replaces the content and marks the part dirty; clean parts + round-trip byte-identically. + """ + + __slots__ = ("name", "_package", "_data", "dirty", "is_new") + + def __init__(self, name: str, package: "OpcPackage", *, is_new: bool = False): + self.name = name + self._package = package + self._data: bytes | None = None + self.dirty = is_new + self.is_new = is_new + + def read(self) -> bytes: + if self._data is None: + self._data = self._package._read_source(self.name) + return self._data + + def write(self, data: bytes) -> None: + self._data = bytes(data) + self.dirty = True + + def __repr__(self) -> str: # pragma: no cover - debug aid + state = "new" if self.is_new else ("dirty" if self.dirty else "clean") + return f"" + + +class Relationships: + """A part's ``_rels`` graph (``_rels/.rels``). + + Minimal, allocation-light XML handling: relationships are parsed with + lxml on demand and re-serialized only when mutated. + """ + + NS = "http://schemas.openxmlformats.org/package/2006/relationships" + + def __init__(self, part: OpcPart): + from lxml import etree + + self._part = part + self._root = etree.fromstring(part.read()) + + def __iter__(self) -> Iterator[dict]: + for rel in self._root: + if rel.tag == f"{{{self.NS}}}Relationship": + yield { + "id": rel.get("Id"), + "type": rel.get("Type"), + "target": rel.get("Target"), + "mode": rel.get("TargetMode", "Internal"), + } + + def target_of(self, rel_id: str) -> str | None: + for rel in self: + if rel["id"] == rel_id: + return rel["target"] + return None + + def by_type(self, type_suffix: str) -> list[dict]: + """Relationships whose Type ends with *type_suffix* (e.g. ``/chart``).""" + return [rel for rel in self if rel["type"].endswith(type_suffix)] + + def resolve(self, base_part: str, target: str) -> str: + """Resolve a relationship target to an absolute part name.""" + if target.startswith("/"): + return target[1:] + base_dir = posixpath.dirname(base_part) + return posixpath.normpath(posixpath.join(base_dir, target)) + + def add( + self, rel_id: str, rel_type: str, target: str, *, external: bool = False + ) -> None: + from lxml import etree + + rel = etree.SubElement(self._root, f"{{{self.NS}}}Relationship") + rel.set("Id", rel_id) + rel.set("Type", rel_type) + rel.set("Target", target) + if external: + rel.set("TargetMode", "External") + self.flush() + + def remove(self, rel_id: str) -> bool: + for rel in list(self._root): + if rel.get("Id") == rel_id: + self._root.remove(rel) + self.flush() + return True + return False + + def next_id(self) -> str: + used = {rel["id"] for rel in self} + n = len(used) + 1 + while f"rId{n}" in used: + n += 1 + return f"rId{n}" + + def flush(self) -> None: + from lxml import etree + + self._part.write( + etree.tostring( + self._root, xml_declaration=True, encoding="UTF-8", standalone=True + ) + ) + + +class OpcPackage: + """The package: parts + content types + relationship graphs. + + Open with :meth:`open`, mutate parts, then :meth:`save` / + :meth:`to_bytes`. Untouched parts keep their exact original bytes. + """ + + def __init__(self, source_bytes: bytes): + self._source = source_bytes + self._zip = zipfile.ZipFile(io.BytesIO(source_bytes)) + total = sum(i.file_size for i in self._zip.infolist()) + if total > _MAX_UNCOMPRESSED: + raise ContextifierError( + f"Package inflates to {total} bytes (> {_MAX_UNCOMPRESSED}); refusing" + ) + self._parts: dict[str, OpcPart] = { + info.filename: OpcPart(info.filename, self) + for info in self._zip.infolist() + if not info.is_dir() + } + self._removed: set[str] = set() + self._rels_cache: dict[str, Relationships] = {} + + # -- construction -------------------------------------------------------- + + @classmethod + def open(cls, source: str | Path | bytes | BinaryIO) -> "OpcPackage": + if isinstance(source, (str, Path)): + data = Path(source).read_bytes() + elif isinstance(source, (bytes, bytearray)): + data = bytes(source) + else: + data = source.read() + if data[:4] != b"PK\x03\x04": + raise ContextifierError("Not a ZIP/OPC package (bad magic)") + return cls(data) + + def sniff_format(self) -> str | None: + """Best-effort format detection from the package layout.""" + if self.has_part("xl/workbook.xml"): + return "xlsx" + if self.has_part("word/document.xml"): + return "docx" + if self.has_part("ppt/presentation.xml"): + return "pptx" + return None + + # -- part access ---------------------------------------------------------- + + def _read_source(self, name: str) -> bytes: + return self._zip.read(name) + + @property + def part_names(self) -> list[str]: + return sorted(self._parts) + + def has_part(self, name: str) -> bool: + return name in self._parts + + def get_part(self, name: str) -> OpcPart: + try: + return self._parts[name] + except KeyError: + raise KeyError(f"No such part: {name!r}") from None + + __getitem__ = get_part + + def add_part(self, name: str, data: bytes) -> OpcPart: + part = OpcPart(name, self, is_new=True) + part.write(data) + self._parts[name] = part + self._removed.discard(name) + return part + + def remove_part(self, name: str) -> None: + self._parts.pop(name, None) + self._rels_cache.pop(self._rels_name_for(name), None) + self._removed.add(name) + + # -- relationships & content types --------------------------------------- + + @staticmethod + def _rels_name_for(part_name: str) -> str: + if part_name == "": # package-level rels + return "_rels/.rels" + directory = posixpath.dirname(part_name) + base = posixpath.basename(part_name) + return ( + posixpath.join(directory, "_rels", f"{base}.rels") + if directory + else f"_rels/{base}.rels" + ) + + def rels_for(self, part_name: str) -> Relationships | None: + """The relationships of *part_name* ('' = package root), or None.""" + rels_name = self._rels_name_for(part_name) + if rels_name in self._rels_cache: + return self._rels_cache[rels_name] + if not self.has_part(rels_name): + return None + rels = Relationships(self.get_part(rels_name)) + self._rels_cache[rels_name] = rels + return rels + + def content_type_of(self, part_name: str) -> str | None: + from lxml import etree + + root = etree.fromstring(self.get_part(_CONTENT_TYPES).read()) + ns = "http://schemas.openxmlformats.org/package/2006/content-types" + for el in root: + if el.tag == f"{{{ns}}}Override" and el.get("PartName") == f"/{part_name}": + return el.get("ContentType") + ext = part_name.rsplit(".", 1)[-1].lower() + for el in root: + if ( + el.tag == f"{{{ns}}}Default" + and (el.get("Extension") or "").lower() == ext + ): + return el.get("ContentType") + return None + + def set_content_type_override(self, part_name: str, content_type: str) -> None: + from lxml import etree + + ct = self.get_part(_CONTENT_TYPES) + root = etree.fromstring(ct.read()) + ns = "http://schemas.openxmlformats.org/package/2006/content-types" + for el in root: + if el.tag == f"{{{ns}}}Override" and el.get("PartName") == f"/{part_name}": + el.set("ContentType", content_type) + break + else: + override = etree.SubElement(root, f"{{{ns}}}Override") + override.set("PartName", f"/{part_name}") + override.set("ContentType", content_type) + ct.write( + etree.tostring( + root, xml_declaration=True, encoding="UTF-8", standalone=True + ) + ) + + # -- saving ---------------------------------------------------------------- + + def to_bytes(self) -> bytes: + """Serialize the package. Clean parts are byte-identical.""" + buf = io.BytesIO() + original_infos = {i.filename: i for i in self._zip.infolist()} + with zipfile.ZipFile(buf, "w") as out: + # Original entry order first (Office is order-tolerant, but + # keeping it minimizes diffs), then new parts. + for name in [ + *(n for n in original_infos if n in self._parts), + *(n for n in self._parts if n not in original_infos), + ]: + part = self._parts[name] + src_info = original_infos.get(name) + if part.dirty or src_info is None: + data = part.read() + compress = ( + src_info.compress_type + if src_info is not None + else zipfile.ZIP_DEFLATED + ) + out.writestr( + zipfile.ZipInfo(name, date_time=(1980, 1, 1, 0, 0, 0)), + data, + compress_type=compress, + ) + else: + # Clean part: original decompressed bytes, original + # compression method, original timestamp. COPY the + # ZipInfo — writestr mutates it (header_offset/sizes), + # which would corrupt the source reader on a second + # to_bytes() call. + info_copy = zipfile.ZipInfo( + src_info.filename, date_time=src_info.date_time + ) + info_copy.compress_type = src_info.compress_type + info_copy.external_attr = src_info.external_attr + info_copy.internal_attr = src_info.internal_attr + info_copy.create_system = src_info.create_system + out.writestr(info_copy, self._zip.read(name)) + return buf.getvalue() + + def save(self, target: str | Path | BinaryIO | None = None) -> bytes: + data = self.to_bytes() + if isinstance(target, (str, Path)): + Path(target).write_bytes(data) + elif target is not None: + target.write(data) + return data + + def close(self) -> None: + self._zip.close() + + def __enter__(self) -> "OpcPackage": + return self + + def __exit__(self, *exc) -> None: + self.close() diff --git a/contextifier/raw/pptx.py b/contextifier/raw/pptx.py new file mode 100644 index 0000000..2ac93a7 --- /dev/null +++ b/contextifier/raw/pptx.py @@ -0,0 +1,653 @@ +# contextifier/raw/pptx.py +""" +PptxRawDocument — the raw (lossless, writable) model for .pptx decks. + +Slides are thin views over their XML parts: reading never dirties, and +edits are surgical, so untouched parts round-trip byte-identical (the +OPC byte-preservation contract). Two operations exist specifically to +fix what template-driven pipelines (edit2docs-style) get wrong today: + +* :meth:`RawSlide.replace_content` — swap a slide's XML for freshly + generated markup while *pulling the original native objects along*: + chart / table / diagram graphicFrames (and optionally pictures) are + lifted out of the old tree and re-appended into the new one, ids + renumbered, relationships untouched. Native charts stay native + instead of being rasterized or dropped. +* :meth:`PptxRawDocument.remove_slide` — deletes the slide *and* + reference-counts every part it pulled in (charts, embedded workbooks, + images, notes), removing the ones no remaining slide uses. No more + orphan-part bloat in the package. +""" + +from __future__ import annotations + +import copy +import posixpath +from dataclasses import dataclass +from typing import TYPE_CHECKING, Iterator + +from contextifier.raw.base import RawDocumentBase +from contextifier.raw.chart import ChartModel, find_chart_parts +from contextifier.raw.opc import OpcPackage +from contextifier.raw.xmlpart import NS, qn + +if TYPE_CHECKING: # pragma: no cover + from lxml.etree import _Element + +__all__ = [ + "PptxRawDocument", + "RawSlide", + "RawShapeInfo", + "RawTable", + "RawTableCell", +] + +_PRESENTATION = "ppt/presentation.xml" + +#: graphicData/@uri tail → shape kind +_GRAPHIC_KIND = { + "table": "table", + "chart": "chart", + "chartex": "chart", + "diagram": "diagram", +} + +#: uri tails of graphicFrames replace_content must always carry over +_NATIVE_FRAME_TAILS = ("table", "chart", "chartex", "diagram") + +#: nv*Pr wrappers whose first p:cNvPr identifies the shape +_NV_PR_TAGS = frozenset( + qn(t) + for t in ( + "p:nvSpPr", + "p:nvPicPr", + "p:nvGraphicFramePr", + "p:nvGrpSpPr", + "p:nvCxnSpPr", + ) +) + + +@dataclass +class RawShapeInfo: + """Inventory entry for one shape on a slide.""" + + id: int + name: str + kind: str # "text" | "picture" | "table" | "chart" | "group" | "diagram" | "other" + text: str | None + + +# -- XML helpers --------------------------------------------------------------- + + +def _para_text(para: "_Element") -> str: + return "".join(t.text or "" for t in para.iter(qn("a:t"))) + + +def _body_text(tx_body: "_Element") -> str: + return "\n".join(_para_text(p) for p in tx_body.findall(qn("a:p"))) + + +def _replace_para_text(para: "_Element", new_text: str) -> None: + """Set a paragraph's text under the format-preserving contract. + + The first ``a:r`` keeps its ``a:rPr`` (formatting) and receives the + new text; every *other* plain-text run is dropped; anything that is + not a plain run — ``a:fld``, ``a:br``, ``a:pPr``, ``a:endParaRPr`` — + stays exactly where it was. + """ + from lxml import etree + + a_r, a_t = qn("a:r"), qn("a:t") + runs = [child for child in para if child.tag == a_r] + if runs: + first = runs[0] + t = first.find(a_t) + if t is None: + t = etree.SubElement(first, a_t) + t.text = new_text + for extra in runs[1:]: + para.remove(extra) + return + run = etree.SubElement(para, a_r) + etree.SubElement(run, a_t).text = new_text + end = para.find(qn("a:endParaRPr")) + if end is not None: + end.addprevious(run) + + +def _shape_cnvpr(shape_el: "_Element") -> "_Element | None": + """The shape's own ``p:cNvPr`` (never a nested child's).""" + for child in shape_el: + if child.tag in _NV_PR_TAGS: + return child.find(qn("p:cNvPr")) + return None + + +def _graphic_kind(frame_el: "_Element") -> str | None: + data = frame_el.find("a:graphic/a:graphicData", NS) + if data is None: + return None + uri = data.get("uri") or "" + return _GRAPHIC_KIND.get(uri.rsplit("/", 1)[-1]) + + +def _walk_shapes(container: "_Element") -> Iterator["_Element"]: + """Shape elements in document order, descending into p:grpSp.""" + shape_tags = ( + qn("p:sp"), + qn("p:pic"), + qn("p:graphicFrame"), + qn("p:grpSp"), + qn("p:cxnSp"), + ) + for child in container: + if child.tag in shape_tags: + yield child + if child.tag == qn("p:grpSp"): + yield from _walk_shapes(child) + + +def _first_rid(el: "_Element", attr: str) -> str | None: + """First ``attr`` (a qualified name) found on *el* or a descendant.""" + for node in el.iter(): + rid = node.get(attr) + if rid: + return rid + return None + + +# -- tables -------------------------------------------------------------------- + + +class RawTableCell: + """One ``a:tc`` — read/format-preserving-write of its text.""" + + def __init__(self, table: "RawTable", tc_el: "_Element"): + self._table = table + self._tc = tc_el + + @property + def text(self) -> str: + tx = self._tc.find(qn("a:txBody")) + return _body_text(tx) if tx is not None else "" + + def set_text(self, text: str) -> None: + """Replace the first paragraph's text (first run's ``a:rPr`` and + any non-text elements are preserved — same rules as + :meth:`RawSlide.set_text`).""" + from lxml import etree + + tx = self._tc.find(qn("a:txBody")) + if tx is None: + tx = etree.SubElement(self._tc, qn("a:txBody")) + etree.SubElement(tx, qn("a:bodyPr")) + etree.SubElement(tx, qn("a:lstStyle")) + para = tx.find(qn("a:p")) + if para is None: + para = etree.SubElement(tx, qn("a:p")) + _replace_para_text(para, text) + self._table._slide._mark_dirty() + + +class RawTable: + """A native DrawingML table (``a:tbl``) hosted in a graphicFrame. + + v0.4 scope: cell text + row insert/delete. Column operations would + require ``a:gridCol`` surgery and are deliberately out of scope. + """ + + def __init__(self, slide: "RawSlide", frame_el: "_Element"): + self._slide = slide + self._frame = frame_el + cnvpr = _shape_cnvpr(frame_el) + self.shape_id: int = int(cnvpr.get("id")) if cnvpr is not None else -1 + tbl = frame_el.find("a:graphic/a:graphicData/a:tbl", NS) + if tbl is None: # pragma: no cover - guarded by caller's uri check + raise ValueError("graphicFrame does not contain an a:tbl") + self._tbl = tbl + + # -- geometry --------------------------------------------------------------- + + @property + def _rows(self) -> list["_Element"]: + return self._tbl.findall(qn("a:tr")) + + @property + def n_rows(self) -> int: + return len(self._rows) + + @property + def n_cols(self) -> int: + return len(self._tbl.findall("a:tblGrid/a:gridCol", NS)) + + # -- cells ------------------------------------------------------------------ + + def cell(self, r: int, c: int) -> RawTableCell: + rows = self._rows + if not 0 <= r < len(rows): + raise IndexError(f"row {r} out of range (table has {len(rows)} rows)") + cells = rows[r].findall(qn("a:tc")) + if not 0 <= c < len(cells): + raise IndexError(f"col {c} out of range (row has {len(cells)} cells)") + return RawTableCell(self, cells[c]) + + # -- rows ------------------------------------------------------------------- + + def insert_row(self, idx: int) -> None: + """Insert an empty row at *idx*, cloning the row above (or the + first row) as the style template. Column count stays consistent + with ``a:tblGrid`` because the template row already matches it.""" + from lxml import etree + + rows = self._rows + if not rows: + raise ValueError("cannot insert into a table with no template row") + if not 0 <= idx <= len(rows): + raise IndexError(f"insert index {idx} out of range (0..{len(rows)})") + template = rows[idx - 1] if idx > 0 else rows[0] + new_row = copy.deepcopy(template) + for tc in new_row.findall(qn("a:tc")): + tx = tc.find(qn("a:txBody")) + if tx is not None: + for para in tx.findall(qn("a:p")): + tx.remove(para) + etree.SubElement(tx, qn("a:p")) + if idx == len(rows): + rows[-1].addnext(new_row) + else: + rows[idx].addprevious(new_row) + self._slide._mark_dirty() + + def delete_row(self, idx: int) -> None: + rows = self._rows + if not 0 <= idx < len(rows): + raise IndexError(f"row {idx} out of range (table has {len(rows)} rows)") + self._tbl.remove(rows[idx]) + self._slide._mark_dirty() + + +# -- slides -------------------------------------------------------------------- + + +class RawSlide: + """One slide part, addressed by shape id.""" + + def __init__(self, doc: "PptxRawDocument", part_name: str, index: int): + self._doc = doc + self.part_name = part_name + self.index = index + + def __repr__(self) -> str: # pragma: no cover - debug aid + return f"" + + @property + def _xp(self): + return self._doc.xml_part(self.part_name) + + def _mark_dirty(self) -> None: + self._xp.mark_dirty() + + @property + def _sp_tree(self) -> "_Element": + tree = self._xp.find("p:cSld/p:spTree") + if tree is None: + raise ValueError(f"{self.part_name} has no p:cSld/p:spTree") + return tree + + # -- inventory --------------------------------------------------------------- + + @property + def shapes(self) -> list[RawShapeInfo]: + """All shapes (document order, groups flattened in place).""" + out: list[RawShapeInfo] = [] + for el in _walk_shapes(self._sp_tree): + cnvpr = _shape_cnvpr(el) + shape_id = int(cnvpr.get("id")) if cnvpr is not None else -1 + name = (cnvpr.get("name") or "") if cnvpr is not None else "" + kind, text = "other", None + if el.tag == qn("p:sp"): + tx = el.find(qn("p:txBody")) + if tx is not None: + kind, text = "text", _body_text(tx) + elif el.tag == qn("p:pic"): + kind = "picture" + elif el.tag == qn("p:grpSp"): + kind = "group" + elif el.tag == qn("p:graphicFrame"): + kind = _graphic_kind(el) or "other" + out.append(RawShapeInfo(id=shape_id, name=name, kind=kind, text=text)) + return out + + def _find_shape(self, shape_id: int) -> "_Element": + for el in _walk_shapes(self._sp_tree): + cnvpr = _shape_cnvpr(el) + if cnvpr is not None and cnvpr.get("id") == str(shape_id): + return el + raise KeyError(f"No shape with id={shape_id} on {self.part_name}") + + # -- text -------------------------------------------------------------------- + + def get_text(self, shape_id: int) -> str: + el = self._find_shape(shape_id) + tx = el.find(qn("p:txBody")) + if tx is None: + raise ValueError(f"Shape id={shape_id} has no text body") + return _body_text(tx) + + def set_text(self, shape_id: int, new_text: str, para: int = 0) -> None: + """Replace paragraph *para*'s text, preserving the first run's + formatting (``a:rPr``) and any non-text elements (``a:fld``, + ``a:br``); other plain-text runs in the paragraph are removed.""" + el = self._find_shape(shape_id) + tx = el.find(qn("p:txBody")) + if tx is None: + raise ValueError(f"Shape id={shape_id} has no text body") + paras = tx.findall(qn("a:p")) + if not 0 <= para < len(paras): + raise IndexError( + f"paragraph {para} out of range (shape has {len(paras)} paragraphs)" + ) + _replace_para_text(paras[para], new_text) + self._mark_dirty() + + # -- tables / charts / notes --------------------------------------------------- + + @property + def tables(self) -> list[RawTable]: + return [ + RawTable(self, el) + for el in _walk_shapes(self._sp_tree) + if el.tag == qn("p:graphicFrame") and _graphic_kind(el) == "table" + ] + + @property + def chart_part_names(self) -> list[str]: + """Chart parts referenced from this slide's relationships.""" + return find_chart_parts(self._doc.package, self.part_name) + + @property + def charts(self) -> list[ChartModel]: + """Lazy :class:`ChartModel` views (implemented in milestone C3).""" + return [ + ChartModel(self._doc.xml_part(name), self._doc.package) + for name in self.chart_part_names + ] + + @property + def notes_text(self) -> str | None: + """Text of the notes slide's body placeholder, or ``None``.""" + rels = self._doc.package.rels_for(self.part_name) + if rels is None: + return None + notes_rels = rels.by_type("/notesSlide") + if not notes_rels: + return None + notes_part = rels.resolve(self.part_name, notes_rels[0]["target"]) + xp = self._doc.xml_part(notes_part) + for sp in xp.findall(".//p:sp"): + ph = sp.find("p:nvSpPr/p:nvPr/p:ph", NS) + tx = sp.find(qn("p:txBody")) + if tx is not None and ph is not None and ph.get("type") == "body": + return _body_text(tx) + texts = [_body_text(tx) for tx in xp.root.iter(qn("p:txBody"))] + return "\n".join(t for t in texts if t) + + # -- content replacement -------------------------------------------------------- + + def replace_content( + self, + new_slide_xml: bytes, + *, + preserve_native: bool = True, + preserve_pictures: bool = True, + ) -> list[str]: + """Replace this slide's XML with *new_slide_xml*, carrying the + original native objects over. + + With ``preserve_native=True`` every chart / table / diagram + graphicFrame of the ORIGINAL slide (and, when + ``preserve_pictures``, every ``p:pic`` that references embedded + media) is deep-copied into the new tree's ``p:spTree``, with + their ``p:cNvPr/@id`` renumbered past any id used by the new + XML. The slide keeps its part name, so its relationships part is + untouched and every carried-over ``r:id`` / ``r:embed`` keeps + resolving. + + *new_slide_xml* must be a complete ```` document and may + only reference relationship ids that already exist in this + slide's rels (or none at all) — this method never edits the rels + part, so unknown ``r:id`` / ``r:embed`` values in the new XML + would dangle. + + Returns descriptions of the preserved elements, e.g. + ``["table", "chart:chart1.xml", "picture:image1.png"]``. + """ + from lxml import etree + + new_root = etree.fromstring(new_slide_xml) + if new_root.tag != qn("p:sld"): + raise ValueError("new_slide_xml must be a complete document") + new_tree = new_root.find("p:cSld/p:spTree", NS) + if new_tree is None: + raise ValueError("new_slide_xml has no p:cSld/p:spTree") + + preserved: list[str] = [] + if preserve_native: + rels = self._doc.package.rels_for(self.part_name) + keep: list[tuple["_Element", str]] = [] + for el in _walk_shapes(self._sp_tree): + if el.tag == qn("p:graphicFrame"): + data = el.find("a:graphic/a:graphicData", NS) + uri_tail = ( + (data.get("uri") or "").rsplit("/", 1)[-1] + if data is not None + else "" + ) + if uri_tail not in _NATIVE_FRAME_TAILS: + continue + desc = _GRAPHIC_KIND[uri_tail] + if desc == "chart": + rid = _first_rid(data, qn("r:id")) + target = rels.target_of(rid) if (rels and rid) else None + if target: + desc = f"chart:{posixpath.basename(target)}" + keep.append((el, desc)) + elif el.tag == qn("p:pic") and preserve_pictures: + blip = el.find("p:blipFill/a:blip", NS) + rid = blip.get(qn("r:embed")) if blip is not None else None + if not rid: + continue + target = rels.target_of(rid) if rels else None + desc = ( + f"picture:{posixpath.basename(target)}" if target else "picture" + ) + keep.append((el, desc)) + + used_ids = { + int(c.get("id")) + for c in new_root.iter(qn("p:cNvPr")) + if (c.get("id") or "").isdigit() + } + next_id = max(used_ids, default=1) + 1 + for el, desc in keep: + clone = copy.deepcopy(el) + for cnvpr in clone.iter(qn("p:cNvPr")): + cnvpr.set("id", str(next_id)) + next_id += 1 + new_tree.append(clone) + preserved.append(desc) + + xp = self._xp + xp._root = new_root # swap the facade's tree in place + xp.mark_dirty() + return preserved + + +# -- document ------------------------------------------------------------------ + + +class PptxRawDocument(RawDocumentBase): + """Raw model for a .pptx package.""" + + format = "pptx" + + @property + def slides(self) -> list[RawSlide]: + """Slides in presentation order (``p:sldIdLst``).""" + pres = self.xml_part(_PRESENTATION) + rels = self.package.rels_for(_PRESENTATION) + sld_id_lst = pres.find("p:sldIdLst") + if sld_id_lst is None or rels is None: + return [] + out: list[RawSlide] = [] + for sld_id in sld_id_lst.findall(qn("p:sldId")): + rid = sld_id.get(qn("r:id")) + target = rels.target_of(rid) if rid else None + if target is None: + continue + out.append(RawSlide(self, rels.resolve(_PRESENTATION, target), len(out))) + return out + + # -- slide removal ------------------------------------------------------------ + + def remove_slide(self, index: int) -> None: + """Remove the slide at *index* and everything only it used. + + Beyond dropping the ``p:sldId`` entry and the slide part itself + (plus its rels and notes slide), this reference-counts every + part transitively reachable from the removed slide — charts, + embedded chart workbooks, images, chart colors/style parts — + against the relationships of everything still in the package, + and deletes the now-orphaned ones, including their content-type + overrides. Parts shared with surviving slides (or anchored by + the presentation / masters, like layouts and the notes master) + are untouched, and surviving slide parts stay byte-identical. + """ + slides = self.slides + if not 0 <= index < len(slides): + raise IndexError(f"slide index {index} out of range (0..{len(slides) - 1})") + part_name = slides[index].part_name + + pres = self.xml_part(_PRESENTATION) + pres_rels = self.package.rels_for(_PRESENTATION) + if pres_rels is None: # pragma: no cover - malformed package + raise ValueError("presentation has no relationships part") + rid = next( + ( + rel["id"] + for rel in pres_rels.by_type("/slide") + if pres_rels.resolve(_PRESENTATION, rel["target"]) == part_name + ), + None, + ) + sld_id_lst = pres.find("p:sldIdLst") + if sld_id_lst is not None and rid is not None: + for sld_id in list(sld_id_lst): + if sld_id.get(qn("r:id")) == rid: + sld_id_lst.remove(sld_id) + pres.mark_dirty() + if rid is not None: + pres_rels.remove(rid) + + # The slide dies; so does its notes slide. + doomed = [part_name] + slide_rels = self.package.rels_for(part_name) + if slide_rels is not None: + doomed += [ + slide_rels.resolve(part_name, rel["target"]) + for rel in slide_rels.by_type("/notesSlide") + ] + + # Everything transitively reachable from the doomed parts is an + # orphan *candidate* — collected before any rels disappear. + candidates: set[str] = set() + stack, visited = list(doomed), set(doomed) + while stack: + src = stack.pop() + rels = self.package.rels_for(src) + if rels is None: + continue + for rel in rels: + if rel["mode"] == "External": + continue + target = rels.resolve(src, rel["target"]) + if target not in visited: + visited.add(target) + candidates.add(target) + stack.append(target) + + removed: list[str] = [] + for name in doomed: + removed += self._delete_part(name) + + # Sweep candidates no remaining part references, to a fixpoint + # (deleting a chart un-anchors its embedded workbook, etc.). + while True: + referenced = self._referenced_parts() + orphans = [ + c + for c in sorted(candidates) + if self.package.has_part(c) and c not in referenced + ] + if not orphans: + break + for name in orphans: + removed += self._delete_part(name) + candidates.discard(name) + + self._drop_content_type_overrides(removed) + + # -- internals ------------------------------------------------------------------ + + def _delete_part(self, name: str) -> list[str]: + """Remove *name* and its rels part; returns what was removed.""" + removed = [] + if self.package.has_part(name): + self.package.remove_part(name) + removed.append(name) + rels_name = OpcPackage._rels_name_for(name) + if self.package.has_part(rels_name): + self.package.remove_part(rels_name) + removed.append(rels_name) + self._xml_parts.pop(name, None) + self._xml_parts.pop(rels_name, None) + return removed + + def _referenced_parts(self) -> set[str]: + """Internal targets of every relationships part still present.""" + referenced: set[str] = set() + for rels_name in list(self.package.part_names): + if not rels_name.endswith(".rels"): + continue + directory, base = posixpath.split(rels_name) + owner_dir = posixpath.dirname(directory) + owner_base = base[: -len(".rels")] + owner = posixpath.join(owner_dir, owner_base) if owner_base else "" + rels = self.package.rels_for(owner) + if rels is None: + continue + for rel in rels: + if rel["mode"] == "External": + continue + referenced.add(rels.resolve(owner, rel["target"])) + return referenced + + def _drop_content_type_overrides(self, part_names: list[str]) -> None: + from lxml import etree + + ct_part = self.package.get_part("[Content_Types].xml") + root = etree.fromstring(ct_part.read()) + doomed = {f"/{name}" for name in part_names} + changed = False + for el in list(root): + if el.tag == qn("ct:Override") and el.get("PartName") in doomed: + root.remove(el) + changed = True + if changed: + ct_part.write( + etree.tostring( + root, xml_declaration=True, encoding="UTF-8", standalone=True + ) + ) diff --git a/contextifier/raw/xlsx.py b/contextifier/raw/xlsx.py new file mode 100644 index 0000000..48b8ed3 --- /dev/null +++ b/contextifier/raw/xlsx.py @@ -0,0 +1,629 @@ +# contextifier/raw/xlsx.py +""" +XLSX raw document model — SpreadsheetML semantics on the OPC container. + +:class:`XlsxRawDocument` gives addressable, *surgical* read/write access +to a workbook. Only the worksheet XML you actually touch (plus, when +needed, ``xl/workbook.xml`` for ``calcPr``) is ever re-serialized — +every other part rides the byte-preservation contract untouched, so +charts, chart styles, pivot tables, sparkline ``extLst`` blocks, custom +XML, themes and styles all survive a save (openpyxl round-trips destroy +all of these). + +Design decisions: + +* **String writes use inline strings** (``t="inlineStr"``), never the + shared-string table — so ``xl/sharedStrings.xml`` is never rewritten + (and never *created*; a workbook without one stays without one). +* **Overwriting a formula cell removes the formula.** ``set_cell`` is + "the user supplies the value now": the ```` element is dropped and + the literal is written. Use :meth:`RawSheet.get_formula` to inspect + formulas before overwriting. +* **Stale caches recalculate on open.** Whenever a ``set_cell`` / + ``append_rows`` overwrites a formula cell — or any formula exists + anywhere in the workbook (its cached result may now be stale) — the + workbook's ```` gains ``fullCalcOnLoad="1"`` (created after + the last of sheets/definedNames if absent) so Excel recalculates. +* **Values pass through as stored.** No date-system handling (1900 vs + 1904 is ignored); date cells read back as their serial numbers. +""" + +from __future__ import annotations + +import math +import re +from typing import TYPE_CHECKING, Iterator + +from contextifier.raw.base import RawDocumentBase +from contextifier.raw.xmlpart import XmlPart, qn + +if TYPE_CHECKING: # pragma: no cover + from lxml.etree import _Element + + from contextifier.raw.chart import ChartModel + from contextifier.raw.opc import OpcPackage + +__all__ = ["XlsxRawDocument", "RawSheet", "SheetCollection"] + +#: attribute key for ``xml:space`` +_XML_SPACE = "{http://www.w3.org/XML/1998/namespace}space" + +_REF_RE = re.compile(r"^\$?([A-Za-z]{1,3})\$?([1-9][0-9]*)$") +_RANGE_RE = re.compile( + r"^\$?([A-Za-z]{1,3})\$?([1-9][0-9]*)(?::\$?([A-Za-z]{1,3})\$?([1-9][0-9]*))?$" +) + +#: workbook children that must precede (CT_Workbook sequence) +_PRE_CALC_PR = ( + "fileVersion", + "fileSharing", + "workbookPr", + "workbookProtection", + "bookViews", + "sheets", + "functionGroups", + "externalReferences", + "definedNames", +) + + +# -- A1-reference helpers ---------------------------------------------------- + + +def col_letters_to_index(letters: str) -> int: + """``"A"`` → 1, ``"B"`` → 2, ..., ``"AA"`` → 27.""" + idx = 0 + for ch in letters.upper(): + idx = idx * 26 + (ord(ch) - 64) + return idx + + +def col_index_to_letters(idx: int) -> str: + """1 → ``"A"``, 27 → ``"AA"``.""" + if idx < 1: + raise ValueError(f"Column index must be >= 1, got {idx}") + out: list[str] = [] + while idx: + idx, rem = divmod(idx - 1, 26) + out.append(chr(65 + rem)) + return "".join(reversed(out)) + + +def parse_ref(ref: str) -> tuple[int, int]: + """``"B3"`` → ``(3, 2)`` (row, column), 1-based. ``$`` anchors OK.""" + m = _REF_RE.match(ref.strip()) + if m is None: + raise ValueError(f"Not a cell reference: {ref!r}") + return int(m.group(2)), col_letters_to_index(m.group(1)) + + +def _row_number(row_el: "_Element", prev: int) -> int: + r = row_el.get("r") + return int(r) if r else prev + 1 + + +def _cell_column(c_el: "_Element", prev: int) -> int: + r = c_el.get("r") + if r: + m = _REF_RE.match(r) + if m is not None: + return col_letters_to_index(m.group(1)) + return prev + 1 + + +class RawSheet: + """One worksheet: addressable cells over the live lxml tree. + + Reads never dirty the part; every mutation marks it dirty so + ``save()`` re-serializes exactly this worksheet and nothing else. + """ + + def __init__(self, doc: "XlsxRawDocument", name: str, part_name: str): + self._doc = doc + self.name = name + self.part_name = part_name + + def __repr__(self) -> str: # pragma: no cover - debug aid + return f"" + + @property + def _xp(self) -> XmlPart: + return self._doc.xml_part(self.part_name) + + def _sheet_data(self) -> "_Element": + sd = self._xp.find("s:sheetData") + if sd is None: # degenerate but legal: create the container + sd = self._xp.root.makeelement(qn("s:sheetData"), {}) + self._xp.root.append(sd) + return sd + + # -- reading --------------------------------------------------------------- + + def get_cell(self, ref: str) -> object | None: + """The cell's value (formula cells: the *cached* value), or None. + + Types follow the cell's ``t``: shared/inline strings → str, + ``b`` → bool, ``str``/``e`` → str, numeric → int when integral + else float. Missing row/cell (or a value-less cell) → None. + """ + row_idx, col_idx = parse_ref(ref) + row = self._find_row(row_idx) + if row is None: + return None + cell = self._find_cell(row, col_idx) + if cell is None: + return None + return self._read_value(cell) + + def get_formula(self, ref: str) -> str | None: + """The cell's formula text (````, without the leading ``=``), + or None if the cell has no formula.""" + row_idx, col_idx = parse_ref(ref) + row = self._find_row(row_idx) + if row is None: + return None + cell = self._find_cell(row, col_idx) + if cell is None: + return None + f = cell.find(qn("s:f")) + return None if f is None else (f.text or "") + + @property + def dimensions(self) -> tuple[int, int]: + """``(max_row, max_col)`` actually present in sheetData; (0, 0) + for an empty sheet.""" + max_row = max_col = 0 + prev_row = 0 + for row_el in self._sheet_data(): + if row_el.tag != qn("s:row"): + continue + prev_row = _row_number(row_el, prev_row) + max_row = max(max_row, prev_row) + prev_col = 0 + for c_el in row_el: + if c_el.tag != qn("s:c"): + continue + prev_col = _cell_column(c_el, prev_col) + max_col = max(max_col, prev_col) + return (max_row, max_col) + + @property + def merged_ranges(self) -> list[str]: + """The sheet's merged ranges, e.g. ``["A5:B5"]``.""" + return [ + mc.get("ref") + for mc in self._xp.findall("s:mergeCells/s:mergeCell") + if mc.get("ref") + ] + + def iter_rows( + self, min_row: int = 1, max_row: int | None = None + ) -> Iterator[tuple[str, object | None]]: + """Yield ``(ref, value)`` for every stored cell in row order.""" + prev_row = 0 + for row_el in self._sheet_data(): + if row_el.tag != qn("s:row"): + continue + prev_row = _row_number(row_el, prev_row) + if prev_row < min_row: + continue + if max_row is not None and prev_row > max_row: + break # rows are kept sorted by @r + prev_col = 0 + for c_el in row_el: + if c_el.tag != qn("s:c"): + continue + prev_col = _cell_column(c_el, prev_col) + ref = c_el.get("r") or f"{col_index_to_letters(prev_col)}{prev_row}" + yield ref, self._read_value(c_el) + + # -- writing --------------------------------------------------------------- + + def set_cell(self, ref: str, value: object) -> None: + """Surgically write one cell, preserving everything else. + + * str → ``t="inlineStr"`` (sharedStrings is never touched) + * bool → ``t="b"``, int/float → plain numeric, None → value + removed (the cell and its style stay) + * The cell's style attribute ``s`` is preserved. + * An existing ```` formula is **removed** — set_cell means + "this literal is the value now". ``fullCalcOnLoad`` is then + ensured so Excel recalculates any dependents on open. + """ + row_idx, col_idx = parse_ref(ref) + row = self._find_row(row_idx) + if row is None: + row = self._insert_row(row_idx) + cell = self._find_cell(row, col_idx) + if cell is None: + cell = self._insert_cell(row, row_idx, col_idx) + had_formula = cell.find(qn("s:f")) is not None + self._write_value(cell, value) + self._xp.mark_dirty() + self._extend_dimension(row_idx, col_idx) + if had_formula or self._doc._any_formula_exists(): + self._doc._ensure_full_calc_on_load() + + def append_rows(self, rows: list[list]) -> None: + """Append *rows* after the last existing row. ``None`` entries + leave their cell unstored (sparse), matching Excel semantics.""" + sd = self._sheet_data() + last = self.dimensions[0] + for offset, values in enumerate(rows, start=1): + row_idx = last + offset + row_el = sd.makeelement(qn("s:row"), {"r": str(row_idx)}) + sd.append(row_el) + for col_idx, value in enumerate(values, start=1): + if value is None: + continue + ref = f"{col_index_to_letters(col_idx)}{row_idx}" + c_el = row_el.makeelement(qn("s:c"), {"r": ref}) + row_el.append(c_el) + self._write_value(c_el, value) + if values: + self._extend_dimension(row_idx, len(values)) + else: + self._extend_dimension(row_idx, 1) + if rows: + self._xp.mark_dirty() + if self._doc._any_formula_exists(): + self._doc._ensure_full_calc_on_load() + + # -- internals --------------------------------------------------------------- + + def _find_row(self, row_idx: int) -> "_Element | None": + prev = 0 + for row_el in self._sheet_data(): + if row_el.tag != qn("s:row"): + continue + prev = _row_number(row_el, prev) + if prev == row_idx: + return row_el + if prev > row_idx: + return None # rows sorted by @r + return None + + def _find_cell(self, row_el: "_Element", col_idx: int) -> "_Element | None": + prev = 0 + for c_el in row_el: + if c_el.tag != qn("s:c"): + continue + prev = _cell_column(c_el, prev) + if prev == col_idx: + return c_el + if prev > col_idx: + return None # cells sorted by column + return None + + def _insert_row(self, row_idx: int) -> "_Element": + """Create ```` keeping sheetData sorted by @r + (Excel requires ascending rows).""" + sd = self._sheet_data() + new = sd.makeelement(qn("s:row"), {"r": str(row_idx)}) + prev = 0 + for row_el in sd: + if row_el.tag != qn("s:row"): + continue + prev = _row_number(row_el, prev) + if prev > row_idx: + row_el.addprevious(new) + return new + sd.append(new) + return new + + def _insert_cell( + self, row_el: "_Element", row_idx: int, col_idx: int + ) -> "_Element": + """Create ```` keeping the row's cells column-sorted.""" + ref = f"{col_index_to_letters(col_idx)}{row_idx}" + new = row_el.makeelement(qn("s:c"), {"r": ref}) + prev = 0 + for c_el in row_el: + if c_el.tag != qn("s:c"): + continue + prev = _cell_column(c_el, prev) + if prev > col_idx: + c_el.addprevious(new) + return new + row_el.append(new) + return new + + def _read_value(self, c_el: "_Element") -> object | None: + t = c_el.get("t", "n") + if t == "inlineStr": + is_el = c_el.find(qn("s:is")) + if is_el is None: + return None + return "".join(tel.text or "" for tel in is_el.iter(qn("s:t"))) + v = c_el.find(qn("s:v")) + if v is None or v.text is None: + return None + text = v.text + if t == "s": + strings = self._doc._shared_strings() + try: + return strings[int(text)] + except (ValueError, IndexError): + return None + if t == "b": + return text.strip() in ("1", "true", "TRUE") + if t in ("str", "e"): + return text + # default: numeric + try: + return int(text) + except ValueError: + f = float(text) + return int(f) if f.is_integer() else f + + def _write_value(self, c_el: "_Element", value: object) -> None: + """Replace the cell's content, preserving its style (@s). + + Any existing formula is removed — overwriting means the caller + supplies the value from now on.""" + for tag in ("s:f", "s:v", "s:is"): + el = c_el.find(qn(tag)) + if el is not None: + c_el.remove(el) + if value is None: + c_el.attrib.pop("t", None) + return + if isinstance(value, bool): # before int: bool is an int subclass + c_el.set("t", "b") + v = c_el.makeelement(qn("s:v"), {}) + v.text = "1" if value else "0" + c_el.append(v) + elif isinstance(value, (int, float)): + if isinstance(value, float) and not math.isfinite(value): + raise ValueError(f"Cannot store non-finite float {value!r}") + c_el.attrib.pop("t", None) # default numeric type + v = c_el.makeelement(qn("s:v"), {}) + v.text = repr(value) if isinstance(value, float) else str(value) + c_el.append(v) + elif isinstance(value, str): + c_el.set("t", "inlineStr") + is_el = c_el.makeelement(qn("s:is"), {}) + t_el = is_el.makeelement(qn("s:t"), {_XML_SPACE: "preserve"}) + t_el.text = value + is_el.append(t_el) + c_el.append(is_el) + else: + raise TypeError( + f"Unsupported cell value type {type(value).__name__!r} " + "(str, bool, int, float or None)" + ) + + def _extend_dimension(self, row_idx: int, col_idx: int) -> None: + """Grow the sheet's ```` to include the cell.""" + dim = self._xp.find("s:dimension") + if dim is None: + return + m = _RANGE_RE.match(dim.get("ref") or "") + if m is None: + return + min_c, min_r = col_letters_to_index(m.group(1)), int(m.group(2)) + if m.group(3): + max_c, max_r = col_letters_to_index(m.group(3)), int(m.group(4)) + else: + max_c, max_r = min_c, min_r + new_min_c, new_min_r = min(min_c, col_idx), min(min_r, row_idx) + new_max_c, new_max_r = max(max_c, col_idx), max(max_r, row_idx) + if (new_min_c, new_min_r, new_max_c, new_max_r) != ( + min_c, + min_r, + max_c, + max_r, + ): + dim.set( + "ref", + f"{col_index_to_letters(new_min_c)}{new_min_r}" + f":{col_index_to_letters(new_max_c)}{new_max_r}", + ) + self._xp.mark_dirty() + + +class SheetCollection: + """Mapping-like sheet accessor: by name (``raw.sheets["Sales"]``) or + by position (``raw.sheets[0]``).""" + + def __init__(self, doc: "XlsxRawDocument"): + self._doc = doc + + def __getitem__(self, key: int | str) -> RawSheet: + entries = self._doc._sheet_entries() + if isinstance(key, int): + name, part = entries[key] # IndexError speaks for itself + else: + for name, part in entries: + if name == key: + break + else: + raise KeyError(f"No sheet named {key!r}") + return self._doc._sheet(name, part) + + def __len__(self) -> int: + return len(self._doc._sheet_entries()) + + def __iter__(self) -> Iterator[RawSheet]: + for name, part in self._doc._sheet_entries(): + yield self._doc._sheet(name, part) + + def __contains__(self, name: object) -> bool: + return any(n == name for n, _ in self._doc._sheet_entries()) + + +class XlsxRawDocument(RawDocumentBase): + """Lossless, writable view of an .xlsx workbook.""" + + format = "xlsx" + + def __init__(self, package: "OpcPackage"): + super().__init__(package) + self._workbook_name = self._locate_workbook() + self._sheet_cache: dict[str, RawSheet] = {} + self._shared_cache: list[str] | None = None + self._formulas_exist: bool | None = None + + def _locate_workbook(self) -> str: + rels = self.package.rels_for("") + if rels is not None: + for rel in rels.by_type("/officeDocument"): + return rels.resolve("", rel["target"]) + return "xl/workbook.xml" + + @property + def workbook(self) -> XmlPart: + """The ``xl/workbook.xml`` facade.""" + return self.xml_part(self._workbook_name) + + # -- sheets --------------------------------------------------------------- + + @property + def sheet_names(self) -> list[str]: + """Sheet names in workbook (tab) order.""" + return [s.get("name") or "" for s in self.workbook.findall("s:sheets/s:sheet")] + + @property + def sheets(self) -> SheetCollection: + return SheetCollection(self) + + def _sheet_entries(self) -> list[tuple[str, str]]: + """``(name, worksheet part name)`` in workbook order.""" + rels = self.package.rels_for(self._workbook_name) + entries: list[tuple[str, str]] = [] + for sheet in self.workbook.findall("s:sheets/s:sheet"): + rid = sheet.get(qn("r:id")) + target = rels.target_of(rid) if (rels is not None and rid) else None + if target is None: + continue + entries.append( + (sheet.get("name") or "", rels.resolve(self._workbook_name, target)) + ) + return entries + + def _sheet(self, name: str, part_name: str) -> RawSheet: + if part_name not in self._sheet_cache: + self._sheet_cache[part_name] = RawSheet(self, name, part_name) + return self._sheet_cache[part_name] + + # -- workbook-level reads ---------------------------------------------------- + + @property + def defined_names(self) -> dict[str, str]: + """``definedName`` → refers-to text (read-only view).""" + out: dict[str, str] = {} + for dn in self.workbook.findall("s:definedNames/s:definedName"): + name = dn.get("name") + if name: + out[name] = dn.text or "" + return out + + # -- charts --------------------------------------------------------------- + + @property + def chart_part_names(self) -> list[str]: + """Chart part names reachable from the sheets' drawings, in + sheet order (duplicates removed).""" + from contextifier.raw.chart import find_chart_parts + + names: list[str] = [] + seen: set[str] = set() + for _, ws_name in self._sheet_entries(): + if not self.package.has_part(ws_name): + continue + drawing = self.xml_part(ws_name).find("s:drawing") + if drawing is None: + continue + rid = drawing.get(qn("r:id")) + rels = self.package.rels_for(ws_name) + target = rels.target_of(rid) if (rels is not None and rid) else None + if target is None: + continue + drawing_part = rels.resolve(ws_name, target) + for chart_name in find_chart_parts(self.package, drawing_part): + if chart_name not in seen: + seen.add(chart_name) + names.append(chart_name) + return names + + @property + def charts(self) -> list["ChartModel"]: + """ChartModel per chart part (raises NotImplementedError until + the C3 chart milestone lands; use :attr:`chart_part_names` for + discovery in the meantime).""" + from contextifier.raw.chart import ChartModel + + models: list["ChartModel"] = [] + for name in self.chart_part_names: + try: + models.append(ChartModel(self.xml_part(name), self.package)) + except NotImplementedError: + raise NotImplementedError( + "ChartModel is not implemented yet (milestone C3); " + "chart_part_names still lists the chart parts" + ) from None + return models + + # -- shared strings ------------------------------------------------------------ + + def _shared_strings(self) -> list[str]: + """The shared-string table (read-only, cached). A workbook with + no sharedStrings part simply yields an empty table — writes use + inline strings, so the part is never created.""" + if self._shared_cache is None: + name = "xl/sharedStrings.xml" + rels = self.package.rels_for(self._workbook_name) + if rels is not None: + for rel in rels.by_type("/sharedStrings"): + name = rels.resolve(self._workbook_name, rel["target"]) + break + if not self.package.has_part(name): + self._shared_cache = [] + else: + root = self.xml_part(name).root + self._shared_cache = [ + "".join(t.text or "" for t in si.iter(qn("s:t"))) + for si in root.findall(qn("s:si")) + ] + return self._shared_cache + + # -- formula staleness --------------------------------------------------------- + + def _any_formula_exists(self) -> bool: + """Whether any worksheet contains a formula (cached; used to + decide if edits make cached results stale).""" + if self._formulas_exist is None: + self._formulas_exist = any( + self._part_has_formula(part) for _, part in self._sheet_entries() + ) + return self._formulas_exist + + def _part_has_formula(self, part_name: str) -> bool: + if not self.package.has_part(part_name): + return False + xp = self._xml_parts.get(part_name) + if xp is not None and xp.loaded: + return next(xp.root.iter(qn("s:f")), None) is not None + # Unparsed part: cheap byte scan (an element in the default + # spreadsheetml namespace; false positives merely add calcPr). + data = self.package.get_part(part_name).read() + return b"" in data or b" None: + """Guarantee ```` in workbook.xml so + Excel recalculates stale formula caches on open. No-op (and no + rewrite) when the flag is already set.""" + wb = self.workbook + calc = wb.find("s:calcPr") + if calc is None: + calc = wb.root.makeelement(qn("s:calcPr"), {"fullCalcOnLoad": "1"}) + preceding = {qn(f"s:{tag}") for tag in _PRE_CALC_PR} + pos = 0 + for i, child in enumerate(wb.root): + if child.tag in preceding: + pos = i + 1 + wb.root.insert(pos, calc) + wb.mark_dirty() + elif calc.get("fullCalcOnLoad") != "1": + calc.set("fullCalcOnLoad", "1") + wb.mark_dirty() diff --git a/contextifier/raw/xmlpart.py b/contextifier/raw/xmlpart.py new file mode 100644 index 0000000..0e81a07 --- /dev/null +++ b/contextifier/raw/xmlpart.py @@ -0,0 +1,121 @@ +# contextifier/raw/xmlpart.py +""" +Lazy lxml view over an :class:`~contextifier.raw.opc.OpcPart`. + +Format models hold ``XmlPart`` facades for the parts they understand. +The tree is parsed on first access and serialized back into the package +only when the model mutated it (``mark_dirty()`` + ``flush()``), so +opening a document and reading a few values never rewrites anything — +the byte-preservation contract stays intact by construction. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: # pragma: no cover + from lxml.etree import _Element + + from contextifier.raw.opc import OpcPart + +__all__ = ["NS", "qn", "XmlPart"] + +# The OOXML namespace registry shared by the raw layer. Prefixes follow +# the conventional short names used across ECMA-376 documentation. +NS: dict[str, str] = { + # wordprocessingml + "w": "http://schemas.openxmlformats.org/wordprocessingml/2006/main", + # drawingml + "a": "http://schemas.openxmlformats.org/drawingml/2006/main", + "c": "http://schemas.openxmlformats.org/drawingml/2006/chart", + "cx": "http://schemas.microsoft.com/office/drawing/2014/chartex", + "xdr": "http://schemas.openxmlformats.org/drawingml/2006/spreadsheetDrawing", + "wp": "http://schemas.openxmlformats.org/drawingml/2006/wordprocessingDrawing", + # presentationml + "p": "http://schemas.openxmlformats.org/presentationml/2006/main", + # spreadsheetml + "s": "http://schemas.openxmlformats.org/spreadsheetml/2006/main", + # shared + "r": "http://schemas.openxmlformats.org/officeDocument/2006/relationships", + "ct": "http://schemas.openxmlformats.org/package/2006/content-types", + "pr": "http://schemas.openxmlformats.org/package/2006/relationships", + "mc": "http://schemas.openxmlformats.org/markup-compatibility/2006", +} + + +def qn(tag: str) -> str: + """``"w:p"`` → ``"{http://…/wordprocessingml/2006/main}p"``.""" + prefix, _, local = tag.partition(":") + if not local: + return tag + try: + return f"{{{NS[prefix]}}}{local}" + except KeyError: + raise KeyError(f"Unknown XML namespace prefix: {prefix!r}") from None + + +class XmlPart: + """A parsed XML part with explicit dirty tracking. + + ``root`` parses lazily; call :meth:`mark_dirty` after mutating the + tree and :meth:`flush` (usually via the owning document's ``save``) + to serialize back into the package part. + """ + + __slots__ = ("part", "_root", "_dirty") + + def __init__(self, part: "OpcPart"): + self.part = part + self._root: "_Element | None" = None + self._dirty = False + + @property + def name(self) -> str: + return self.part.name + + @property + def root(self) -> "_Element": + if self._root is None: + from lxml import etree + + self._root = etree.fromstring(self.part.read()) + return self._root + + @property + def loaded(self) -> bool: + return self._root is not None + + @property + def dirty(self) -> bool: + return self._dirty + + def mark_dirty(self) -> None: + self._dirty = True + + def flush(self) -> None: + """Serialize the tree into the part iff this facade dirtied it.""" + if self._dirty and self._root is not None: + from lxml import etree + + self.part.write( + etree.tostring( + self._root, xml_declaration=True, encoding="UTF-8", standalone=True + ) + ) + self._dirty = False + + # -- conveniences ----------------------------------------------------------- + + def find(self, path: str): + """`find` with the shared namespace map (``"s:sheetData/s:row"``).""" + return self.root.find(path, NS) + + def findall(self, path: str): + return self.root.findall(path, NS) + + def iter(self, tag: str): + return self.root.iter(qn(tag)) + + def __repr__(self) -> str: # pragma: no cover - debug aid + state = "dirty" if self._dirty else ("loaded" if self.loaded else "lazy") + return f"" diff --git a/contextifier/services/chart_service.py b/contextifier/services/chart_service.py index e0c3391..54ddc5b 100644 --- a/contextifier/services/chart_service.py +++ b/contextifier/services/chart_service.py @@ -29,10 +29,10 @@ import logging import re -from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple +from typing import TYPE_CHECKING, Dict, List, Optional, Tuple from contextifier.config import ProcessingConfig, ChartConfig -from contextifier.types import ChartData, ChartSeries +from contextifier.types import ChartData if TYPE_CHECKING: from contextifier.services.tag_service import TagService @@ -216,8 +216,10 @@ def _format_as_html_table(self, data: ChartData) -> str: rows.append(f"{''.join(header_cells)}") # Data rows - num_rows = len(data.categories) if data.categories else ( - max((len(s.values) for s in data.series), default=0) + num_rows = ( + len(data.categories) + if data.categories + else (max((len(s.values) for s in data.series), default=0)) ) for i in range(num_rows): cells: List[str] = [] @@ -239,8 +241,7 @@ def _format_as_text_table(self, data: ChartData) -> str: lines: List[str] = [] for i, cat in enumerate(data.categories or []): values = [ - str(s.values[i]) if i < len(s.values) else "" - for s in data.series + str(s.values[i]) if i < len(s.values) else "" for s in data.series ] lines.append(f"{cat}: {', '.join(values)}") diff --git a/contextifier/services/image_service.py b/contextifier/services/image_service.py index 12d4286..3a01d4c 100644 --- a/contextifier/services/image_service.py +++ b/contextifier/services/image_service.py @@ -142,12 +142,14 @@ def save( if size_mb > max_mb: self._logger.warning( "Image skipped: %.2f MB exceeds limit of %.2f MB", - size_mb, max_mb, + size_mb, + max_mb, ) return None should_skip = ( - skip_duplicate if skip_duplicate is not None + skip_duplicate + if skip_duplicate is not None else self._image_config.skip_duplicate ) @@ -280,7 +282,9 @@ def _build_tag(self, file_path: str) -> str: if self._tag_service is not None: return self._tag_service.create_image_tag(file_path) # Fallback: build directly from config (backward compat) - return f"{self._tag_config.image_prefix}{file_path}{self._tag_config.image_suffix}" + return ( + f"{self._tag_config.image_prefix}{file_path}{self._tag_config.image_suffix}" + ) def _generate_filename(self, image_data: bytes) -> str: """Generate a filename using the configured naming strategy.""" diff --git a/contextifier/services/metadata_service.py b/contextifier/services/metadata_service.py index 6152b35..ee27e25 100644 --- a/contextifier/services/metadata_service.py +++ b/contextifier/services/metadata_service.py @@ -76,9 +76,7 @@ def __init__(self, config: ProcessingConfig) -> None: self._logger = logging.getLogger("contextifier.services.metadata") # Select label set - self._labels = ( - _LABELS_KO if self._meta_config.language == "ko" else _LABELS_EN - ) + self._labels = _LABELS_KO if self._meta_config.language == "ko" else _LABELS_EN def format_metadata(self, metadata: Optional[DocumentMetadata]) -> str: """ diff --git a/contextifier/services/storage/local.py b/contextifier/services/storage/local.py index 1e01d33..9074774 100644 --- a/contextifier/services/storage/local.py +++ b/contextifier/services/storage/local.py @@ -36,7 +36,10 @@ def save(self, data: bytes, file_path: str) -> bool: resolved_path = os.path.realpath(file_path) # Path-traversal guard: resolved path must be inside base dir - if not resolved_path.startswith(resolved_base + os.sep) and resolved_path != resolved_base: + if ( + not resolved_path.startswith(resolved_base + os.sep) + and resolved_path != resolved_base + ): raise StorageError( f"Path traversal blocked: {file_path!r} resolves outside base directory", ) diff --git a/contextifier/services/table_service.py b/contextifier/services/table_service.py index 4f9a073..302e132 100644 --- a/contextifier/services/table_service.py +++ b/contextifier/services/table_service.py @@ -17,10 +17,10 @@ import html as html_mod import logging -from typing import List, Optional +from typing import List from contextifier.config import ProcessingConfig, TableConfig -from contextifier.types import OutputFormat, TableCell, TableData +from contextifier.types import OutputFormat, TableData class TableService: diff --git a/contextifier/services/tag_service.py b/contextifier/services/tag_service.py index 0a2df3c..f7fca21 100644 --- a/contextifier/services/tag_service.py +++ b/contextifier/services/tag_service.py @@ -17,7 +17,7 @@ import re import logging -from typing import List, Optional, Tuple +from typing import List, Tuple from contextifier.config import ProcessingConfig, TagConfig from contextifier.types import TagType @@ -67,7 +67,9 @@ def __init__(self, config: ProcessingConfig) -> None: def create_page_tag(self, page_number: int) -> str: """Create a page number tag. E.g., '[Page Number: 1]'""" - return f"{self._tag_config.page_prefix}{page_number}{self._tag_config.page_suffix}" + return ( + f"{self._tag_config.page_prefix}{page_number}{self._tag_config.page_suffix}" + ) def create_slide_tag(self, slide_number: int) -> str: """Create a slide number tag. E.g., '[Slide Number: 1]'""" @@ -141,8 +143,7 @@ def find_slide_tags(self, text: str) -> List[Tuple[int, int, int]]: def find_sheet_tags(self, text: str) -> List[Tuple[int, int, str]]: """Find all sheet tags. Returns [(start, end, sheet_name), ...]""" return [ - (m.start(), m.end(), m.group(1)) - for m in self._sheet_pattern.finditer(text) + (m.start(), m.end(), m.group(1)) for m in self._sheet_pattern.finditer(text) ] def has_page_markers(self, text: str) -> bool: diff --git a/contextifier/types.py b/contextifier/types.py index c376f72..2f61ba1 100644 --- a/contextifier/types.py +++ b/contextifier/types.py @@ -32,26 +32,29 @@ # File & Format Types # ═══════════════════════════════════════════════════════════════════════════════ + @unique class FileCategory(str, Enum): """Classification of file types by their nature.""" - DOCUMENT = "document" # PDF, DOCX, DOC, RTF, HWP, HWPX + + DOCUMENT = "document" # PDF, DOCX, DOC, RTF, HWP, HWPX PRESENTATION = "presentation" # PPT, PPTX - SPREADSHEET = "spreadsheet" # XLSX, XLS - TEXT = "text" # TXT, MD - CODE = "code" # PY, JS, TS, Java, ... - CONFIG = "config" # JSON, YAML, XML, TOML, ... - DATA = "data" # CSV, TSV - SCRIPT = "script" # SH, BAT, PS1, ... - LOG = "log" # LOG - WEB = "web" # HTML, HTM, XHTML - IMAGE = "image" # JPG, PNG, GIF, ... + SPREADSHEET = "spreadsheet" # XLSX, XLS + TEXT = "text" # TXT, MD + CODE = "code" # PY, JS, TS, Java, ... + CONFIG = "config" # JSON, YAML, XML, TOML, ... + DATA = "data" # CSV, TSV + SCRIPT = "script" # SH, BAT, PS1, ... + LOG = "log" # LOG + WEB = "web" # HTML, HTM, XHTML + IMAGE = "image" # JPG, PNG, GIF, ... UNKNOWN = "unknown" @unique class OutputFormat(str, Enum): """Output format for table rendering.""" + HTML = "html" MARKDOWN = "markdown" TEXT = "text" @@ -60,6 +63,7 @@ class OutputFormat(str, Enum): @unique class ImageFormat(str, Enum): """Supported image formats.""" + PNG = "png" JPEG = "jpeg" JPG = "jpg" @@ -73,15 +77,17 @@ class ImageFormat(str, Enum): @unique class NamingStrategy(str, Enum): """Strategy for naming saved files.""" - HASH = "hash" # Content-based hash (deduplication) - UUID = "uuid" # Random UUID + + HASH = "hash" # Content-based hash (deduplication) + UUID = "uuid" # Random UUID SEQUENTIAL = "sequential" # Counter-based - TIMESTAMP = "timestamp" # Time-based + TIMESTAMP = "timestamp" # Time-based @unique class StorageType(str, Enum): """Storage backend types.""" + LOCAL = "local" MINIO = "minio" S3 = "s3" @@ -92,6 +98,7 @@ class StorageType(str, Enum): @unique class TagType(str, Enum): """Types of structural tags in extracted text.""" + PAGE = "page" SLIDE = "slide" SHEET = "sheet" @@ -100,6 +107,7 @@ class TagType(str, Enum): @unique class PipelineStage(str, Enum): """Named stages in the processing pipeline.""" + CONVERT = "convert" PREPROCESS = "preprocess" EXTRACT_METADATA = "extract_metadata" @@ -111,6 +119,7 @@ class PipelineStage(str, Enum): # File Data Structures # ═══════════════════════════════════════════════════════════════════════════════ + class FileContext(TypedDict): """ Standardized file input for all handlers. @@ -122,22 +131,25 @@ class FileContext(TypedDict): ``file_stream`` may be ``None`` to save memory; use ``BaseConverter._get_stream()`` to obtain a seekable stream. """ - file_path: str # Absolute path to the original file - file_name: str # Filename with extension - file_extension: str # Lowercase extension without dot - file_category: str # FileCategory value - file_data: bytes # Raw binary content + + file_path: str # Absolute path to the original file + file_name: str # Filename with extension + file_extension: str # Lowercase extension without dot + file_category: str # FileCategory value + file_data: bytes # Raw binary content file_stream: Optional[io.BytesIO] # Seekable binary stream (may be None) - file_size: int # Size in bytes + file_size: int # Size in bytes # ═══════════════════════════════════════════════════════════════════════════════ # Document Metadata # ═══════════════════════════════════════════════════════════════════════════════ + @unique class MetadataField(str, Enum): """Standard metadata field names.""" + TITLE = "title" SUBJECT = "subject" AUTHOR = "author" @@ -160,6 +172,7 @@ class DocumentMetadata: Every handler's MetadataExtractor produces this same structure, ensuring uniform metadata handling regardless of source format. """ + title: Optional[str] = None subject: Optional[str] = None author: Optional[str] = None @@ -178,9 +191,18 @@ def to_dict(self) -> Dict[str, Any]: """Convert to dictionary, omitting None values.""" result: Dict[str, Any] = {} for f in [ - "title", "subject", "author", "keywords", "comments", - "last_saved_by", "create_time", "last_saved_time", - "page_count", "word_count", "category", "revision", + "title", + "subject", + "author", + "keywords", + "comments", + "last_saved_by", + "create_time", + "last_saved_time", + "page_count", + "word_count", + "category", + "revision", ]: val = getattr(self, f) if val is not None: @@ -197,8 +219,16 @@ def from_dict(cls, data: Dict[str, Any]) -> "DocumentMetadata": """Create from dictionary.""" kwargs: Dict[str, Any] = {} for f in [ - "title", "subject", "author", "keywords", "comments", - "last_saved_by", "page_count", "word_count", "category", "revision", + "title", + "subject", + "author", + "keywords", + "comments", + "last_saved_by", + "page_count", + "word_count", + "category", + "revision", ]: if f in data: kwargs[f] = data[f] @@ -219,9 +249,18 @@ def from_dict(cls, data: Dict[str, Any]) -> "DocumentMetadata": def is_empty(self) -> bool: """Check if all fields are empty/None.""" for f in [ - "title", "subject", "author", "keywords", "comments", - "last_saved_by", "create_time", "last_saved_time", - "page_count", "word_count", "category", "revision", + "title", + "subject", + "author", + "keywords", + "comments", + "last_saved_by", + "create_time", + "last_saved_time", + "page_count", + "word_count", + "category", + "revision", ]: if getattr(self, f) is not None: return False @@ -232,9 +271,11 @@ def is_empty(self) -> bool: # Table Data Structures # ═══════════════════════════════════════════════════════════════════════════════ + @dataclass class TableCell: """Single cell in a table.""" + content: str = "" row_span: int = 1 col_span: int = 1 @@ -253,6 +294,7 @@ class TableData: using the same structure — PDF tables, DOCX tables, Excel sheets, CSV data all normalize to this representation. """ + rows: List[List[TableCell]] = field(default_factory=list) num_rows: int = 0 num_cols: int = 0 @@ -266,9 +308,11 @@ class TableData: # Chart Data Structures # ═══════════════════════════════════════════════════════════════════════════════ + @dataclass class ChartSeries: """One data series in a chart.""" + name: Optional[str] = None values: List[Any] = field(default_factory=list) @@ -280,6 +324,7 @@ class ChartData: Whether from DOCX, PPTX, XLSX, or HWP — chart data normalizes here. """ + chart_type: Optional[str] = None title: Optional[str] = None categories: List[str] = field(default_factory=list) @@ -291,6 +336,7 @@ class ChartData: # Extraction Result # ═══════════════════════════════════════════════════════════════════════════════ + @dataclass class ExtractionResult: """ @@ -299,6 +345,7 @@ class ExtractionResult: Every handler returns this same structure, making downstream consumers (chunking, OCR, user code) format-agnostic. """ + text: str = "" metadata: Optional[DocumentMetadata] = None tables: List[TableData] = field(default_factory=list) @@ -328,9 +375,11 @@ def has_images(self) -> bool: # Chunk Result # ═══════════════════════════════════════════════════════════════════════════════ + @dataclass class ChunkMetadata: """Position metadata for a single chunk.""" + chunk_index: int = 0 page_number: Optional[int] = None line_start: int = 0 @@ -345,6 +394,7 @@ class ChunkMetadata: @dataclass class Chunk: """A single text chunk with optional position metadata.""" + text: str = "" metadata: Optional[ChunkMetadata] = None @@ -353,6 +403,7 @@ class Chunk: # Preprocessed Data # ═══════════════════════════════════════════════════════════════════════════════ + @dataclass class PreprocessedData: """ @@ -361,11 +412,16 @@ class PreprocessedData: The preprocessor transforms the converted format object into cleaned data ready for content extraction. """ - content: Any = None # Primary processed content - raw_content: Any = None # Original input (for reference) - encoding: str = "utf-8" # Detected encoding - resources: Dict[str, Any] = field(default_factory=dict) # Extracted resources (images, etc.) - properties: Dict[str, Any] = field(default_factory=dict) # Discovered properties during preprocessing + + content: Any = None # Primary processed content + raw_content: Any = None # Original input (for reference) + encoding: str = "utf-8" # Detected encoding + resources: Dict[str, Any] = field( + default_factory=dict + ) # Extracted resources (images, etc.) + properties: Dict[str, Any] = field( + default_factory=dict + ) # Discovered properties during preprocessing # ═══════════════════════════════════════════════════════════════════════════════ @@ -380,14 +436,47 @@ class PreprocessedData: FileCategory.PRESENTATION: frozenset(["ppt", "pptx"]), FileCategory.SPREADSHEET: frozenset(["xlsx", "xls"]), FileCategory.TEXT: frozenset(["txt", "md", "markdown"]), - FileCategory.CODE: frozenset([ - "py", "js", "ts", "java", "cpp", "c", "h", "cs", "go", "rs", - "php", "rb", "swift", "kt", "scala", "dart", "r", "sql", - "css", "jsx", "tsx", "vue", "svelte", - ]), - FileCategory.CONFIG: frozenset([ - "json", "yaml", "yml", "xml", "toml", "ini", "cfg", "conf", "properties", "env", - ]), + FileCategory.CODE: frozenset( + [ + "py", + "js", + "ts", + "java", + "cpp", + "c", + "h", + "cs", + "go", + "rs", + "php", + "rb", + "swift", + "kt", + "scala", + "dart", + "r", + "sql", + "css", + "jsx", + "tsx", + "vue", + "svelte", + ] + ), + FileCategory.CONFIG: frozenset( + [ + "json", + "yaml", + "yml", + "xml", + "toml", + "ini", + "cfg", + "conf", + "properties", + "env", + ] + ), FileCategory.DATA: frozenset(["csv", "tsv"]), FileCategory.SCRIPT: frozenset(["sh", "bat", "ps1", "zsh", "fish"]), FileCategory.LOG: frozenset(["log"]), diff --git a/pyproject.toml b/pyproject.toml index 85a0de4..b19a3b5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "contextifier" -version = "0.3.0" +version = "0.4.0" description = "Convert raw documents into AI-understandable context with intelligent text extraction, table detection, and semantic chunking" readme = "README.md" requires-python = ">=3.12" @@ -50,6 +50,7 @@ dependencies = [ "pdfminer.six>=20231228", "pdf2image>=1.17.0", # Office — DOCX / PPTX / Excel + "lxml>=5.0", "python-docx>=1.1.0", "docx2pdf>=0.1.8", "python-pptx>=1.0.0", diff --git a/tests/conftest.py b/tests/conftest.py index 29741df..576e60e 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -11,9 +11,6 @@ from __future__ import annotations -import io -import os -import tempfile from pathlib import Path from typing import Any, Dict, Optional from unittest.mock import MagicMock diff --git a/tests/integration/test_document_processor.py b/tests/integration/test_document_processor.py index 2e377d3..736c72e 100644 --- a/tests/integration/test_document_processor.py +++ b/tests/integration/test_document_processor.py @@ -7,8 +7,6 @@ from __future__ import annotations -import os -import tempfile from pathlib import Path import pytest diff --git a/tests/regression/test_bug_fixes.py b/tests/regression/test_bug_fixes.py index 763fd0d..3cd5216 100644 --- a/tests/regression/test_bug_fixes.py +++ b/tests/regression/test_bug_fixes.py @@ -7,7 +7,6 @@ from __future__ import annotations -import html as html_mod import pytest @@ -102,7 +101,7 @@ class TestP1_HandlerFinalMethods: def test_process_is_final(self) -> None: from contextifier.handlers.base import BaseHandler - process_method = getattr(BaseHandler.process, "__final__", None) + getattr(BaseHandler.process, "__final__", None) # @final sets __final__ attribute (Python 3.11+) or we check via typing # The decorator is applied — we verify it doesn't raise at import assert BaseHandler.process is not None diff --git a/tests/service/knowledge/__init__.py b/tests/service/knowledge/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/unit/chunking/test_chunker.py b/tests/unit/chunking/test_chunker.py index cc32836..dfadf52 100644 --- a/tests/unit/chunking/test_chunker.py +++ b/tests/unit/chunking/test_chunker.py @@ -7,7 +7,6 @@ from contextifier.config import ProcessingConfig from contextifier.chunking.chunker import TextChunker -from contextifier.errors import ChunkingError @pytest.fixture() diff --git a/tests/unit/handlers/test_configurable_thresholds.py b/tests/unit/handlers/test_configurable_thresholds.py index ab6a24e..35db586 100644 --- a/tests/unit/handlers/test_configurable_thresholds.py +++ b/tests/unit/handlers/test_configurable_thresholds.py @@ -125,7 +125,7 @@ def test_render_dpi_used_in_scan(self, mock_fitz): ext._image_service.save_and_tag.return_value = "[Image: scan.png]" # Call _extract_scan_pages - result = ext._extract_scan_pages(mock_doc) + ext._extract_scan_pages(mock_doc) # Verify Matrix was called with custom DPI zoom (300/72 ≈ 4.167) zoom = 300 / 72.0 diff --git a/tests/unit/handlers/test_csv_delimiter_confidence.py b/tests/unit/handlers/test_csv_delimiter_confidence.py index b768516..6cd5dd4 100644 --- a/tests/unit/handlers/test_csv_delimiter_confidence.py +++ b/tests/unit/handlers/test_csv_delimiter_confidence.py @@ -3,7 +3,6 @@ from __future__ import annotations -import pytest from contextifier.handlers.csv.preprocessor import ( CsvParsedData, diff --git a/tests/unit/handlers/test_delegation.py b/tests/unit/handlers/test_delegation.py index bc4ac50..a138128 100644 --- a/tests/unit/handlers/test_delegation.py +++ b/tests/unit/handlers/test_delegation.py @@ -3,7 +3,7 @@ from __future__ import annotations -from unittest.mock import MagicMock, patch +from unittest.mock import MagicMock import pytest diff --git a/tests/unit/handlers/test_doc_fib.py b/tests/unit/handlers/test_doc_fib.py index 229f603..820d055 100644 --- a/tests/unit/handlers/test_doc_fib.py +++ b/tests/unit/handlers/test_doc_fib.py @@ -13,14 +13,10 @@ from __future__ import annotations import struct -import pytest from contextifier.handlers.doc._fib import ( - PieceDescriptor, _clean_doc_text, _parse_clx, - _parse_plc_pcd, - _read_pieces, detect_tables_from_text, parse_fib_text, ) @@ -87,7 +83,7 @@ def _build_clx_unicode(text: str, body_offset: int = 0x0200) -> tuple[bytes, int Returns (table_stream_bytes, fc_clx, lcb_clx). """ - text_bytes = text.encode("utf-16-le") + text.encode("utf-16-le") char_count = len(text) # Build PlcPcd: 2 CPs + 1 PCD @@ -152,7 +148,6 @@ def test_single_compressed_piece(self): text = "Hello World" table_stream, fc_clx, lcb_clx = _build_clx_compressed(text) - body_offset = 0x0200 word_data = _build_word_data( ccp_text=len(text), fc_clx=fc_clx, @@ -168,7 +163,6 @@ def test_single_unicode_piece(self): text = "한글 테스트" table_stream, fc_clx, lcb_clx = _build_clx_unicode(text) - body_offset = 0x0200 word_data = _build_word_data( ccp_text=len(text), fc_clx=fc_clx, diff --git a/tests/unit/handlers/test_html_handler.py b/tests/unit/handlers/test_html_handler.py index c90ba9a..89db3c6 100644 --- a/tests/unit/handlers/test_html_handler.py +++ b/tests/unit/handlers/test_html_handler.py @@ -11,7 +11,6 @@ from contextifier.handlers.html.converter import HtmlConvertedData from contextifier.handlers.html.preprocessor import ( HtmlPreprocessor, - _MAX_IMAGE_DECODE_BYTES, ) diff --git a/tests/unit/handlers/test_hwpx_charts.py b/tests/unit/handlers/test_hwpx_charts.py index 8dd5025..ba160d0 100644 --- a/tests/unit/handlers/test_hwpx_charts.py +++ b/tests/unit/handlers/test_hwpx_charts.py @@ -12,18 +12,13 @@ import io import zipfile -import xml.etree.ElementTree as ET -from unittest import mock -import pytest from contextifier.handlers.hwpx._section import ( _parse_ooxml_chart, - _process_chart_ref, _format_chart_simple, parse_hwpx_section, ) -from contextifier.handlers.hwpx._constants import OOXML_CHART_NS # ═══════════════════════════════════════════════════════════════════════════ diff --git a/tests/unit/handlers/test_pdf_scan_ocr.py b/tests/unit/handlers/test_pdf_scan_ocr.py index 5ee66e9..83281e3 100644 --- a/tests/unit/handlers/test_pdf_scan_ocr.py +++ b/tests/unit/handlers/test_pdf_scan_ocr.py @@ -3,7 +3,7 @@ from __future__ import annotations -from unittest.mock import MagicMock, patch +from unittest.mock import MagicMock import pytest diff --git a/tests/unit/handlers/test_ppt_tables_charts.py b/tests/unit/handlers/test_ppt_tables_charts.py index 12558c0..64c0d66 100644 --- a/tests/unit/handlers/test_ppt_tables_charts.py +++ b/tests/unit/handlers/test_ppt_tables_charts.py @@ -15,7 +15,6 @@ from typing import List, Tuple from unittest.mock import MagicMock -import pytest from contextifier.handlers.ppt.content_extractor import ( PptContentExtractor, @@ -23,7 +22,7 @@ _detect_tabular_text, _rows_to_table_data, ) -from contextifier.types import PreprocessedData, TableCell +from contextifier.types import PreprocessedData # ═════════════════════════════════════════════════════════════════════════════ diff --git a/tests/unit/handlers/test_pptx_group_shapes.py b/tests/unit/handlers/test_pptx_group_shapes.py index 7e57d47..d19dbd5 100644 --- a/tests/unit/handlers/test_pptx_group_shapes.py +++ b/tests/unit/handlers/test_pptx_group_shapes.py @@ -10,10 +10,9 @@ from unittest import mock -import pytest from contextifier.handlers.pptx.content_extractor import PptxContentExtractor -from contextifier.types import ChartData, ChartSeries, PreprocessedData, TableData, TableCell +from contextifier.types import ChartData, PreprocessedData def _make_shape( diff --git a/tests/unit/handlers/test_rtf_table_merge.py b/tests/unit/handlers/test_rtf_table_merge.py index 228a77d..9fafaf5 100644 --- a/tests/unit/handlers/test_rtf_table_merge.py +++ b/tests/unit/handlers/test_rtf_table_merge.py @@ -9,20 +9,16 @@ from __future__ import annotations -import pytest from unittest import mock from contextifier.handlers.rtf._table_parser import ( _parse_cell_definitions, - _extract_cells_with_merge, _build_table_data, _is_real_table, _ParsedCell, - _CellDef, extract_tables, single_column_to_text, ) -from contextifier.types import TableCell, TableData # ═══════════════════════════════════════════════════════════════════════════ diff --git a/tests/unit/handlers/test_text_handler.py b/tests/unit/handlers/test_text_handler.py index 27b883a..273f55f 100644 --- a/tests/unit/handlers/test_text_handler.py +++ b/tests/unit/handlers/test_text_handler.py @@ -4,7 +4,6 @@ from __future__ import annotations import pytest -from unittest.mock import MagicMock from contextifier.config import ProcessingConfig from contextifier.handlers.text.handler import TextHandler diff --git a/tests/unit/handlers/test_tpe_reuse.py b/tests/unit/handlers/test_tpe_reuse.py index c313f3b..4723e1e 100644 --- a/tests/unit/handlers/test_tpe_reuse.py +++ b/tests/unit/handlers/test_tpe_reuse.py @@ -12,7 +12,6 @@ from __future__ import annotations import unittest -from unittest.mock import MagicMock, patch from contextifier.handlers.base import BaseHandler from contextifier.config import ProcessingConfig diff --git a/tests/unit/handlers/test_xls_images_charts.py b/tests/unit/handlers/test_xls_images_charts.py index da60b58..3728138 100644 --- a/tests/unit/handlers/test_xls_images_charts.py +++ b/tests/unit/handlers/test_xls_images_charts.py @@ -11,10 +11,8 @@ from __future__ import annotations -import io from unittest.mock import MagicMock, patch -import pytest from contextifier.handlers.xls.content_extractor import ( XlsContentExtractor, @@ -156,7 +154,7 @@ def test_deduplication(self, mock_olefile_mod): resources={"file_data": b"\xd0\xcf" + b"\x00" * 100}, ) extractor = XlsContentExtractor(image_service=image_service) - tags = extractor.extract_images(preprocessed) + extractor.extract_images(preprocessed) # Should only save once (dedup by md5) assert image_service.save_and_tag.call_count == 1 diff --git a/tests/unit/handlers/test_xlsx_charts.py b/tests/unit/handlers/test_xlsx_charts.py index a4634e6..bbe6f1b 100644 --- a/tests/unit/handlers/test_xlsx_charts.py +++ b/tests/unit/handlers/test_xlsx_charts.py @@ -10,7 +10,6 @@ from unittest import mock -import pytest from contextifier.handlers.xlsx.content_extractor import XlsxContentExtractor from contextifier.types import ChartData, ChartSeries, PreprocessedData diff --git a/tests/unit/handlers/test_zip_bomb_defense.py b/tests/unit/handlers/test_zip_bomb_defense.py index 43bfbd2..650c7bc 100644 --- a/tests/unit/handlers/test_zip_bomb_defense.py +++ b/tests/unit/handlers/test_zip_bomb_defense.py @@ -8,7 +8,7 @@ import pytest -from contextifier.pipeline.converter import check_zip_bomb, MAX_ZIP_DECOMPRESSED_BYTES +from contextifier.pipeline.converter import check_zip_bomb from contextifier.errors import ConversionError diff --git a/tests/unit/integrations/test_langchain_loader.py b/tests/unit/integrations/test_langchain_loader.py index d9d480c..ed6d707 100644 --- a/tests/unit/integrations/test_langchain_loader.py +++ b/tests/unit/integrations/test_langchain_loader.py @@ -1,8 +1,6 @@ # tests/unit/integrations/test_langchain_loader.py """Tests for ContextifierLoader — LangChain integration.""" -import os -import tempfile from unittest.mock import MagicMock, patch import pytest @@ -11,7 +9,6 @@ from contextifier.integrations.langchain_loader import ContextifierLoader from contextifier.config import ProcessingConfig from contextifier.document_processor import ChunkResult -from contextifier.types import Chunk, ChunkMetadata # ── Fixtures ────────────────────────────────────────────────────────────── diff --git a/tests/unit/ocr/test_ocr_parallel.py b/tests/unit/ocr/test_ocr_parallel.py index a856d9a..558fec7 100644 --- a/tests/unit/ocr/test_ocr_parallel.py +++ b/tests/unit/ocr/test_ocr_parallel.py @@ -14,7 +14,6 @@ import os import tempfile -import time import unittest from unittest.mock import MagicMock diff --git a/tests/unit/ocr/test_ocr_prompt_language.py b/tests/unit/ocr/test_ocr_prompt_language.py index 6062447..b2dca4b 100644 --- a/tests/unit/ocr/test_ocr_prompt_language.py +++ b/tests/unit/ocr/test_ocr_prompt_language.py @@ -10,7 +10,6 @@ from __future__ import annotations -import pytest from contextifier.config import OCRConfig from contextifier.ocr.base import ( diff --git a/tests/unit/ocr/test_tesseract_engine.py b/tests/unit/ocr/test_tesseract_engine.py index 860e16e..16c94ce 100644 --- a/tests/unit/ocr/test_tesseract_engine.py +++ b/tests/unit/ocr/test_tesseract_engine.py @@ -5,11 +5,9 @@ from __future__ import annotations -import os import sys from unittest import mock -import pytest from contextifier.ocr.engines.tesseract_engine import TesseractOCREngine diff --git a/tests/unit/raw/__init__.py b/tests/unit/raw/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/unit/raw/test_chart_model.py b/tests/unit/raw/test_chart_model.py new file mode 100644 index 0000000..cb14fad --- /dev/null +++ b/tests/unit/raw/test_chart_model.py @@ -0,0 +1,462 @@ +# tests/unit/raw/test_chart_model.py +"""ChartModel (C3) — chart read/write over real Office files. + +Every test works on genuine packages: charts authored by openpyxl +(xlsx-hosted, no caches) and python-pptx (pptx-hosted, full caches + +embedded workbook), plus hand-built packages for the chartEx and +missing-cache corners. Write paths are verified by REOPENING the saved +bytes with python-pptx — proving the rewrite is Office-readable, not +just self-consistent. +""" + +from __future__ import annotations + +import io +import re +import zipfile + +import pytest +from openpyxl import Workbook, load_workbook +from openpyxl.chart import BarChart, LineChart, Reference +from pptx import Presentation +from pptx.chart.data import CategoryChartData, XyChartData +from pptx.enum.chart import XL_CHART_TYPE +from pptx.util import Inches + +from contextifier.raw.chart import ChartSeriesData, load_chart +from contextifier.raw.opc import OpcPackage, RawUnsupportedError + +_CHART_PART_RE = re.compile(r"(?:xl|ppt|word)/charts/chart\d+\.xml$") + + +def _chart_part_names(pkg: OpcPackage) -> list[str]: + return sorted(n for n in pkg.part_names if _CHART_PART_RE.fullmatch(n)) + + +# -- builders: real files ------------------------------------------------------ + + +def _xlsx_with_charts() -> bytes: + """openpyxl workbook hosting a BarChart + a LineChart.""" + wb = Workbook() + ws = wb.active + for row in [["Cat", "S1", "S2"], ["A", 1, 4], ["B", 2, 5], ["C", 3, 6]]: + ws.append(row) + data = Reference(ws, min_col=2, min_row=1, max_col=3, max_row=4) + cats = Reference(ws, min_col=1, min_row=2, max_row=4) + + bar = BarChart() + bar.type = "col" + bar.title = "Bar Title" + bar.add_data(data, titles_from_data=True) + bar.set_categories(cats) + ws.add_chart(bar, "E5") + + line = LineChart() + line.title = "Line Title" + line.add_data(data, titles_from_data=True) + line.set_categories(cats) + ws.add_chart(line, "E20") + + buf = io.BytesIO() + wb.save(buf) + return buf.getvalue() + + +def _pptx_with_charts(specs) -> bytes: + """One slide + chart per (chart_type, categories, [(name, values)...]).""" + prs = Presentation() + for chart_type, cats, series in specs: + slide = prs.slides.add_slide(prs.slide_layouts[6]) + chart_data = CategoryChartData() + chart_data.categories = cats + for name, values in series: + chart_data.add_series(name, values) + slide.shapes.add_chart( + chart_type, Inches(1), Inches(1), Inches(6), Inches(4), chart_data + ) + buf = io.BytesIO() + prs.save(buf) + return buf.getvalue() + + +def _pptx_charts(data: bytes) -> list: + """Reopen saved bytes with python-pptx and collect the chart objects.""" + prs = Presentation(io.BytesIO(data)) + return [ + shape.chart for slide in prs.slides for shape in slide.shapes if shape.has_chart + ] + + +def _mini_package(parts: dict[str, bytes | str]) -> bytes: + """Hand-built OPC zip: content types + root rels + given parts.""" + buf = io.BytesIO() + with zipfile.ZipFile(buf, "w", zipfile.ZIP_DEFLATED) as z: + z.writestr( + "[Content_Types].xml", + '' + '' + '' + '' + "", + ) + z.writestr( + "_rels/.rels", + '' + '', + ) + for name, content in parts.items(): + z.writestr(name, content) + return buf.getvalue() + + +_CHARTEX_XML = ( + '' + '' + "" + '' + 'Sheet1!$A$2:$A$4' + 'AB' + 'C' + 'Sheet1!$B$2:$B$4' + '13' + "" + "" + "" + "" + "" + "Funnel Title" + "" + "" + '' + "Sheet1!$B$1S1" + '' + "" + "" + "" + "" +) + +_NO_CACHE_CHART_XML = ( + '' + '' + "" + '' + # ser 0: openpyxl-style — bare c:f references, no caches anywhere + "" + '' + "Sheet1!$B$1" + "Sheet1!$A$2:$A$4" + "Sheet1!$B$2:$B$4" + "" + # ser 1: numCache shell with a ptCount but zero c:pt children + "" + '' + "Sheet1!$C$2:$C$4" + '' + "" + "" + "" + "" +) + + +# -- 1. openpyxl-hosted reading ------------------------------------------------- + + +class TestOpenpyxlHostedReading: + def test_kind_title_and_series_shape(self): + pkg = OpcPackage.open(_xlsx_with_charts()) + names = _chart_part_names(pkg) + assert len(names) == 2 + + models = {m.kind: m for m in (load_chart(pkg, n) for n in names)} + assert set(models) == {"column", "line"} + assert models["column"].title == "Bar Title" + assert models["line"].title == "Line Title" + + for model in models.values(): + series = model.series + assert len(series) == 2 + # openpyxl writes bare c:f references into the host workbook, + # never caches — the documented cache-based read is empty. + for ser in series: + assert ser.name is None + assert ser.categories == [] + assert ser.values == [] + + +# -- 2. python-pptx-hosted read + set_data round-trip --------------------------- + + +class TestPptxSetDataRoundTrip: + def _package(self) -> tuple[OpcPackage, dict[str, str]]: + src = _pptx_with_charts( + [ + ( + XL_CHART_TYPE.COLUMN_CLUSTERED, + ["East", "West", "Mid"], + [("Q1", (19.2, 21.4, 16.7)), ("Q2", (22.3, 28.6, 15.2))], + ), + (XL_CHART_TYPE.PIE, ["A", "B", "C"], [("Share", (1, 2, 3))]), + ] + ) + pkg = OpcPackage.open(src) + by_kind = {load_chart(pkg, n).kind: n for n in _chart_part_names(pkg)} + assert set(by_kind) == {"column", "pie"} + return pkg, by_kind + + def test_reads_pptx_caches(self): + pkg, by_kind = self._package() + series = load_chart(pkg, by_kind["column"]).series + assert [s.name for s in series] == ["Q1", "Q2"] + assert series[0].categories == ["East", "West", "Mid"] + assert series[0].values == pytest.approx([19.2, 21.4, 16.7]) + assert series[1].values == pytest.approx([22.3, 28.6, 15.2]) + + def test_column_set_data_reopens_in_office_model(self): + pkg, by_kind = self._package() + model = load_chart(pkg, by_kind["column"]) + new_cats = ["N", "S", "E", "W"] + model.set_data(new_cats, [("Alpha", [1, 2, 3, 4]), ("Beta", [5, 6, 7, 8])]) + out = pkg.to_bytes() + + chart = next( + c + for c in _pptx_charts(out) + if c.chart_type == XL_CHART_TYPE.COLUMN_CLUSTERED + ) + assert list(chart.plots[0].categories) == new_cats + assert [(s.name, list(s.values)) for s in chart.series] == [ + ("Alpha", [1, 2, 3, 4]), + ("Beta", [5, 6, 7, 8]), + ] + + def test_pie_set_data_reopens_in_office_model(self): + pkg, by_kind = self._package() + model = load_chart(pkg, by_kind["pie"]) + model.set_data(["Red", "Blue"], [("Split", [60, 40])]) + out = pkg.to_bytes() + + chart = next(c for c in _pptx_charts(out) if c.chart_type == XL_CHART_TYPE.PIE) + assert list(chart.plots[0].categories) == ["Red", "Blue"] + assert [(s.name, list(s.values)) for s in chart.series] == [("Split", [60, 40])] + + def test_embedded_workbook_is_regenerated(self): + pkg, by_kind = self._package() + model = load_chart(pkg, by_kind["column"]) + workbook_part = model.embedded_workbook_part() + assert workbook_part is not None + assert workbook_part.name.endswith(".xlsx") + + model.set_data(["N", "S"], [("Alpha", [1, 2]), ("Beta", [5, 6])]) + out = pkg.to_bytes() + + reopened = OpcPackage.open(out) + wb_bytes = load_chart(reopened, by_kind["column"]).embedded_workbook_part() + wb = load_workbook(io.BytesIO(wb_bytes.read())) + rows = list(wb["Sheet1"].values) + assert rows == [ + (None, "Alpha", "Beta"), + ("N", 1, 5), + ("S", 2, 6), + ] + + +# -- 3. set_title round-trip ------------------------------------------------------ + + +class TestSetTitle: + def test_set_title_roundtrip_via_pptx(self): + src = _pptx_with_charts( + [(XL_CHART_TYPE.COLUMN_CLUSTERED, ["A", "B"], [("S", (1, 2))])] + ) + pkg = OpcPackage.open(src) + name = _chart_part_names(pkg)[0] + model = load_chart(pkg, name) + assert model.title is None # python-pptx writes no explicit title + + model.set_title("First Draft") # create path + model.set_title("Quarterly Revenue") # replace path + assert model.title == "Quarterly Revenue" + out = pkg.to_bytes() + + chart = _pptx_charts(out)[0] + assert chart.has_title + assert chart.chart_title.text_frame.text == "Quarterly Revenue" + assert load_chart(OpcPackage.open(out), name).title == "Quarterly Revenue" + + +# -- 4. series count growth / shrink --------------------------------------------- + + +class TestSeriesCountChanges: + def test_growth_one_to_three(self): + src = _pptx_with_charts( + [(XL_CHART_TYPE.COLUMN_CLUSTERED, ["A", "B", "C"], [("Only", (1, 2, 3))])] + ) + pkg = OpcPackage.open(src) + name = _chart_part_names(pkg)[0] + load_chart(pkg, name).set_data( + ["X", "Y"], + [("S1", [1, 2]), ("S2", [3, 4]), ("S3", [5, 6])], + ) + out = pkg.to_bytes() + + chart = _pptx_charts(out)[0] + assert list(chart.plots[0].categories) == ["X", "Y"] + assert [(s.name, list(s.values)) for s in chart.series] == [ + ("S1", [1, 2]), + ("S2", [3, 4]), + ("S3", [5, 6]), + ] + + def test_shrink_three_to_one(self): + src = _pptx_with_charts( + [ + ( + XL_CHART_TYPE.COLUMN_CLUSTERED, + ["A", "B"], + [("S1", (1, 2)), ("S2", (3, 4)), ("S3", (5, 6))], + ) + ] + ) + pkg = OpcPackage.open(src) + name = _chart_part_names(pkg)[0] + # ChartSeriesData input form (its .categories are ignored). + load_chart(pkg, name).set_data( + ["X", "Y"], [ChartSeriesData("Solo", values=[9, 8])] + ) + out = pkg.to_bytes() + + chart = _pptx_charts(out)[0] + assert list(chart.plots[0].categories) == ["X", "Y"] + assert [(s.name, list(s.values)) for s in chart.series] == [("Solo", [9, 8])] + + +# -- 5. validation + chartEx ------------------------------------------------------ + + +class TestValidationAndChartEx: + def test_ragged_series_raises_value_error(self): + src = _pptx_with_charts( + [(XL_CHART_TYPE.COLUMN_CLUSTERED, ["A", "B"], [("S", (1, 2))])] + ) + pkg = OpcPackage.open(src) + model = load_chart(pkg, _chart_part_names(pkg)[0]) + with pytest.raises(ValueError, match="categories"): + model.set_data(["A", "B", "C"], [("S", [1, 2])]) + + def test_chartex_reads_and_write_guards(self): + data = _mini_package({"ppt/charts/chartEx1.xml": _CHARTEX_XML}) + model = load_chart(OpcPackage.open(data), "ppt/charts/chartEx1.xml") + + assert model.kind == "chartex:funnel" + assert model.title == "Funnel Title" # empty a:t run skipped + series = model.series + assert len(series) == 1 + assert series[0].name == "S1" + assert series[0].categories == ["A", "B", "C"] + assert series[0].values == [1.0, None, 3.0] # idx 1 is a gap + + with pytest.raises(RawUnsupportedError): + model.set_data(["A"], [("S", [1])]) + with pytest.raises(RawUnsupportedError): + model.set_title("New") + + +# -- 6. byte preservation ---------------------------------------------------------- + + +class TestBytePreservation: + def test_set_data_touches_only_chart_xml_and_its_workbook(self): + src = _pptx_with_charts( + [ + ( + XL_CHART_TYPE.COLUMN_CLUSTERED, + ["A", "B"], + [("S1", (1, 2)), ("S2", (3, 4))], + ), + (XL_CHART_TYPE.PIE, ["A", "B", "C"], [("Share", (1, 2, 3))]), + ] + ) + pkg = OpcPackage.open(src) + by_kind = {load_chart(pkg, n).kind: n for n in _chart_part_names(pkg)} + model = load_chart(pkg, by_kind["column"]) + workbook_name = model.embedded_workbook_part().name + + model.set_data(["X", "Y"], [("A", [1, 2])]) + out = pkg.to_bytes() + + before = zipfile.ZipFile(io.BytesIO(src)) + after = zipfile.ZipFile(io.BytesIO(out)) + assert sorted(before.namelist()) == sorted(after.namelist()) + changed = {n for n in before.namelist() if before.read(n) != after.read(n)} + assert changed == {by_kind["column"], workbook_name} + + +# -- 7. missing caches -------------------------------------------------------------- + + +class TestMissingCaches: + def test_no_cache_and_empty_cache_read_safely(self): + data = _mini_package({"xl/charts/chart1.xml": _NO_CACHE_CHART_XML}) + model = load_chart(OpcPackage.open(data), "xl/charts/chart1.xml") + + assert model.kind == "column" + series = model.series + assert len(series) == 2 + # ser 0: refs only, no caches anywhere → everything empty + assert series[0].name is None + assert series[0].categories == [] + assert series[0].values == [] + # ser 1: numCache with ptCount=3 and zero pts → three Nones + assert series[1].values == [None, None, None] + + +# -- extra: scatter xVal/yVal write path ------------------------------------------- + + +class TestScatterXY: + def _scatter_pptx(self) -> bytes: + prs = Presentation() + slide = prs.slides.add_slide(prs.slide_layouts[6]) + chart_data = XyChartData() + ser = chart_data.add_series("S1") + ser.add_data_point(1, 2) + ser.add_data_point(3, 4) + slide.shapes.add_chart( + XL_CHART_TYPE.XY_SCATTER, + Inches(1), + Inches(1), + Inches(6), + Inches(4), + chart_data, + ) + buf = io.BytesIO() + prs.save(buf) + return buf.getvalue() + + def test_scatter_set_data_writes_xval_yval(self): + pkg = OpcPackage.open(self._scatter_pptx()) + name = _chart_part_names(pkg)[0] + model = load_chart(pkg, name) + assert model.kind == "scatter" + assert model.series[0].values == [2.0, 4.0] # yVal + assert model.series[0].categories == ["1", "3"] # xVal, as str + + model.set_data([10, 20, 30], [("S1", [1, 2, 3])]) + out = pkg.to_bytes() + + reopened = load_chart(OpcPackage.open(out), name) + assert reopened.series[0].categories == ["10", "20", "30"] + assert reopened.series[0].values == [1.0, 2.0, 3.0] + # all-numeric x categories → numRef inside c:xVal + ser_el = reopened.xml.find("c:chart/c:plotArea/c:scatterChart/c:ser") + assert ser_el is not None + from contextifier.raw.xmlpart import NS + + assert ser_el.find("c:xVal/c:numRef/c:numCache", NS) is not None + assert ser_el.find("c:yVal/c:numRef/c:numCache", NS) is not None diff --git a/tests/unit/raw/test_docx_raw.py b/tests/unit/raw/test_docx_raw.py new file mode 100644 index 0000000..ecfe4cc --- /dev/null +++ b/tests/unit/raw/test_docx_raw.py @@ -0,0 +1,535 @@ +# tests/unit/raw/test_docx_raw.py +"""DocxRawDocument — addressing parity with python-docx, run-preserving +edits (the edit2docs P0-3 fix), and the byte-preservation contract. + +Every fixture document is built with python-docx so the parity claims +are checked against the library whose addressing edit2docs already +uses. +""" + +from __future__ import annotations + +import io +import struct +import zipfile +import zlib + +import docx +import pytest +from docx.opc.constants import RELATIONSHIP_TYPE as RT +from docx.oxml.ns import qn as docx_qn +from docx.oxml.parser import OxmlElement +from lxml import etree + +from contextifier.raw import open_raw +from contextifier.raw.docx import DocxRawDocument +from contextifier.raw.xmlpart import NS, qn + +W_DRAWING = qn("w:drawing") +W_HYPERLINK = qn("w:hyperlink") + + +# -- helpers ------------------------------------------------------------------ + + +def _png_1x1() -> bytes: + """A deterministic, valid 1x1 RGB PNG (no Pillow needed).""" + + def chunk(tag: bytes, data: bytes) -> bytes: + return ( + struct.pack(">I", len(data)) + + tag + + data + + struct.pack(">I", zlib.crc32(tag + data) & 0xFFFFFFFF) + ) + + ihdr = struct.pack(">IIBBBBB", 1, 1, 8, 2, 0, 0, 0) + idat = zlib.compress(b"\x00\xff\x00\x00") + return ( + b"\x89PNG\r\n\x1a\n" + + chunk(b"IHDR", ihdr) + + chunk(b"IDAT", idat) + + chunk(b"IEND", b"") + ) + + +def _doc_bytes(document) -> bytes: + buf = io.BytesIO() + document.save(buf) + return buf.getvalue() + + +def _raw(document) -> DocxRawDocument: + return open_raw(_doc_bytes(document), extension="docx") + + +def _reopen(raw: DocxRawDocument): + """Round-trip through save and reopen with python-docx.""" + return docx.Document(io.BytesIO(raw.to_bytes())) + + +def _parts(data: bytes) -> dict[str, bytes]: + with zipfile.ZipFile(io.BytesIO(data)) as z: + return {n: z.read(n) for n in z.namelist() if not n.endswith("/")} + + +def _add_hyperlink(paragraph, text: str, url: str) -> str: + """python-docx has no public hyperlink writer; build the oxml.""" + r_id = paragraph.part.relate_to(url, RT.HYPERLINK, is_external=True) + h = OxmlElement("w:hyperlink") + h.set(docx_qn("r:id"), r_id) + r = OxmlElement("w:r") + t = OxmlElement("w:t") + t.text = text + r.append(t) + h.append(r) + paragraph._p.append(h) + return r_id + + +def _set_cell_shading(cell, fill: str) -> None: + shd = OxmlElement("w:shd") + shd.set(docx_qn("w:val"), "clear") + shd.set(docx_qn("w:fill"), fill) + cell._tc.get_or_add_tcPr().append(shd) + + +def _cell_shading_fill(cell) -> str | None: + shd = cell._tc.find(f"{docx_qn('w:tcPr')}/{docx_qn('w:shd')}") + return shd.get(docx_qn("w:fill")) if shd is not None else None + + +# -- 1. addressing parity ------------------------------------------------------- + + +class TestAddressingParity: + @pytest.fixture() + def built(self): + d = docx.Document() + d.add_paragraph("Intro") + d.add_paragraph("Section title", style="Heading 1") + t1 = d.add_table(rows=2, cols=3) + for r in range(2): + for c in range(3): + t1.cell(r, c).text = f"t1 r{r}c{c}" + d.add_paragraph("Middle") + t2 = d.add_table(rows=3, cols=3) + t2.cell(0, 0).merge(t2.cell(0, 1)) # horizontal: gridSpan + t2.cell(1, 2).merge(t2.cell(2, 2)) # vertical: vMerge + seen: set[int] = set() + for r in range(3): + for c in range(3): + cell = t2.cell(r, c) + if id(cell._tc) not in seen: + seen.add(id(cell._tc)) + cell.text = f"t2 r{r}c{c}" + d.add_paragraph("End") + data = _doc_bytes(d) + return open_raw(data, extension="docx"), docx.Document(io.BytesIO(data)) + + def test_paragraph_indices_and_texts_match_python_docx(self, built): + raw, ref = built + assert [p.text for p in raw.paragraphs] == [p.text for p in ref.paragraphs] + assert [p.index for p in raw.paragraphs] == list(range(len(ref.paragraphs))) + + def test_paragraph_style_ids(self, built): + raw, ref = built + assert raw.paragraphs[0].style == "Normal" + assert raw.paragraphs[1].style == ref.paragraphs[1].style.style_id + assert raw.paragraphs[1].style == "Heading1" + + def test_table_dims(self, built): + raw, ref = built + assert len(raw.tables) == len(ref.tables) == 2 + for rt, ft in zip(raw.tables, ref.tables): + assert rt.n_rows == len(ft.rows) + assert rt.n_cols == len(ft.columns) + + def test_every_grid_cell_matches_including_merges(self, built): + raw, ref = built + for t, ft in enumerate(ref.tables): + rt = raw.tables[t] + for r in range(len(ft.rows)): + for c in range(len(ft.columns)): + assert rt.cell(r, c).text == ft.rows[r].cells[c].text, ( + f"table {t} cell ({r},{c})" + ) + + def test_merged_positions_share_one_cell(self, built): + raw, _ = built + t2 = raw.tables[1] + assert t2.cell(0, 0).element is t2.cell(0, 1).element # gridSpan + assert t2.cell(1, 2).element is t2.cell(2, 2).element # vMerge + + def test_cell_out_of_range(self, built): + raw, _ = built + with pytest.raises(IndexError): + raw.tables[0].cell(2, 0) + with pytest.raises(IndexError): + raw.tables[0].cell(0, 3) + + +# -- 2. run preservation (P0-3 fix) -------------------------------------------- + + +class TestRunPreservingParagraphReplace: + @pytest.fixture() + def built(self): + d = docx.Document() + d.add_paragraph("plain first paragraph") + p = d.add_paragraph() + p.add_run("bold lead").bold = True + p.add_run(" middle ") + p.add_run().add_picture(io.BytesIO(_png_1x1())) + p.add_run("italic tail").italic = True + return d + + def test_replace_keeps_formatting_and_image(self, built): + raw = _raw(built) + raw.set_paragraph_text(1, "REPLACED") + ref = _reopen(raw) + + p = ref.paragraphs[1] + assert p.text == "REPLACED" + # first run carried the text and kept its rPr + assert p.runs[0].text == "REPLACED" + assert p.runs[0].bold is True + # image run survived, other pure-text runs are gone + assert len(p._p.findall(f".//{docx_qn('w:drawing')}")) == 1 + assert len(p.runs) == 2 + + def test_media_part_untouched(self, built): + original = _parts(_doc_bytes(built)) + raw = _raw(built) + raw.set_paragraph_text(1, "REPLACED") + edited = _parts(raw.to_bytes()) + media = [n for n in original if n.startswith("word/media/")] + assert media, "fixture should embed an image" + for name in media: + assert edited[name] == original[name] + + def test_whitespace_edges_get_space_preserve(self, built): + raw = _raw(built) + raw.set_paragraph_text(0, " padded ") + assert _reopen(raw).paragraphs[0].text == " padded " + + def test_replace_on_run_free_paragraph_creates_run(self): + d = docx.Document() + d.add_paragraph() # no runs at all + raw = _raw(d) + raw.set_paragraph_text(0, "fresh") + assert _reopen(raw).paragraphs[0].text == "fresh" + + +# -- 3. cell image preservation --------------------------------------------------- + + +class TestCellSetTextPreservesImages: + @pytest.fixture() + def built(self): + d = docx.Document() + t = d.add_table(rows=1, cols=2) + cell = t.cell(0, 0) + cell.text = "before" + cell.add_paragraph().add_run().add_picture(io.BytesIO(_png_1x1())) + cell.add_paragraph("after") + t.cell(0, 1).text = "sibling" + return d + + def test_set_text_keeps_image_paragraph(self, built): + raw = _raw(built) + assert raw.tables[0].cell(0, 0).paragraph_count == 3 + raw.tables[0].cell(0, 0).set_text("X") + ref = _reopen(raw) + + cell = ref.tables[0].cell(0, 0) + texts = [p.text for p in cell.paragraphs] + assert texts == ["X", "", ""] + assert cell.text.strip() == "X" + # the image paragraph kept its drawing, in position + assert len(cell.paragraphs[1]._p.findall(f".//{docx_qn('w:drawing')}")) == 1 + # neighbours untouched + assert ref.tables[0].cell(0, 1).text == "sibling" + + def test_media_part_untouched(self, built): + original = _parts(_doc_bytes(built)) + raw = _raw(built) + raw.tables[0].cell(0, 0).set_text("X") + edited = _parts(raw.to_bytes()) + for name in (n for n in original if n.startswith("word/media/")): + assert edited[name] == original[name] + + +# -- 4. hyperlink policy ------------------------------------------------------------ + + +class TestHyperlinkPolicy: + def _build(self): + d = docx.Document() + p = d.add_paragraph("Go to ") + rid = _add_hyperlink(p, "click here", "https://example.com/") + return d, rid + + def test_replace_keeps_hyperlink_element_and_rel(self): + d, rid = self._build() + raw = _raw(d) + assert raw.paragraphs[0].text == "Go to click here" + raw.set_paragraph_text(0, "REPLACED") + + p_el = raw.paragraphs[0].element + links = p_el.findall(W_HYPERLINK) + assert len(links) == 1 and links[0].get(qn("r:id")) == rid + assert raw.paragraphs[0].text == "REPLACED" + + ref = _reopen(raw) + assert ref.paragraphs[0].text == "REPLACED" + assert rid in ref.part.rels # relationship intact + + def test_hyperlink_carries_text_when_no_direct_run(self): + d = docx.Document() + p = d.add_paragraph() + _add_hyperlink(p, "only link", "https://example.com/") + raw = _raw(d) + raw.set_paragraph_text(0, "NEW") + p_el = raw.paragraphs[0].element + assert len(p_el.findall(W_HYPERLINK)) == 1 + assert _reopen(raw).paragraphs[0].text == "NEW" + + def test_strip_empty_hyperlinks(self): + d, rid = self._build() + raw = _raw(d) + raw.set_paragraph_text(0, "REPLACED") + assert raw.strip_empty_hyperlinks(0) == 1 + assert raw.paragraphs[0].element.findall(W_HYPERLINK) == [] + + ref = _reopen(raw) + assert ref.paragraphs[0].text == "REPLACED" + assert rid not in ref.part.rels # orphaned rel cleaned up + + def test_strip_leaves_nonempty_hyperlinks(self): + d, _ = self._build() + raw = _raw(d) + assert raw.strip_empty_hyperlinks(0) == 0 + assert len(raw.paragraphs[0].element.findall(W_HYPERLINK)) == 1 + + +# -- 5. row editing ------------------------------------------------------------------ + + +class TestRowEditing: + @pytest.fixture() + def built(self): + d = docx.Document() + t = d.add_table(rows=2, cols=3) + for c in range(3): + t.cell(0, c).text = f"head{c}" + _set_cell_shading(t.cell(0, c), "FF0000") + t.cell(1, c).text = f"data{c}" + return d + + def test_insert_row_copies_template_formatting(self, built): + raw = _raw(built) + raw.tables[0].insert_row(1) # template = row 0 (shaded) + ref = _reopen(raw) + + t = ref.tables[0] + assert (len(t.rows), len(t.columns)) == (3, 3) + for c in range(3): + assert t.cell(1, c).text == "" # text cleared + assert _cell_shading_fill(t.cell(1, c)) == "FF0000" # props kept + assert t.cell(0, c).text == f"head{c}" + assert t.cell(2, c).text == f"data{c}" + + def test_insert_row_at_start_and_end(self, built): + raw = _raw(built) + raw.tables[0].insert_row(0) + raw.tables[0].insert_row(raw.tables[0].n_rows) + ref = _reopen(raw) + assert len(ref.tables[0].rows) == 4 + assert ref.tables[0].cell(1, 0).text == "head0" + assert ref.tables[0].cell(3, 0).text == "" + + def test_delete_row(self, built): + raw = _raw(built) + raw.tables[0].delete_row(0) + ref = _reopen(raw) + assert len(ref.tables[0].rows) == 1 + assert ref.tables[0].cell(0, 0).text == "data0" + + def test_row_index_bounds(self, built): + raw = _raw(built) + with pytest.raises(IndexError): + raw.tables[0].insert_row(5) + with pytest.raises(IndexError): + raw.tables[0].delete_row(2) + + +# -- 6. paragraph insert/delete ----------------------------------------------------- + + +class TestParagraphInsertDelete: + @pytest.fixture() + def built(self): + d = docx.Document() + for text in ("A", "B", "C"): + d.add_paragraph(text) + return d + + def test_insert_after_middle(self, built): + raw = _raw(built) + new = raw.insert_paragraph_after(1, "B2") + assert new.index == 2 + assert [p.text for p in _reopen(raw).paragraphs] == ["A", "B", "B2", "C"] + + def test_insert_at_body_start(self, built): + raw = _raw(built) + new = raw.insert_paragraph_after(-1, "Start") + assert new.index == 0 + assert [p.text for p in _reopen(raw).paragraphs] == ["Start", "A", "B", "C"] + + def test_insert_after_last_keeps_sectpr_last(self, built): + raw = _raw(built) + raw.insert_paragraph_after(2, "End", style="Heading1") + data = raw.to_bytes() + + assert [p.text for p in docx.Document(io.BytesIO(data)).paragraphs] == [ + "A", + "B", + "C", + "End", + ] + raw2 = open_raw(data, extension="docx") + assert raw2.paragraphs[3].style == "Heading1" + body = etree.fromstring(_parts(data)["word/document.xml"]).find("w:body", NS) + assert body[-1].tag == qn("w:sectPr") # sectPr still the last body child + + def test_delete_first_and_last(self, built): + raw = _raw(built) + raw.delete_paragraph(2) # last paragraph, adjacent to sectPr + raw.delete_paragraph(0) + assert [p.text for p in _reopen(raw).paragraphs] == ["B"] + + def test_index_bounds(self, built): + raw = _raw(built) + with pytest.raises(IndexError): + raw.set_paragraph_text(3, "x") + with pytest.raises(IndexError): + raw.delete_paragraph(-1) + with pytest.raises(IndexError): + raw.insert_paragraph_after(3, "x") + + +# -- 7. nested tables ----------------------------------------------------------------- + + +class TestNestedTables: + @pytest.fixture() + def built(self): + d = docx.Document() + t = d.add_table(rows=2, cols=2) + host = t.cell(0, 0) + host.text = "host" + nested = host.add_table(rows=2, cols=2) + for r in range(2): + for c in range(2): + nested.cell(r, c).text = f"n{r}{c}" + return d + + def test_nested_table_is_addressable(self, built): + raw = _raw(built) + assert len(raw.tables) == 1 # body-level only, matching python-docx + nested = raw.tables[0].nested_tables(0, 0) + assert len(nested) == 1 + assert (nested[0].n_rows, nested[0].n_cols) == (2, 2) + assert nested[0].cell(1, 1).text == "n11" + assert raw.tables[0].nested_tables(0, 1) == [] + + def test_nested_cell_edit(self, built): + raw = _raw(built) + raw.tables[0].nested_tables(0, 0)[0].cell(0, 0).set_text("edited") + ref = _reopen(raw) + assert ref.tables[0].cell(0, 0).tables[0].cell(0, 0).text == "edited" + + def test_host_set_text_leaves_nested_table_intact(self, built): + raw = _raw(built) + raw.tables[0].cell(0, 0).set_text("HOST") + ref = _reopen(raw) + host = ref.tables[0].cell(0, 0) + assert host.paragraphs[0].text == "HOST" + assert len(host.tables) == 1 + assert host.tables[0].cell(0, 0).text == "n00" + + +# -- 8. byte preservation --------------------------------------------------------------- + + +class TestBytePreservation: + def test_single_paragraph_edit_touches_only_document_xml(self): + d = docx.Document() + d.add_paragraph("hello") + d.add_paragraph().add_run().add_picture(io.BytesIO(_png_1x1())) + d.add_table(rows=1, cols=1).cell(0, 0).text = "cell" + original = _parts(_doc_bytes(d)) + + raw = open_raw(_doc_bytes(d), extension="docx") + raw.set_paragraph_text(0, "changed") + edited = _parts(raw.to_bytes()) + + assert sorted(original) == sorted(edited) + assert edited["word/document.xml"] != original["word/document.xml"] + for name in original: + if name != "word/document.xml": + assert edited[name] == original[name], f"collateral change in {name}" + + +# -- 9. body order ------------------------------------------------------------------------- + + +class TestBodyOrder: + def test_document_order_of_paragraphs_and_tables(self): + d = docx.Document() + d.add_paragraph("p0") + d.add_paragraph("p1") + d.add_table(rows=1, cols=1) + d.add_paragraph("p2") + d.add_table(rows=1, cols=1) + d.add_paragraph("p3") + raw = _raw(d) + assert raw.body_order() == [ + ("p", 0), + ("p", 1), + ("tbl", 0), + ("p", 2), + ("tbl", 1), + ("p", 3), + ] + + +# -- misc surface --------------------------------------------------------------------------- + + +class TestMiscSurface: + def test_format_and_charts_empty(self): + d = docx.Document() + d.add_paragraph("x") + raw = _raw(d) + assert raw.format == "docx" + assert raw.chart_part_names == [] + assert raw.charts == [] + + def test_headers_and_footers_text(self): + d = docx.Document() + d.add_paragraph("body") + section = d.sections[0] + section.header.paragraphs[0].text = "the header" + section.footer.paragraphs[0].text = "the footer" + raw = _raw(d) + assert list(raw.headers.values()) == ["the header"] + assert list(raw.footers.values()) == ["the footer"] + assert all(n.startswith("word/header") for n in raw.headers) + assert all(n.startswith("word/footer") for n in raw.footers) + + def test_open_raw_dispatches_to_docx_model(self): + d = docx.Document() + raw = open_raw(_doc_bytes(d)) # no extension hint — sniffed + assert isinstance(raw, DocxRawDocument) diff --git a/tests/unit/raw/test_opc_contract.py b/tests/unit/raw/test_opc_contract.py new file mode 100644 index 0000000..483702d --- /dev/null +++ b/tests/unit/raw/test_opc_contract.py @@ -0,0 +1,178 @@ +# tests/unit/raw/test_opc_contract.py +"""The byte-preservation contract — the foundation of the raw layer. + +If these fail, nothing else in contextifier.raw can be trusted. +""" + +from __future__ import annotations + +import io +import zipfile + +import pytest + +from contextifier.errors import ContextifierError +from contextifier.raw.opc import OpcPackage +from contextifier.raw.xmlpart import NS, XmlPart, qn + + +def _mini_package() -> bytes: + """A minimal OPC-shaped zip with several parts of varied content.""" + buf = io.BytesIO() + with zipfile.ZipFile(buf, "w", zipfile.ZIP_DEFLATED) as z: + z.writestr( + "[Content_Types].xml", + '' + '' + '' + '' + '' + "", + ) + z.writestr( + "_rels/.rels", + '' + '' + '' + "", + ) + z.writestr("xl/workbook.xml", "") + # Deliberately quirky bytes: BOM, weird whitespace, high compression value + z.writestr("xl/quirky.xml", b"\xef\xbb\xbftext") + z.writestr("media/blob.bin", bytes(range(256)) * 100) + z.writestr( + "xl/_rels/workbook.xml.rels", + '' + '' + '' + "", + ) + z.writestr("xl/charts/chart1.xml", "") + return buf.getvalue() + + +class TestBytePreservation: + def test_untouched_roundtrip_is_content_identical(self): + src = _mini_package() + out = OpcPackage.open(src).to_bytes() + a, b = zipfile.ZipFile(io.BytesIO(src)), zipfile.ZipFile(io.BytesIO(out)) + assert sorted(a.namelist()) == sorted(b.namelist()) + for name in a.namelist(): + assert a.read(name) == b.read(name), f"content drift in {name}" + + def test_single_edit_touches_only_that_part(self): + src = _mini_package() + pkg = OpcPackage.open(src) + pkg.get_part("xl/workbook.xml").write( + b"" + ) + out = pkg.to_bytes() + a, b = zipfile.ZipFile(io.BytesIO(src)), zipfile.ZipFile(io.BytesIO(out)) + for name in a.namelist(): + if name == "xl/workbook.xml": + assert b.read(name) == b"" + else: + assert a.read(name) == b.read(name), f"collateral change in {name}" + + def test_quirky_bytes_survive_verbatim(self): + """BOM + odd whitespace must not be normalized for clean parts.""" + src = _mini_package() + pkg = OpcPackage.open(src) + pkg.get_part("xl/workbook.xml").write(b"") # dirty a sibling + out = zipfile.ZipFile(io.BytesIO(pkg.to_bytes())) + assert out.read("xl/quirky.xml") == b"\xef\xbb\xbftext" + + def test_lazy_xmlpart_read_does_not_dirty(self): + """Parsing (without mutation) must never rewrite a part.""" + src = _mini_package() + pkg = OpcPackage.open(src) + xp = XmlPart(pkg.get_part("xl/quirky.xml")) + assert xp.root.tag == "a" # parses fine + xp.flush() # no-op: not dirty + out = zipfile.ZipFile(io.BytesIO(pkg.to_bytes())) + assert out.read("xl/quirky.xml") == b"\xef\xbb\xbftext" + + def test_add_and_remove_parts(self): + pkg = OpcPackage.open(_mini_package()) + pkg.add_part("xl/new.xml", b"") + pkg.remove_part("media/blob.bin") + out = zipfile.ZipFile(io.BytesIO(pkg.to_bytes())) + assert out.read("xl/new.xml") == b"" + assert "media/blob.bin" not in out.namelist() + + +class TestPackageBasics: + def test_open_from_path_bytes_and_stream(self, tmp_path): + data = _mini_package() + p = tmp_path / "x.zip" + p.write_bytes(data) + assert OpcPackage.open(p).part_names == OpcPackage.open(data).part_names + assert OpcPackage.open(io.BytesIO(data)).has_part("xl/workbook.xml") + + def test_bad_magic_rejected(self): + with pytest.raises(ContextifierError, match="ZIP"): + OpcPackage.open(b"not a zip at all") + + def test_sniff_format(self): + assert OpcPackage.open(_mini_package()).sniff_format() == "xlsx" + + def test_content_type_lookup_and_override(self): + pkg = OpcPackage.open(_mini_package()) + assert pkg.content_type_of("xl/workbook.xml") == "application/wb+xml" + assert pkg.content_type_of("media/blob.bin") == "application/octet-stream" + pkg.set_content_type_override("xl/charts/chart1.xml", "application/chart+xml") + assert pkg.content_type_of("xl/charts/chart1.xml") == "application/chart+xml" + + +class TestRelationships: + def test_rels_iteration_and_resolution(self): + pkg = OpcPackage.open(_mini_package()) + rels = pkg.rels_for("xl/workbook.xml") + assert rels is not None + charts = rels.by_type("/chart") + assert charts[0]["id"] == "rId7" + assert ( + rels.resolve("xl/workbook.xml", charts[0]["target"]) + == "xl/charts/chart1.xml" + ) + assert rels.target_of("rId7") == "charts/chart1.xml" + + def test_package_root_rels(self): + pkg = OpcPackage.open(_mini_package()) + root_rels = pkg.rels_for("") + assert ( + root_rels is not None and root_rels.target_of("rId1") == "xl/workbook.xml" + ) + + def test_add_remove_and_next_id(self): + pkg = OpcPackage.open(_mini_package()) + rels = pkg.rels_for("xl/workbook.xml") + rid = rels.next_id() + assert rid != "rId7" + rels.add(rid, "http://x/image", "media/blob.bin") + assert rels.target_of(rid) == "media/blob.bin" + assert rels.remove(rid) is True + assert rels.target_of(rid) is None + + +class TestQn: + def test_qn_expands_known_prefixes(self): + assert qn("w:p") == f"{{{NS['w']}}}p" + assert qn("c:ser") == f"{{{NS['c']}}}ser" + + def test_qn_rejects_unknown_prefix(self): + with pytest.raises(KeyError, match="prefix"): + qn("zz:nope") + + +class TestRepeatedSave: + def test_double_to_bytes_is_stable(self): + """writestr must not mutate source ZipInfo state (regression: + header_offset corruption made the second save raise BadZipFile).""" + pkg = OpcPackage.open(_mini_package()) + pkg.get_part("xl/workbook.xml").write(b"") + first = pkg.to_bytes() + second = pkg.to_bytes() # must not raise + a, b = zipfile.ZipFile(io.BytesIO(first)), zipfile.ZipFile(io.BytesIO(second)) + for name in a.namelist(): + assert a.read(name) == b.read(name) diff --git a/tests/unit/raw/test_open_raw_integration.py b/tests/unit/raw/test_open_raw_integration.py new file mode 100644 index 0000000..930412f --- /dev/null +++ b/tests/unit/raw/test_open_raw_integration.py @@ -0,0 +1,146 @@ +# tests/unit/raw/test_open_raw_integration.py +"""Cross-format integration: the README's raw-layer promises, end to end. + +These run the exact user-facing flows — open_raw dispatch, chart editing +through the format models' ``charts`` property, combined cell+chart +edits, and DocumentProcessor.open_raw — with Office-library reopens as +the acceptance check. +""" + +from __future__ import annotations + +import io + +import pytest + +from contextifier import DocumentProcessor, open_raw +from contextifier.raw import RawUnsupportedError +from contextifier.raw.docx import DocxRawDocument +from contextifier.raw.pptx import PptxRawDocument +from contextifier.raw.xlsx import XlsxRawDocument + + +def _xlsx_with_chart() -> bytes: + from openpyxl import Workbook + from openpyxl.chart import BarChart, Reference + + wb = Workbook() + ws = wb.active + ws.title = "Sales" + ws.append(["Quarter", "Amount"]) + for row in [["Q1", 120], ["Q2", 135], ["Q3", 150]]: + ws.append(row) + chart = BarChart() + chart.title = "Sales by Quarter" + chart.add_data( + Reference(ws, min_col=2, min_row=1, max_row=4), titles_from_data=True + ) + chart.set_categories(Reference(ws, min_col=1, min_row=2, max_row=4)) + ws.add_chart(chart, "D2") + buf = io.BytesIO() + wb.save(buf) + return buf.getvalue() + + +def _pptx_with_chart() -> bytes: + from pptx import Presentation + from pptx.chart.data import CategoryChartData + from pptx.enum.chart import XL_CHART_TYPE + from pptx.util import Inches + + prs = Presentation() + slide = prs.slides.add_slide(prs.slide_layouts[5]) + data = CategoryChartData() + data.categories = ["East", "West"] + data.add_series("Units", (10, 20)) + slide.shapes.add_chart( + XL_CHART_TYPE.COLUMN_CLUSTERED, Inches(1), Inches(1), Inches(6), Inches(4), data + ) + buf = io.BytesIO() + prs.save(buf) + return buf.getvalue() + + +class TestDispatch: + def test_open_raw_from_path_and_bytes(self, tmp_path): + data = _xlsx_with_chart() + path = tmp_path / "s.xlsx" + path.write_bytes(data) + assert isinstance(open_raw(path), XlsxRawDocument) + assert isinstance(open_raw(data), XlsxRawDocument) # sniffed + assert isinstance(open_raw(_pptx_with_chart()), PptxRawDocument) + + def test_docx_dispatch(self): + from docx import Document + + buf = io.BytesIO() + Document().save(buf) + assert isinstance(open_raw(buf.getvalue()), DocxRawDocument) + + def test_unsupported_raises(self): + import zipfile + + buf = io.BytesIO() + with zipfile.ZipFile(buf, "w") as z: + z.writestr("hello.txt", "not an office file") + with pytest.raises(RawUnsupportedError, match="supported"): + open_raw(buf.getvalue(), extension="zip") + + def test_document_processor_open_raw(self, tmp_path): + path = tmp_path / "s.xlsx" + path.write_bytes(_xlsx_with_chart()) + raw = DocumentProcessor().open_raw(str(path)) + assert isinstance(raw, XlsxRawDocument) + + +class TestReadmeFlowXlsx: + def test_cell_plus_chart_edit_roundtrip(self, tmp_path): + """The README quick-start, verbatim behavior.""" + from openpyxl import load_workbook + + raw = open_raw(_xlsx_with_chart()) + raw.sheets["Sales"].set_cell("B3", 142) + charts = raw.charts + assert len(charts) == 1 and charts[0].kind in ("bar", "column") + assert charts[0].title == "Sales by Quarter" + charts[0].set_data( + categories=["Q1", "Q2", "Q3"], + series=[("Sales", [120, 142, 150])], + ) + out = tmp_path / "edited.xlsx" + raw.save(out) + + wb = load_workbook(out) + assert wb["Sales"]["B3"].value == 142 + # chart part still present and parseable by our own model + raw2 = open_raw(out) + assert [round(v) for v in raw2.charts[0].series[0].values] == [120, 142, 150] + + def test_ai_view_still_works_on_edited_file(self, tmp_path): + """The two views coexist: raw-edit, then extract_text reads it.""" + raw = open_raw(_xlsx_with_chart()) + raw.sheets["Sales"].set_cell("A5", "Q4-new") + out = tmp_path / "for_extract.xlsx" + raw.save(out) + text = DocumentProcessor().extract_text(str(out)) + assert "Q4-new" in text + + +class TestReadmeFlowPptx: + def test_chart_edit_via_slide(self, tmp_path): + from pptx import Presentation + + raw = open_raw(_pptx_with_chart()) + slide = raw.slides[0] + assert slide.chart_part_names, "chart not discovered on slide" + chart = slide.charts[0] + chart.set_data(categories=["North", "South"], series=[("Units", [7, 9])]) + out = tmp_path / "deck.pptx" + raw.save(out) + + prs = Presentation(str(out)) + found = [s for s in prs.slides[0].shapes if getattr(s, "has_chart", False)] + assert found + plot = found[0].chart.plots[0] + assert list(plot.categories) == ["North", "South"] + assert [round(v) for v in plot.series[0].values] == [7, 9] diff --git a/tests/unit/raw/test_pptx_raw.py b/tests/unit/raw/test_pptx_raw.py new file mode 100644 index 0000000..a2a7645 --- /dev/null +++ b/tests/unit/raw/test_pptx_raw.py @@ -0,0 +1,369 @@ +# tests/unit/raw/test_pptx_raw.py +"""PptxRawDocument — C4 of the raw layer. + +Decks are built with python-pptx (the reference implementation) and +reopened with it after every raw edit, so each assertion is a +round-trip through an independent OOXML reader. +""" + +from __future__ import annotations + +import base64 +import io +import zipfile + +import pytest + +pptx = pytest.importorskip("pptx") + +from pptx import Presentation # noqa: E402 +from pptx.chart.data import CategoryChartData # noqa: E402 +from pptx.enum.chart import XL_CHART_TYPE # noqa: E402 +from pptx.enum.shapes import MSO_SHAPE_TYPE # noqa: E402 +from pptx.util import Inches # noqa: E402 + +from contextifier.raw import open_raw # noqa: E402 + +# 1x1 px valid PNG — fallback when Pillow is unavailable. +_TINY_PNG = base64.b64decode( + "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR4nGP4" + "z8DwHwAFBQIAX8jx0gAAAABJRU5ErkJggg==" +) + + +def _png_bytes() -> bytes: + try: + from PIL import Image + except ImportError: + return _TINY_PNG + buf = io.BytesIO() + Image.new("RGB", (8, 8), (200, 30, 30)).save(buf, format="PNG") + return buf.getvalue() + + +def build_deck() -> bytes: + """Slide 1: title (bold run) + body + notes. Slide 2: table, chart, + picture.""" + prs = Presentation() + + s1 = prs.slides.add_slide(prs.slide_layouts[1]) + s1.shapes.title.text = "Hello Title" + s1.shapes.title.text_frame.paragraphs[0].runs[0].font.bold = True + s1.placeholders[1].text = "Body text" + s1.notes_slide.notes_text_frame.text = "Speaker notes here" + + s2 = prs.slides.add_slide(prs.slide_layouts[6]) + tbl = s2.shapes.add_table(2, 3, Inches(0.5), Inches(0.5), Inches(6), Inches(1.5)) + for r in range(2): + for c in range(3): + tbl.table.cell(r, c).text = f"R{r}C{c}" + chart_data = CategoryChartData() + chart_data.categories = ["Q1", "Q2", "Q3"] + chart_data.add_series("Sales", (10.0, 20.0, 30.0)) + s2.shapes.add_chart( + XL_CHART_TYPE.COLUMN_CLUSTERED, + Inches(0.5), + Inches(2.5), + Inches(4), + Inches(3), + chart_data, + ) + s2.shapes.add_picture( + io.BytesIO(_png_bytes()), Inches(5), Inches(2.5), Inches(1), Inches(1) + ) + + buf = io.BytesIO() + prs.save(buf) + return buf.getvalue() + + +def build_three_slide_deck() -> bytes: + """Three text slides; slide 2 additionally owns the deck's ONLY + chart (with embedded workbook) and a notes slide.""" + prs = Presentation() + for i, txt in enumerate(["One", "Two", "Three"]): + s = prs.slides.add_slide(prs.slide_layouts[6]) + box = s.shapes.add_textbox(Inches(1), Inches(1), Inches(3), Inches(1)) + box.text_frame.text = txt + if i == 1: + chart_data = CategoryChartData() + chart_data.categories = ["A", "B"] + chart_data.add_series("S", (1.0, 2.0)) + s.shapes.add_chart( + XL_CHART_TYPE.COLUMN_CLUSTERED, + Inches(1), + Inches(2), + Inches(3), + Inches(3), + chart_data, + ) + s.notes_slide.notes_text_frame.text = "gone with the slide" + buf = io.BytesIO() + prs.save(buf) + return buf.getvalue() + + +# A complete handwritten slide: one text box saying REPLACED. +MINIMAL_SLIDE_XML = ( + b'' + b'' + b"" + b'' + b"" + b"" + b'' + b"" + b'' + b'' + b'' + b"" + b"REPLACED" + b"" + b"" + b"" + b"" +) + + +class TestInventory: + def test_shapes_kinds_and_texts(self): + data = build_deck() + raw = open_raw(data) + assert raw.format == "pptx" + assert len(raw.slides) == 2 + assert raw.slides[0].part_name == "ppt/slides/slide1.xml" + + s1 = raw.slides[0].shapes + assert [s.kind for s in s1] == ["text", "text"] + assert [s.text for s in s1] == ["Hello Title", "Body text"] + + s2 = raw.slides[1].shapes + assert [s.kind for s in s2] == ["table", "chart", "picture"] + assert all(s.text is None for s in s2) + assert all(s.name for s in s2) # python-pptx names every shape + + def test_slide_and_shape_order_matches_python_pptx(self): + data = build_deck() + raw = open_raw(data) + prs = Presentation(io.BytesIO(data)) + assert len(raw.slides) == len(prs.slides) + for raw_slide, pp_slide in zip(raw.slides, prs.slides): + assert [s.id for s in raw_slide.shapes] == [ + sh.shape_id for sh in pp_slide.shapes + ] + + +class TestText: + def test_set_text_roundtrip_preserves_first_run_formatting(self): + raw = open_raw(build_deck()) + slide = raw.slides[0] + title_id = next(s.id for s in slide.shapes if s.text == "Hello Title") + slide.set_text(title_id, "New Title") + assert slide.get_text(title_id) == "New Title" + + prs = Presentation(io.BytesIO(raw.to_bytes())) + title = prs.slides[0].shapes.title + assert title.text == "New Title" + run = title.text_frame.paragraphs[0].runs[0] + assert run.font.bold is True # a:rPr of the original run survived + + def test_get_and_set_text_reject_shapes_without_txbody(self): + raw = open_raw(build_deck()) + slide = raw.slides[1] + pic_id = next(s.id for s in slide.shapes if s.kind == "picture") + with pytest.raises(ValueError, match="text body"): + slide.get_text(pic_id) + with pytest.raises(ValueError, match="text body"): + slide.set_text(pic_id, "nope") + + def test_set_text_paragraph_out_of_range(self): + raw = open_raw(build_deck()) + slide = raw.slides[0] + title_id = next(s.id for s in slide.shapes if s.text == "Hello Title") + with pytest.raises(IndexError): + slide.set_text(title_id, "x", para=5) + + +class TestTables: + def test_dims_and_cell_read(self): + raw = open_raw(build_deck()) + tables = raw.slides[1].tables + assert len(tables) == 1 + t = tables[0] + assert (t.n_rows, t.n_cols) == (2, 3) + assert t.cell(0, 0).text == "R0C0" + assert t.cell(1, 2).text == "R1C2" + with pytest.raises(IndexError): + t.cell(2, 0) + + def test_cell_set_text_and_insert_row_roundtrip(self): + raw = open_raw(build_deck()) + t = raw.slides[1].tables[0] + t.cell(0, 0).set_text("HDR") + t.insert_row(1) # empty row cloned from row 0 + assert (t.n_rows, t.n_cols) == (3, 3) + + prs = Presentation(io.BytesIO(raw.to_bytes())) + table = next(sh for sh in prs.slides[1].shapes if sh.has_table).table + assert len(table.rows) == 3 + assert len(table.columns) == 3 + assert table.cell(0, 0).text == "HDR" + assert [table.cell(1, c).text for c in range(3)] == ["", "", ""] + assert [table.cell(2, c).text for c in range(3)] == ["R1C0", "R1C1", "R1C2"] + + def test_delete_row_roundtrip(self): + raw = open_raw(build_deck()) + raw.slides[1].tables[0].delete_row(0) + prs = Presentation(io.BytesIO(raw.to_bytes())) + table = next(sh for sh in prs.slides[1].shapes if sh.has_table).table + assert len(table.rows) == 1 + assert table.cell(0, 0).text == "R1C0" + + +class TestChartsAndNotes: + def test_chart_part_names(self): + raw = open_raw(build_deck()) + assert raw.slides[1].chart_part_names == ["ppt/charts/chart1.xml"] + assert raw.slides[0].chart_part_names == [] + + def test_notes_text(self): + raw = open_raw(build_deck()) + assert raw.slides[0].notes_text == "Speaker notes here" + assert raw.slides[1].notes_text is None + + +class TestReplaceContent: + def test_preserve_native_keeps_chart_table_picture(self): + raw = open_raw(build_deck()) + preserved = raw.slides[1].replace_content( + MINIMAL_SLIDE_XML, preserve_native=True + ) + assert "table" in preserved + assert "chart:chart1.xml" in preserved + assert any(p.startswith("picture:") for p in preserved) + + prs = Presentation(io.BytesIO(raw.to_bytes())) + slide = prs.slides[1] + texts = [sh.text_frame.text for sh in slide.shapes if sh.has_text_frame] + assert "REPLACED" in texts + + chart_shapes = [sh for sh in slide.shapes if sh.has_chart] + assert len(chart_shapes) == 1 + chart = chart_shapes[0].chart # the native chart still loads... + assert len(chart.plots) == 1 # ...and its plot area is readable + assert list(chart.plots[0].categories) == ["Q1", "Q2", "Q3"] + + table_shapes = [sh for sh in slide.shapes if sh.has_table] + assert len(table_shapes) == 1 + assert table_shapes[0].table.cell(0, 0).text == "R0C0" + + pictures = [ + sh for sh in slide.shapes if sh.shape_type == MSO_SHAPE_TYPE.PICTURE + ] + assert len(pictures) == 1 + + ids = [sh.shape_id for sh in slide.shapes] + assert len(ids) == len(set(ids)) # renumbering avoided collisions + + def test_no_preserve_drops_frames_but_keeps_parts(self): + raw = open_raw(build_deck()) + preserved = raw.slides[1].replace_content( + MINIMAL_SLIDE_XML, preserve_native=False + ) + assert preserved == [] + + out = raw.to_bytes() + prs = Presentation(io.BytesIO(out)) + slide = prs.slides[1] + assert not any(sh.has_chart for sh in slide.shapes) + assert not any(sh.has_table for sh in slide.shapes) + assert [sh.text_frame.text for sh in slide.shapes if sh.has_text_frame] == [ + "REPLACED" + ] + # No orphan cleanup here — that is remove_slide's job. + names = zipfile.ZipFile(io.BytesIO(out)).namelist() + assert "ppt/charts/chart1.xml" in names + assert any(n.startswith("ppt/embeddings/") for n in names) + + def test_rejects_non_slide_xml(self): + raw = open_raw(build_deck()) + with pytest.raises(ValueError, match="p:sld"): + raw.slides[0].replace_content(b"") + + +class TestRemoveSlide: + def test_orphan_cleanup_and_byte_identity(self): + data = build_three_slide_deck() + zin = zipfile.ZipFile(io.BytesIO(data)) + assert "ppt/charts/chart1.xml" in zin.namelist() + assert any(n.startswith("ppt/embeddings/") for n in zin.namelist()) + + raw = open_raw(data) + raw.remove_slide(1) + out = raw.to_bytes() + z = zipfile.ZipFile(io.BytesIO(out)) + names = set(z.namelist()) + + # The slide, its notes, the chart and its embedded workbook are gone. + for gone in ( + "ppt/slides/slide2.xml", + "ppt/slides/_rels/slide2.xml.rels", + "ppt/notesSlides/notesSlide1.xml", + "ppt/notesSlides/_rels/notesSlide1.xml.rels", + "ppt/charts/chart1.xml", + "ppt/charts/_rels/chart1.xml.rels", + ): + assert gone not in names, f"orphan left behind: {gone}" + assert not any(n.startswith("ppt/embeddings/") for n in names) + + # Shared infrastructure survives the reference count. + assert "ppt/slideLayouts/slideLayout7.xml" in names + assert "ppt/notesMasters/notesMaster1.xml" in names + + # Content types no longer mention the removed parts. + content_types = z.read("[Content_Types].xml").decode() + for part in ( + "/ppt/slides/slide2.xml", + "/ppt/charts/chart1.xml", + "/ppt/notesSlides/notesSlide1.xml", + ): + assert part not in content_types + + # Untouched slides are byte-identical. + assert z.read("ppt/slides/slide1.xml") == zin.read("ppt/slides/slide1.xml") + assert z.read("ppt/slides/slide3.xml") == zin.read("ppt/slides/slide3.xml") + + # And the deck still opens cleanly with the right slides. + prs = Presentation(io.BytesIO(out)) + assert len(prs.slides) == 2 + remaining = [ + next(sh.text_frame.text for sh in s.shapes if sh.has_text_frame) + for s in prs.slides + ] + assert remaining == ["One", "Three"] + + def test_index_out_of_range(self): + raw = open_raw(build_three_slide_deck()) + with pytest.raises(IndexError): + raw.remove_slide(3) + + +class TestBytePreservation: + def test_set_text_touches_only_that_slide_part(self): + data = build_deck() + raw = open_raw(data) + slide = raw.slides[0] + title_id = next(s.id for s in slide.shapes if s.text == "Hello Title") + slide.set_text(title_id, "Changed") + out = raw.to_bytes() + + a = zipfile.ZipFile(io.BytesIO(data)) + b = zipfile.ZipFile(io.BytesIO(out)) + assert sorted(a.namelist()) == sorted(b.namelist()) + for name in a.namelist(): + if name == "ppt/slides/slide1.xml": + assert a.read(name) != b.read(name) + else: + assert a.read(name) == b.read(name), f"collateral change in {name}" diff --git a/tests/unit/raw/test_xlsx_raw.py b/tests/unit/raw/test_xlsx_raw.py new file mode 100644 index 0000000..b20551e --- /dev/null +++ b/tests/unit/raw/test_xlsx_raw.py @@ -0,0 +1,602 @@ +# tests/unit/raw/test_xlsx_raw.py +"""XLSX raw model — surgical cell edits under the byte-preservation contract. + +The headline: an edit through :class:`XlsxRawDocument` re-serializes +exactly the worksheet it touched (plus ``xl/workbook.xml`` when a +``calcPr`` flag is needed) — everything an openpyxl round-trip destroys +(chart style parts, custom XML, sparkline extLst blocks) survives +byte-identical. + +Fixtures are built with openpyxl and then, where openpyxl cannot express +a feature (shared strings, formula cached values, foreign parts), the +package is patched at the ZIP level to look like an Excel-saved file. +Every patch asserts its anchor string so openpyxl output drift fails +loudly instead of silently weakening a test. +""" + +from __future__ import annotations + +import io +import re +import warnings +import zipfile + +import openpyxl +import pytest +from lxml import etree +from openpyxl.chart import BarChart, Reference +from openpyxl.workbook.defined_name import DefinedName + +from contextifier.raw import open_raw +from contextifier.raw.xlsx import ( + XlsxRawDocument, + col_index_to_letters, + col_letters_to_index, + parse_ref, +) +from contextifier.raw.xmlpart import NS + +SHEET1 = "xl/worksheets/sheet1.xml" +SHEET2 = "xl/worksheets/sheet2.xml" +WORKBOOK = "xl/workbook.xml" +WORKBOOK_RELS = "xl/_rels/workbook.xml.rels" +CONTENT_TYPES = "[Content_Types].xml" +SHARED_STRINGS = "xl/sharedStrings.xml" + + +# -- fixture builders -------------------------------------------------------- + + +def _build_workbook(with_image: bool = False) -> bytes: + """Two sheets, strings/numbers/bools/formula, merge, number format, + a bar chart, a defined name — and optionally an image.""" + wb = openpyxl.Workbook() + ws = wb.active + ws.title = "Sales" + ws["A1"] = "Region" + ws["B1"] = "Amount" + ws["C1"] = True + ws["D1"] = 1234.5 + ws["D1"].number_format = "#,##0.00" + ws["A2"] = "East" + ws["B2"] = 100 + ws["A3"] = "West" + ws["B3"] = 250.5 + ws["B4"] = "=SUM(B2:B3)" + ws["A5"] = "Merged" + ws.merge_cells("A5:B5") + chart = BarChart() + chart.add_data( + Reference(ws, min_col=2, min_row=1, max_row=3), titles_from_data=True + ) + chart.set_categories(Reference(ws, min_col=1, min_row=2, max_row=3)) + ws.add_chart(chart, "E7") + wb.defined_names["MyRange"] = DefinedName("MyRange", attr_text="Sales!$A$1") + ws2 = wb.create_sheet("Notes") + ws2["A1"] = "hello" + if with_image: + from openpyxl.drawing.image import Image as XlImage + from PIL import Image as PilImage + + ibuf = io.BytesIO() + PilImage.new("RGB", (4, 4), (200, 30, 30)).save(ibuf, format="PNG") + ibuf.seek(0) + ws.add_image(XlImage(ibuf), "F10") + buf = io.BytesIO() + wb.save(buf) + return buf.getvalue() + + +def _parts(data: bytes) -> dict[str, bytes]: + z = zipfile.ZipFile(io.BytesIO(data)) + return {n: z.read(n) for n in z.namelist()} + + +def _rezip( + data: bytes, + *, + replace: dict[str, bytes] | None = None, + add: dict[str, bytes] | None = None, +) -> bytes: + z = zipfile.ZipFile(io.BytesIO(data)) + out = io.BytesIO() + with zipfile.ZipFile(out, "w", zipfile.ZIP_DEFLATED) as znew: + for n in z.namelist(): + znew.writestr(n, (replace or {}).get(n, z.read(n))) + for n, content in (add or {}).items(): + znew.writestr(n, content) + return out.getvalue() + + +def _diff_parts(a: bytes, b: bytes) -> set[str]: + pa, pb = _parts(a), _parts(b) + assert set(pa) == set(pb), "part list drift" + return {n for n in pa if pa[n] != pb[n]} + + +def _load_wb(data: bytes) -> openpyxl.Workbook: + with warnings.catch_warnings(): + warnings.simplefilter("ignore") # openpyxl warns on features it drops + return openpyxl.load_workbook(io.BytesIO(data)) + + +def _sheet_root(data: bytes, part: str): + return etree.fromstring(_parts(data)[part]) + + +def _cell_el(data: bytes, part: str, ref: str): + for c in _sheet_root(data, part).findall("s:sheetData/s:row/s:c", NS): + if c.get("r") == ref: + return c + return None + + +def _excel_flavoured(base: bytes) -> bytes: + """Patch openpyxl output to look Excel-saved: a shared-string cell, + a cached formula value, and a calcPr *without* fullCalcOnLoad.""" + parts = _parts(base) + s1 = parts[SHEET1] + anchor = b"SUM(B2:B3)" + assert anchor in s1, "openpyxl formula serialization drifted" + s1 = s1.replace(anchor, b"SUM(B2:B3)350.5") + anchor = b'Region' + assert anchor in s1, "openpyxl inline-string serialization drifted" + s1 = s1.replace(anchor, b'0') + wb = re.sub(rb"]*/>", b'', parts[WORKBOOK]) + assert b"fullCalcOnLoad" not in wb + ct = parts[CONTENT_TYPES].replace( + b"", + b'' + b"", + ) + rels = parts[WORKBOOK_RELS].replace( + b"", + b'', + ) + sst = ( + b'' + b'' + b'Region' + ) + return _rezip( + base, + replace={ + SHEET1: s1, + WORKBOOK: wb, + CONTENT_TYPES: ct, + WORKBOOK_RELS: rels, + }, + add={SHARED_STRINGS: sst}, + ) + + +_SPARKLINE_EXT = ( + b'' + b'' + b'' + b"Sales!B2:B3B2" + b"" + b"" +) + +FOREIGN_PARTS = [ + "xl/charts/style1.xml", + "xl/charts/colors1.xml", + "xl/charts/_rels/chart1.xml.rels", + "customXml/item1.xml", + "customXml/itemProps1.xml", + "customXml/_rels/item1.xml.rels", +] + + +def _inject_foreign_features(base: bytes) -> bytes: + """Add the parts openpyxl round-trips destroy: chart style/colors, + custom XML (with rels + content types), and a sparkline extLst + inside the Notes worksheet.""" + parts = _parts(base) + ct = parts[CONTENT_TYPES].replace( + b"", + b'' + b'' + b'' + b"", + ) + wb_rels = parts[WORKBOOK_RELS].replace( + b"", + b'', + ) + s2 = parts[SHEET2] + assert b"" in s2 + s2 = s2.replace(b"", _SPARKLINE_EXT + b"") + chart_rels = ( + b'' + b'' + b'' + b'' + b"" + ) + item_rels = ( + b'' + b'' + b'' + ) + return _rezip( + base, + replace={CONTENT_TYPES: ct, WORKBOOK_RELS: wb_rels, SHEET2: s2}, + add={ + "xl/charts/style1.xml": ( + b'' + ), + "xl/charts/colors1.xml": ( + b'' + ), + "customXml/item1.xml": ( + '홍길동' + ).encode("utf-8"), + "customXml/itemProps1.xml": ( + b'' + ), + "customXml/_rels/item1.xml.rels": item_rels, + "xl/charts/_rels/chart1.xml.rels": chart_rels, + }, + ) + + +# -- reading ------------------------------------------------------------------ + + +class TestReadModel: + def test_open_raw_dispatch(self): + raw = open_raw(_build_workbook()) # bytes: format is sniffed + assert isinstance(raw, XlsxRawDocument) + assert raw.format == "xlsx" + + def test_sheet_accessors(self): + raw = open_raw(_build_workbook(), extension="xlsx") + assert raw.sheet_names == ["Sales", "Notes"] + assert raw.sheets[0].name == "Sales" + assert raw.sheets["Notes"].part_name == SHEET2 + assert len(raw.sheets) == 2 + assert [s.name for s in raw.sheets] == ["Sales", "Notes"] + assert "Sales" in raw.sheets and "Nope" not in raw.sheets + with pytest.raises(KeyError, match="Nope"): + raw.sheets["Nope"] + with pytest.raises(IndexError): + raw.sheets[9] + + def test_values_and_types(self): + raw = open_raw(_excel_flavoured(_build_workbook()), extension="xlsx") + sales = raw.sheets["Sales"] + assert sales.get_cell("A1") == "Region" # shared string, 2 runs + v = sales.get_cell("B2") + assert v == 100 and isinstance(v, int) + v = sales.get_cell("B3") + assert v == 250.5 and isinstance(v, float) + assert sales.get_cell("C1") is True + assert sales.get_cell("A5") == "Merged" # inline string + assert raw.sheets["Notes"].get_cell("A1") == "hello" + + def test_formula_cached_value_and_text(self): + raw = open_raw(_excel_flavoured(_build_workbook()), extension="xlsx") + sales = raw.sheets["Sales"] + assert sales.get_cell("B4") == 350.5 # the cached result + assert sales.get_formula("B4") == "SUM(B2:B3)" + assert sales.get_formula("B3") is None + assert sales.get_formula("Z99") is None + + def test_dimensions_and_merged_ranges(self): + raw = open_raw(_build_workbook(), extension="xlsx") + assert raw.sheets["Sales"].dimensions == (5, 4) + assert raw.sheets["Sales"].merged_ranges == ["A5:B5"] + assert raw.sheets["Notes"].dimensions == (1, 1) + assert raw.sheets["Notes"].merged_ranges == [] + + def test_defined_names(self): + raw = open_raw(_build_workbook(), extension="xlsx") + assert raw.defined_names == {"MyRange": "Sales!$A$1"} + + def test_chart_part_names(self): + raw = open_raw(_build_workbook(), extension="xlsx") + assert raw.chart_part_names == ["xl/charts/chart1.xml"] + + def test_iter_rows(self): + raw = open_raw(_excel_flavoured(_build_workbook()), extension="xlsx") + cells = dict(raw.sheets["Sales"].iter_rows()) + assert cells["A1"] == "Region" + assert cells["B3"] == 250.5 + window = dict(raw.sheets["Sales"].iter_rows(min_row=2, max_row=3)) + assert set(window) == {"A2", "B2", "A3", "B3"} + + def test_missing_cells_are_none(self): + raw = open_raw(_build_workbook(), extension="xlsx") + sales = raw.sheets["Sales"] + assert sales.get_cell("Z99") is None # row doesn't exist + assert sales.get_cell("D5") is None # row exists, cell doesn't + with pytest.raises(ValueError, match="reference"): + sales.get_cell("not-a-ref") + + def test_shared_string_ref_without_table_is_none(self): + """No sharedStrings part at all: constructor stays resilient and + a dangling t="s" cell reads as None instead of raising.""" + parts = _parts(_build_workbook()) + s1 = parts[SHEET1] + anchor = b'Region' + assert anchor in s1 + data = _rezip( + _build_workbook(), + replace={SHEET1: s1.replace(anchor, b'0')}, + ) + raw = open_raw(data, extension="xlsx") + assert raw.sheets["Sales"].get_cell("A1") is None + + def test_ref_helpers(self): + assert col_letters_to_index("A") == 1 + assert col_letters_to_index("AA") == 27 + assert col_index_to_letters(28) == "AB" + for i in (1, 26, 27, 52, 703): + assert col_letters_to_index(col_index_to_letters(i)) == i + assert parse_ref("$B$3") == (3, 2) + with pytest.raises(ValueError): + parse_ref("B0") + + +# -- surgical writes ----------------------------------------------------------- + + +class TestSurgicalWrites: + def test_single_set_cell_diffs_exactly_one_part(self): + base = _build_workbook() + # openpyxl already emits fullCalcOnLoad, so workbook.xml stays clean + assert b'fullCalcOnLoad="1"' in _parts(base)[WORKBOOK] + raw = open_raw(base, extension="xlsx") + raw.sheets["Sales"].set_cell("B3", 999) + out = raw.to_bytes() + + assert _diff_parts(base, out) == {SHEET1} + assert _parts(out)[CONTENT_TYPES] == _parts(base)[CONTENT_TYPES] + + wb = _load_wb(out) + assert wb["Sales"]["B3"].value == 999 + assert wb["Sales"]["A1"].value == "Region" # neighbors untouched + assert wb["Sales"]["D1"].number_format == "#,##0.00" # styles intact + + def test_style_attribute_preserved_on_overwrite(self): + base = _build_workbook() + orig_style = _cell_el(base, SHEET1, "D1").get("s") + assert orig_style is not None # the number format style index + raw = open_raw(base, extension="xlsx") + raw.sheets["Sales"].set_cell("D1", 777) + out = raw.to_bytes() + d1 = _cell_el(out, SHEET1, "D1") + assert d1.get("s") == orig_style + assert d1.get("t") is None # numeric: no type attribute + wb = _load_wb(out) + assert wb["Sales"]["D1"].value == 777 + assert wb["Sales"]["D1"].number_format == "#,##0.00" + + def test_set_none_clears_value_keeps_cell_and_style(self): + base = _build_workbook() + orig_style = _cell_el(base, SHEET1, "D1").get("s") + raw = open_raw(base, extension="xlsx") + raw.sheets["Sales"].set_cell("D1", None) + assert raw.sheets["Sales"].get_cell("D1") is None + out = raw.to_bytes() + d1 = _cell_el(out, SHEET1, "D1") + assert d1 is not None and d1.get("s") == orig_style + assert len(d1) == 0 # no v / is children left + + def test_bool_and_float_write_roundtrip(self): + raw = open_raw(_build_workbook(), extension="xlsx") + sales = raw.sheets["Sales"] + sales.set_cell("C2", False) + sales.set_cell("C3", 0.125) + wb = _load_wb(raw.to_bytes()) + assert wb["Sales"]["C2"].value is False + assert wb["Sales"]["C3"].value == 0.125 + + def test_unsupported_type_raises(self): + raw = open_raw(_build_workbook(), extension="xlsx") + with pytest.raises(TypeError, match="dict"): + raw.sheets["Sales"].set_cell("A1", {"no": "way"}) + + +# -- the killer test: foreign-feature survival --------------------------------- + + +class TestForeignFeatureSurvival: + def test_injected_parts_and_extlst_survive_raw_edit(self): + injected = _inject_foreign_features(_build_workbook()) + raw = open_raw(injected, extension="xlsx") + raw.sheets["Sales"].set_cell("B3", 777) + out = raw.save() + + # the ONLY part that differs is the edited worksheet + assert _diff_parts(injected, out) == {SHEET1} + out_parts, in_parts = _parts(out), _parts(injected) + for name in FOREIGN_PARTS + [SHEET2, CONTENT_TYPES, WORKBOOK_RELS]: + assert out_parts[name] == in_parts[name], f"{name} was not preserved" + assert b"sparklineGroups" in out_parts[SHEET2] + assert "홍길동".encode("utf-8") in out_parts["customXml/item1.xml"] + # and the edit itself landed + assert _load_wb(out)["Sales"]["B3"].value == 777 + + def test_openpyxl_roundtrip_destroys_the_same_features(self): + """The contrast: the exact features we preserve, openpyxl drops.""" + injected = _inject_foreign_features(_build_workbook()) + wb = _load_wb(injected) + buf = io.BytesIO() + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + wb.save(buf) + survivors = set(_parts(buf.getvalue())) + assert "xl/charts/style1.xml" not in survivors + assert "xl/charts/colors1.xml" not in survivors + assert "customXml/item1.xml" not in survivors + s2 = _parts(buf.getvalue()).get(SHEET2, b"") + assert b"sparklineGroups" not in s2 + + def test_sparkline_extlst_survives_editing_its_own_sheet(self): + """Even when the sparkline's OWN sheet is edited (so its XML is + re-serialized), the extLst rides along in the tree.""" + injected = _inject_foreign_features(_build_workbook()) + raw = open_raw(injected, extension="xlsx") + raw.sheets["Notes"].set_cell("A2", "edited") + out = raw.to_bytes() + s2 = _parts(out)[SHEET2] + assert b"sparklineGroups" in s2 + assert b"Sales!B2:B3" in s2 + assert _load_wb(out)["Notes"]["A2"].value == "edited" + + def test_image_media_part_survives(self): + pytest.importorskip("PIL") + base = _build_workbook(with_image=True) + assert "xl/media/image1.png" in _parts(base) + raw = open_raw(base, extension="xlsx") + raw.sheets["Sales"].set_cell("B3", 1) + out = raw.to_bytes() + assert _parts(out)["xl/media/image1.png"] == _parts(base)["xl/media/image1.png"] + + +# -- formulas ------------------------------------------------------------------- + + +class TestFormulaOverwrite: + def test_overwrite_removes_formula_and_forces_recalc(self): + data = _excel_flavoured(_build_workbook()) # calcPr WITHOUT fullCalcOnLoad + raw = open_raw(data, extension="xlsx") + sales = raw.sheets["Sales"] + assert sales.get_formula("B4") == "SUM(B2:B3)" + sales.set_cell("B4", 42) + assert sales.get_formula("B4") is None + assert sales.get_cell("B4") == 42 + out = raw.to_bytes() + assert b"" not in _parts(out)[SHEET1] # formula gone + assert re.search(rb']*fullCalcOnLoad="1"', _parts(out)[WORKBOOK]) + assert _load_wb(out)["Sales"]["B4"].value == 42 # the literal + + def test_calc_pr_created_after_sheets_when_absent(self): + base = _excel_flavoured(_build_workbook()) + stripped = re.sub(rb"]*/>", b"", _parts(base)[WORKBOOK]) + assert b"calcPr" not in stripped + data = _rezip(base, replace={WORKBOOK: stripped}) + raw = open_raw(data, extension="xlsx") + raw.sheets["Sales"].set_cell("B4", 5) + out = raw.to_bytes() + out_wb = _parts(out)[WORKBOOK] + assert b'' in out_wb + assert out_wb.index(b" out_wb.index(b"") + _load_wb(out) # still a valid workbook + + def test_plain_edit_with_formulas_elsewhere_forces_recalc(self): + """B4's =SUM depends on B2 — editing B2 must trigger recalc-on-load + even though B2 itself has no formula.""" + data = _excel_flavoured(_build_workbook()) + raw = open_raw(data, extension="xlsx") + raw.sheets["Sales"].set_cell("B2", 999) + out = raw.to_bytes() + assert re.search(rb'fullCalcOnLoad="1"', _parts(out)[WORKBOOK]) + assert _diff_parts(data, out) == {SHEET1, WORKBOOK} + + +# -- appending ------------------------------------------------------------------ + + +class TestAppendRows: + def test_append_rows_values_and_dimension(self): + base = _build_workbook() + raw = open_raw(base, extension="xlsx") + raw.sheets["Sales"].append_rows([["North", 300], ["South", 400.5, True]]) + out = raw.to_bytes() + + wb = _load_wb(out) + ws = wb["Sales"] + assert ws["A6"].value == "North" and ws["B6"].value == 300 + assert ws["A7"].value == "South" and ws["B7"].value == 400.5 + assert ws["C7"].value is True + + assert b'ref="A1:D7"' in _parts(out)[SHEET1] # dimension extended + raw2 = open_raw(out, extension="xlsx") + assert raw2.sheets["Sales"].dimensions == (7, 4) + assert raw2.sheets["Sales"].get_cell("B7") == 400.5 + + def test_append_sparse_rows_skips_none(self): + raw = open_raw(_build_workbook(), extension="xlsx") + raw.sheets["Notes"].append_rows([[None, "only-b"]]) + out = raw.to_bytes() + assert _cell_el(out, SHEET2, "A2") is None # sparse: A2 not stored + assert _load_wb(out)["Notes"]["B2"].value == "only-b" + + +# -- inline strings --------------------------------------------------------------- + + +class TestInlineStrings: + def test_korean_and_xml_specials_roundtrip(self): + base = _build_workbook() + raw = open_raw(base, extension="xlsx") + text = "한글 데이터 <&> \"quote\" 'tick' 100%" + raw.sheets["Notes"].set_cell("A2", text) + out = raw.to_bytes() + assert _load_wb(out)["Notes"]["A2"].value == text + raw2 = open_raw(out, extension="xlsx") + assert raw2.sheets["Notes"].get_cell("A2") == text + # inlineStr writes never create a shared-string table + assert SHARED_STRINGS not in _parts(out) + + def test_leading_trailing_space_preserved(self): + raw = open_raw(_build_workbook(), extension="xlsx") + raw.sheets["Notes"].set_cell("A3", " padded ") + out = raw.to_bytes() + assert _load_wb(out)["Notes"]["A3"].value == " padded " + + +# -- ordering invariants ------------------------------------------------------------ + + +class TestOrderingInvariants: + def test_new_rows_inserted_sorted_by_r(self): + """Excel requires ascending @r; creating a row between existing + rows must keep the order.""" + raw = open_raw(_build_workbook(), extension="xlsx") + notes = raw.sheets["Notes"] + notes.set_cell("A5", "five") + notes.set_cell("A3", "three") # lands BETWEEN rows 1 and 5 + out = raw.to_bytes() + rows = [ + r.get("r") + for r in _sheet_root(out, SHEET2).findall("s:sheetData/s:row", NS) + ] + assert rows == ["1", "3", "5"] + wb = _load_wb(out) + assert wb["Notes"]["A3"].value == "three" + assert wb["Notes"]["A5"].value == "five" + + def test_new_cells_inserted_sorted_within_row(self): + raw = open_raw(_build_workbook(), extension="xlsx") + sales = raw.sheets["Sales"] + sales.set_cell("D2", 1) + sales.set_cell("C2", 2) # lands BETWEEN B2 and D2 + out = raw.to_bytes() + row2 = _sheet_root(out, SHEET1).findall("s:sheetData/s:row", NS)[1] + assert [c.get("r") for c in row2] == ["A2", "B2", "C2", "D2"] diff --git a/tests/unit/services/test_crypto_service.py b/tests/unit/services/test_crypto_service.py index cffa7de..1504193 100644 --- a/tests/unit/services/test_crypto_service.py +++ b/tests/unit/services/test_crypto_service.py @@ -8,7 +8,7 @@ import pytest -from contextifier.services.crypto_service import decrypt_if_encrypted, is_encrypted +from contextifier.services.crypto_service import decrypt_if_encrypted # ═══════════════════════════════════════════════════════════════════════════ diff --git a/tests/unit/services/test_image_service.py b/tests/unit/services/test_image_service.py index c6f50b6..4f7b8a8 100644 --- a/tests/unit/services/test_image_service.py +++ b/tests/unit/services/test_image_service.py @@ -263,8 +263,10 @@ def worker(name: str) -> None: t1 = threading.Thread(target=worker, args=("t1",)) t2 = threading.Thread(target=worker, args=("t2",)) - t1.start(); t1.join() - t2.start(); t2.join() + t1.start() + t1.join() + t2.start() + t2.join() # Both should get img_0001 (independent counters) assert "img_0001" in paths["t1"] diff --git a/tests/unit/services/test_table_service.py b/tests/unit/services/test_table_service.py index e3d200c..f44473f 100644 --- a/tests/unit/services/test_table_service.py +++ b/tests/unit/services/test_table_service.py @@ -3,11 +3,10 @@ from __future__ import annotations -import pytest -from contextifier.config import ProcessingConfig, TableConfig +from contextifier.config import ProcessingConfig from contextifier.services.table_service import TableService -from contextifier.types import OutputFormat, TableData, TableCell +from contextifier.types import TableData, TableCell def _simple_table() -> TableData: diff --git a/tests/unit/test_cached_processor.py b/tests/unit/test_cached_processor.py index 8cb19c7..73144ba 100644 --- a/tests/unit/test_cached_processor.py +++ b/tests/unit/test_cached_processor.py @@ -3,11 +3,9 @@ from __future__ import annotations -import json from pathlib import Path -from unittest.mock import MagicMock, patch +from unittest.mock import MagicMock -import pytest from contextifier.cached_processor import ( CachedDocumentProcessor, diff --git a/tests/unit/test_config.py b/tests/unit/test_config.py index fd0ce9a..81d95a9 100644 --- a/tests/unit/test_config.py +++ b/tests/unit/test_config.py @@ -8,8 +8,6 @@ from contextifier.config import ( ChunkingConfig, ProcessingConfig, - TagConfig, - ImageConfig, ) from contextifier.errors import ConfigurationError diff --git a/tests/unit/test_security.py b/tests/unit/test_security.py index b573e97..4f77c14 100644 --- a/tests/unit/test_security.py +++ b/tests/unit/test_security.py @@ -117,7 +117,7 @@ def test_nested_dotdot_blocked(self, tmp_path: Path) -> None: def test_normal_subdir_allowed(self, tmp_path: Path) -> None: backend = LocalStorageBackend(str(tmp_path)) full_path = str(tmp_path / "subdir" / "file.txt") - result = backend.save(b"data", full_path) + backend.save(b"data", full_path) assert (tmp_path / "subdir" / "file.txt").exists()