kreuzberg 3.8.1__py3-none-any.whl → 3.9.0__py3-none-any.whl

kreuzberg/_extractors/_image.py

@@ -3,6 +3,7 @@ from __future__ import annotations
 import contextlib
 import os
 import tempfile
+from dataclasses import asdict
 from pathlib import Path
 from typing import TYPE_CHECKING, ClassVar
 
@@ -88,17 +89,17 @@ class ImageExtractor(Extractor):
             config = (
                 self.config.ocr_config if isinstance(self.config.ocr_config, TesseractConfig) else TesseractConfig()
             )
-            result = backend.process_file_sync(path, **config.__dict__)
+            result = backend.process_file_sync(path, **asdict(config))
         elif self.config.ocr_backend == "paddleocr":
             paddle_config = (
                 self.config.ocr_config if isinstance(self.config.ocr_config, PaddleOCRConfig) else PaddleOCRConfig()
             )
-            result = backend.process_file_sync(path, **paddle_config.__dict__)
+            result = backend.process_file_sync(path, **asdict(paddle_config))
         elif self.config.ocr_backend == "easyocr":
             easy_config = (
                 self.config.ocr_config if isinstance(self.config.ocr_config, EasyOCRConfig) else EasyOCRConfig()
             )
-            result = backend.process_file_sync(path, **easy_config.__dict__)
+            result = backend.process_file_sync(path, **asdict(easy_config))
         else:
             raise NotImplementedError(f"Sync OCR not implemented for {self.config.ocr_backend}")
         return self._apply_quality_processing(result)

kreuzberg/_extractors/_pdf.py

@@ -3,6 +3,7 @@ from __future__ import annotations
 import contextlib
 import os
 import tempfile
+from dataclasses import asdict
 from multiprocessing import cpu_count
 from pathlib import Path
 from re import Pattern
@@ -58,9 +59,13 @@ class PDFExtractor(Extractor):
         result: ExtractionResult | None = None
 
         if not self.config.force_ocr:
-            content = await self._extract_pdf_searchable_text(path)
-            if self._validate_extracted_text(content):
-                result = ExtractionResult(content=content, mime_type=PLAIN_TEXT_MIME_TYPE, metadata={}, chunks=[])
+            try:
+                content = await self._extract_pdf_searchable_text(path)
+                if self._validate_extracted_text(content):
+                    result = ExtractionResult(content=content, mime_type=PLAIN_TEXT_MIME_TYPE, metadata={}, chunks=[])
+            except ParsingError:
+                # If searchable text extraction fails, continue to OCR or empty result
+                pass
 
         if not result and self.config.ocr_backend is not None:
             result = await self._extract_pdf_text_with_ocr(path, self.config.ocr_backend)
@@ -73,7 +78,7 @@ class PDFExtractor(Extractor):
         if self.config.extract_tables:
             # GMFT is optional dependency
             try:
-                from kreuzberg._gmft import extract_tables
+                from kreuzberg._gmft import extract_tables  # noqa: PLC0415
 
                 result.tables = await extract_tables(path, self.config.gmft_config)
             except ImportError:
@@ -112,16 +117,19 @@ class PDFExtractor(Extractor):
 
     def extract_path_sync(self, path: Path) -> ExtractionResult:
         """Pure sync implementation of PDF extraction from path."""
-        text = self._extract_pdf_searchable_text_sync(path)
+        try:
+            text = self._extract_pdf_searchable_text_sync(path)
+        except ParsingError:
+            text = ""
 
-        if self.config.force_ocr or not self._validate_extracted_text(text):
+        if (self.config.force_ocr or not self._validate_extracted_text(text)) and self.config.ocr_backend is not None:
             text = self._extract_pdf_with_ocr_sync(path)
 
         tables = []
         if self.config.extract_tables:
             # GMFT is optional dependency
             try:
-                from kreuzberg._gmft import extract_tables_sync
+                from kreuzberg._gmft import extract_tables_sync  # noqa: PLC0415
 
                 tables = extract_tables_sync(path)
             except ImportError:
@@ -248,9 +256,10 @@ class PDFExtractor(Extractor):
             *[backend.process_image(image, **self.config.get_config_dict()) for image in images],
             batch_size=cpu_count(),
         )
-        return ExtractionResult(
-            content="\n".join([v.content for v in ocr_results]), mime_type=PLAIN_TEXT_MIME_TYPE, metadata={}, chunks=[]
-        )
+        # Use list comprehension and join for efficient string building
+        content = "\n".join(result.content for result in ocr_results)
+
+        return ExtractionResult(content=content, mime_type=PLAIN_TEXT_MIME_TYPE, metadata={}, chunks=[])
 
     @staticmethod
     async def _extract_pdf_searchable_text(input_file: Path) -> str:
@@ -269,22 +278,24 @@ class PDFExtractor(Extractor):
         try:
             with pypdfium_file_lock(input_file):
                 document = await run_sync(pypdfium2.PdfDocument, str(input_file))
-                text_parts = []
+                pages_content = []
                 page_errors = []
 
                 for i, page in enumerate(cast("pypdfium2.PdfDocument", document)):
                     try:
                         text_page = page.get_textpage()
-                        text_parts.append(text_page.get_text_bounded())
+                        page_content = text_page.get_text_bounded()
+                        pages_content.append(page_content)
                     except Exception as e:  # noqa: PERF203, BLE001
                         page_errors.append({"page": i + 1, "error": str(e)})
-                        text_parts.append(f"[Error extracting page {i + 1}]")
+                        pages_content.append(f"[Error extracting page {i + 1}]")
 
-                text = "\n".join(text_parts)
+                text = "\n".join(pages_content)
+                has_content = bool(text.strip())
 
-                if page_errors and text_parts:
+                if page_errors and has_content:
                     return normalize_spaces(text)
-                if not text_parts:
+                if not has_content:
                     raise ParsingError(
                         "Could not extract any text from PDF",
                         context=create_error_context(
@@ -315,14 +326,14 @@ class PDFExtractor(Extractor):
         try:
             with pypdfium_file_lock(path):
                 pdf = pypdfium2.PdfDocument(str(path))
-                text_parts = []
+                pages_text = []
                 for page in pdf:
                     text_page = page.get_textpage()
                     text = text_page.get_text_bounded()
-                    text_parts.append(text)
+                    pages_text.append(text)
                     text_page.close()
                     page.close()
-                return "".join(text_parts)
+                return "\n".join(pages_text)
         except Exception as e:
             raise ParsingError(f"Failed to extract PDF text: {e}") from e
         finally:
@@ -378,22 +389,22 @@ class PDFExtractor(Extractor):
             config = (
                 self.config.ocr_config if isinstance(self.config.ocr_config, TesseractConfig) else TesseractConfig()
             )
-            results = backend.process_batch_sync(paths, **config.__dict__)
+            results = backend.process_batch_sync(paths, **asdict(config))
         elif self.config.ocr_backend == "paddleocr":
             paddle_config = (
                 self.config.ocr_config if isinstance(self.config.ocr_config, PaddleOCRConfig) else PaddleOCRConfig()
             )
-            results = backend.process_batch_sync(paths, **paddle_config.__dict__)
+            results = backend.process_batch_sync(paths, **asdict(paddle_config))
         elif self.config.ocr_backend == "easyocr":
             easy_config = (
                 self.config.ocr_config if isinstance(self.config.ocr_config, EasyOCRConfig) else EasyOCRConfig()
             )
-            results = backend.process_batch_sync(paths, **easy_config.__dict__)
+            results = backend.process_batch_sync(paths, **asdict(easy_config))
         else:
             raise NotImplementedError(f"Sync OCR not implemented for {self.config.ocr_backend}")
 
-        text_parts = [r.content for r in results]
-        return "\n\n".join(text_parts)
+        # Use list comprehension and join for efficient string building
+        return "\n\n".join(result.content for result in results)
 
     def _extract_with_playa_sync(self, path: Path, fallback_text: str) -> str:
         """Extract text using playa for better structure preservation."""
@@ -401,14 +412,14 @@ class PDFExtractor(Extractor):
             content = path.read_bytes()
             document = parse(content, max_workers=1)
 
-            text_parts = []
+            # Extract text while preserving structure
+            pages_text = []
             for page in document.pages:
-                # Extract text while preserving structure
                 page_text = page.extract_text()
                 if page_text and page_text.strip():
-                    text_parts.append(page_text)
+                    pages_text.append(page_text)
 
-            if text_parts:
-                return "\n\n".join(text_parts)
+            if pages_text:
+                return "\n\n".join(pages_text)
 
         return fallback_text

kreuzberg/_extractors/_spread_sheet.py

@@ -2,13 +2,16 @@ from __future__ import annotations
 
 import contextlib
 import csv
+import os
 import sys
+import tempfile
 from datetime import date, datetime, time, timedelta
 from io import StringIO
 from pathlib import Path
 from typing import Any
 
 from anyio import Path as AsyncPath
+from PIL import Image
 from python_calamine import CalamineWorkbook
 
 from kreuzberg._extractors._base import Extractor
@@ -68,9 +71,6 @@ class SpreadSheetExtractor(Extractor):
 
     def extract_bytes_sync(self, content: bytes) -> ExtractionResult:
         """Pure sync implementation of extract_bytes."""
-        import os
-        import tempfile
-
         fd, temp_path = tempfile.mkstemp(suffix=".xlsx")
 
         try:
@@ -198,9 +198,9 @@ class SpreadSheetExtractor(Extractor):
         """Enhanced sheet processing with better table structure preservation."""
         try:
             # pandas is optional dependency
-            import pandas as pd
+            import pandas as pd  # noqa: PLC0415
 
-            from kreuzberg._utils._table import enhance_table_markdown
+            from kreuzberg._utils._table import enhance_table_markdown  # noqa: PLC0415
 
             sheet = workbook.get_sheet_by_name(sheet_name)
             data = sheet.to_python()
@@ -218,9 +218,7 @@ class SpreadSheetExtractor(Extractor):
                 return f"## {sheet_name}\n\n*No data*"
 
             # Create a mock TableData for enhanced formatting
-            from PIL import Image
-
-            from kreuzberg._types import TableData
+            from kreuzberg._types import TableData  # noqa: PLC0415
 
             # Create a 1x1 transparent image as placeholder
             placeholder_image = Image.new("RGBA", (1, 1), (0, 0, 0, 0))

kreuzberg/_extractors/_structured.py

@@ -1,8 +1,22 @@
 from __future__ import annotations
 
 import json
+import sys
 from typing import TYPE_CHECKING, Any, ClassVar
 
+if sys.version_info >= (3, 11):
+    import tomllib
+else:
+    try:
+        import tomli as tomllib  # type: ignore[import-not-found]
+    except ImportError:
+        tomllib = None
+
+try:
+    import yaml
+except ImportError:
+    yaml = None
+
 from anyio import Path as AsyncPath
 
 from kreuzberg._extractors._base import Extractor
@@ -14,6 +28,9 @@ from kreuzberg._utils._sync import run_sync
 if TYPE_CHECKING:
     from pathlib import Path
 
+# Define text field keywords as a set for O(1) membership testing
+_TEXT_FIELD_KEYWORDS = frozenset({"title", "name", "subject", "description", "content", "body", "text", "message"})
+
 
 class StructuredDataExtractor(Extractor):
     SUPPORTED_MIME_TYPES: ClassVar[set[str]] = {
@@ -41,41 +58,34 @@ class StructuredDataExtractor(Extractor):
             if self.mime_type in {JSON_MIME_TYPE, "text/json"}:
                 data = json.loads(text_content)
             elif self.mime_type in {TOML_MIME_TYPE, "text/toml"}:
-                try:
-                    import tomllib  # type: ignore[import-not-found]
-                except ImportError:
-                    try:
-                        import tomli as tomllib  # type: ignore[import-not-found]
-                    except ImportError:
-                        return ExtractionResult(
-                            content=normalize_spaces(text_content),
-                            mime_type=PLAIN_TEXT_MIME_TYPE,
-                            metadata={"warning": "tomllib/tomli not available, returning raw text"},
-                            chunks=[],
-                        )
+                if tomllib is None:
+                    return ExtractionResult(
+                        content=normalize_spaces(text_content),
+                        mime_type=PLAIN_TEXT_MIME_TYPE,
+                        metadata={"warning": "tomllib/tomli not available, returning raw text"},
+                        chunks=[],
+                    )
                 data = tomllib.loads(text_content)
             else:
-                try:
-                    import yaml
-
-                    data = yaml.safe_load(text_content)
-                except ImportError:
+                if yaml is None:
                     return ExtractionResult(
                         content=normalize_spaces(text_content),
                         mime_type=PLAIN_TEXT_MIME_TYPE,
                         metadata={"warning": "PyYAML not available, returning raw text"},
                         chunks=[],
                     )
+                data = yaml.safe_load(text_content)
 
             text_parts: list[str] = []
             metadata: dict[str, Any] = {}
 
+            # Use match statement for cleaner code and avoid multiple isinstance calls
             if isinstance(data, dict):
-                text_parts.extend(self._extract_from_dict(data, metadata))
+                text_parts = self._extract_from_dict(data, metadata)
             elif isinstance(data, list):
-                text_parts.extend(self._extract_from_list(data, metadata))
+                text_parts = self._extract_from_list(data, metadata)
             else:
-                text_parts.append(str(data))
+                text_parts = [str(data)]
 
             combined_text = "\n".join(text_parts) if text_parts else text_content
 
@@ -86,7 +96,7 @@ class StructuredDataExtractor(Extractor):
                 chunks=[],
             )
 
-        except (ValueError, TypeError, KeyError, AttributeError, UnicodeDecodeError) as e:
+        except (json.JSONDecodeError, ValueError, TypeError) as e:
             return ExtractionResult(
                 content=normalize_spaces(text_content),
                 mime_type=PLAIN_TEXT_MIME_TYPE,
@@ -107,10 +117,9 @@ class StructuredDataExtractor(Extractor):
             if isinstance(value, str) and value.strip():
                 text_parts.append(f"{full_key}: {value}")
 
-                if any(
-                    text_field in key.lower()
-                    for text_field in ["title", "name", "subject", "description", "content", "body", "text", "message"]
-                ):
+                # Check if key contains any text field keywords efficiently
+                key_lower = key.lower()
+                if any(keyword in key_lower for keyword in _TEXT_FIELD_KEYWORDS):
                     metadata[full_key] = value
 
             elif isinstance(value, (int, float, bool)):

kreuzberg/_gmft.py

@@ -1,14 +1,21 @@
 from __future__ import annotations
 
+import io
 import multiprocessing as mp
 import os
 import queue
 import signal
+import time
 import traceback
 from dataclasses import dataclass, field
 from io import StringIO
+from pathlib import Path
 from typing import TYPE_CHECKING, Any, Literal
 
+import anyio
+import msgspec
+from PIL import Image
+
 from kreuzberg._types import TableData
 from kreuzberg._utils._sync import run_sync
 from kreuzberg.exceptions import MissingDependencyError, ParsingError
@@ -20,7 +27,7 @@ if TYPE_CHECKING:
     from pandas import DataFrame
 
 
-@dataclass(unsafe_hash=True)
+@dataclass(unsafe_hash=True, slots=True)
 class GMFTConfig:
     """Configuration options for GMFT.
 
@@ -131,7 +138,7 @@ class GMFTConfig:
     """
 
 
-async def extract_tables(  # noqa: PLR0915
+async def extract_tables(
     file_path: str | PathLike[str], config: GMFTConfig | None = None, use_isolated_process: bool | None = None
 ) -> list[TableData]:
     """Extracts tables from a PDF file.
@@ -151,9 +158,7 @@ async def extract_tables(  # noqa: PLR0915
     Returns:
         A list of table data dictionaries.
     """
-    from pathlib import Path
-
-    from kreuzberg._utils._cache import get_table_cache
+    from kreuzberg._utils._cache import get_table_cache  # noqa: PLC0415
 
     # Determine if we should use isolated process  # ~keep
     if use_isolated_process is None:
@@ -178,7 +183,7 @@ async def extract_tables(  # noqa: PLR0915
     cache_kwargs = {
         "file_info": str(sorted(file_info.items())),
         "extractor": "gmft",
-        "config": str(sorted(config.__dict__.items())),
+        "config": str(sorted(msgspec.to_builtins(config).items())),
     }
 
     table_cache = get_table_cache()
@@ -187,8 +192,6 @@ async def extract_tables(  # noqa: PLR0915
         return cached_result  # type: ignore[no-any-return]
 
     if table_cache.is_processing(**cache_kwargs):
-        import anyio
-
         event = table_cache.mark_processing(**cache_kwargs)
         await anyio.to_thread.run_sync(event.wait)
 
@@ -208,10 +211,13 @@ async def extract_tables(  # noqa: PLR0915
             return result
 
         try:
-            from gmft.auto import AutoTableDetector, AutoTableFormatter  # type: ignore[attr-defined]
-            from gmft.detectors.tatr import TATRDetectorConfig  # type: ignore[attr-defined]
-            from gmft.formatters.tatr import TATRFormatConfig
-            from gmft.pdf_bindings.pdfium import PyPDFium2Document
+            from gmft.auto import (  # type: ignore[attr-defined]  # noqa: PLC0415  # noqa: PLC0415
+                AutoTableDetector,
+                AutoTableFormatter,
+            )
+            from gmft.detectors.tatr import TATRDetectorConfig  # type: ignore[attr-defined]  # noqa: PLC0415
+            from gmft.formatters.tatr import TATRFormatConfig  # noqa: PLC0415  # noqa: PLC0415
+            from gmft.pdf_bindings.pdfium import PyPDFium2Document  # noqa: PLC0415  # noqa: PLC0415
 
             formatter: Any = AutoTableFormatter(  # type: ignore[no-untyped-call]  # type: ignore[no-untyped-call]
                 config=TATRFormatConfig(
@@ -281,9 +287,7 @@ def extract_tables_sync(
     Returns:
         A list of table data dictionaries.
     """
-    from pathlib import Path
-
-    from kreuzberg._utils._cache import get_table_cache
+    from kreuzberg._utils._cache import get_table_cache  # noqa: PLC0415
 
     # Determine if we should use isolated process  # ~keep
     if use_isolated_process is None:
@@ -308,7 +312,7 @@ def extract_tables_sync(
     cache_kwargs = {
         "file_info": str(sorted(file_info.items())),
         "extractor": "gmft",
-        "config": str(sorted(config.__dict__.items())),
+        "config": str(sorted(msgspec.to_builtins(config).items())),
     }
 
     table_cache = get_table_cache()
@@ -324,10 +328,10 @@ def extract_tables_sync(
         return result
 
     try:
-        from gmft.auto import AutoTableDetector, AutoTableFormatter  # type: ignore[attr-defined]
-        from gmft.detectors.tatr import TATRDetectorConfig  # type: ignore[attr-defined]
-        from gmft.formatters.tatr import TATRFormatConfig
-        from gmft.pdf_bindings.pdfium import PyPDFium2Document
+        from gmft.auto import AutoTableDetector, AutoTableFormatter  # type: ignore[attr-defined]  # noqa: PLC0415
+        from gmft.detectors.tatr import TATRDetectorConfig  # type: ignore[attr-defined]  # noqa: PLC0415
+        from gmft.formatters.tatr import TATRFormatConfig  # noqa: PLC0415
+        from gmft.pdf_bindings.pdfium import PyPDFium2Document  # noqa: PLC0415
 
         formatter: Any = AutoTableFormatter(  # type: ignore[no-untyped-call]
             config=TATRFormatConfig(
@@ -396,10 +400,10 @@ def _extract_tables_in_process(
     signal.signal(signal.SIGINT, signal.SIG_IGN)
 
     try:
-        from gmft.auto import AutoTableDetector, AutoTableFormatter  # type: ignore[attr-defined]
-        from gmft.detectors.tatr import TATRDetectorConfig  # type: ignore[attr-defined]
-        from gmft.formatters.tatr import TATRFormatConfig
-        from gmft.pdf_bindings.pdfium import PyPDFium2Document
+        from gmft.auto import AutoTableDetector, AutoTableFormatter  # type: ignore[attr-defined]  # noqa: PLC0415
+        from gmft.detectors.tatr import TATRDetectorConfig  # type: ignore[attr-defined]  # noqa: PLC0415
+        from gmft.formatters.tatr import TATRFormatConfig  # noqa: PLC0415
+        from gmft.pdf_bindings.pdfium import PyPDFium2Document  # noqa: PLC0415
 
         config = GMFTConfig(**config_dict)
 
@@ -435,8 +439,6 @@ def _extract_tables_in_process(
 
             results = []
             for data_frame, cropped_table in zip(dataframes, cropped_tables, strict=False):
-                import io
-
                 img_bytes = io.BytesIO()
                 cropped_image = cropped_table.image()
                 cropped_image.save(img_bytes, format="PNG")
@@ -480,7 +482,7 @@ def _extract_tables_isolated(
         RuntimeError: If extraction fails or times out
     """
     config = config or GMFTConfig()
-    config_dict = config.__dict__.copy()
+    config_dict = msgspec.to_builtins(config)
 
     ctx = mp.get_context("spawn")
     result_queue = ctx.Queue()
@@ -494,7 +496,6 @@ def _extract_tables_isolated(
 
     try:
         # Wait for result with timeout, checking for process death  # ~keep
-        import time
 
         start_time = time.time()
         while True:
@@ -528,12 +529,8 @@ def _extract_tables_isolated(
         if success:
             tables = []
             for table_dict in result:
-                import io
-
-                from PIL import Image
-
                 img = Image.open(io.BytesIO(table_dict["cropped_image_bytes"]))
-                import pandas as pd
+                import pandas as pd  # noqa: PLC0415
 
                 df = pd.read_csv(StringIO(table_dict["df_csv"]))
 
@@ -578,7 +575,7 @@ def _extract_tables_isolated(
 async def _extract_tables_isolated_async(
     file_path: str | PathLike[str],
     config: GMFTConfig | None = None,
-    timeout: float = 300.0,
+    timeout: float = 300.0,  # noqa: ASYNC109
 ) -> list[TableData]:
     """Async version of extract_tables_isolated using asyncio.
 
@@ -593,10 +590,8 @@ async def _extract_tables_isolated_async(
     Raises:
         RuntimeError: If extraction fails or times out
     """
-    import anyio
-
     config = config or GMFTConfig()
-    config_dict = config.__dict__.copy()
+    config_dict = msgspec.to_builtins(config)
 
     ctx = mp.get_context("spawn")
     result_queue = ctx.Queue()
@@ -640,12 +635,8 @@ async def _extract_tables_isolated_async(
         if success:
             tables = []
             for table_dict in result:
-                import io
-
-                from PIL import Image
-
                 img = Image.open(io.BytesIO(table_dict["cropped_image_bytes"]))
-                import pandas as pd
+                import pandas as pd  # noqa: PLC0415
 
                 df = pd.read_csv(StringIO(table_dict["df_csv"]))
 

kreuzberg/_language_detection.py

@@ -23,7 +23,7 @@ except ImportError:
 _CACHE_SIZE = 128
 
 
-@dataclass(frozen=True)
+@dataclass(frozen=True, slots=True)
 class LanguageDetectionConfig:
     """Configuration for language detection.