add serializer, internal marking of chunkers, typing

bridgetmcg · bridgetmcg · commit 10e9ed8a0f28 · 2025-10-22T16:39:21.000-04:00
diff --git a/docling_core/transforms/chunker/__init__.py b/docling_core/transforms/chunker/__init__.py
@@ -6,7 +6,6 @@
 """Define the chunker types."""
 
 from docling_core.transforms.chunker.base import BaseChunk, BaseChunker, BaseMeta
-from docling_core.transforms.chunker.base_code_chunker import CodeChunker
 from docling_core.transforms.chunker.code_chunk_utils.chunk_utils import (
     ChunkBuilder,
     ChunkMetadataBuilder,
@@ -21,9 +20,9 @@
     NoOpCodeChunkingStrategy,
 )
 from docling_core.transforms.chunker.hierarchical_chunker import (
-    ChunkType,
     CodeChunk,
     CodeChunkingStrategy,
+    CodeChunkType,
     CodeDocMeta,
     DocChunk,
     DocMeta,
diff --git a/docling_core/transforms/chunker/base_code_chunker.py b/docling_core/transforms/chunker/base_code_chunker.py
@@ -19,7 +19,7 @@
 from docling_core.types.doc.labels import DocItemLabel
 
 
-class CodeChunker(BaseChunker):
+class _CodeChunker(BaseChunker):
     """Data model for code chunker."""
 
     language: Language
diff --git a/docling_core/transforms/chunker/code_chunk_utils/chunk_utils.py b/docling_core/transforms/chunker/code_chunk_utils/chunk_utils.py
@@ -4,8 +4,8 @@
 from tree_sitter import Node
 
 from docling_core.transforms.chunker.hierarchical_chunker import (
-    ChunkType,
     CodeChunk,
+    CodeChunkType,
     CodeDocMeta,
 )
 from docling_core.types.doc.document import DocumentOrigin
@@ -99,7 +99,7 @@ def build_function_metadata(
             end_line=end_line,
             end_line_signature=signature_end_line,
             origin=self.origin,
-            chunk_type=ChunkType.FUNCTION,
+            chunk_type=CodeChunkType.FUNCTION,
         )
 
     def build_class_metadata(
@@ -119,7 +119,7 @@ def build_class_metadata(
             end_line=end_line,
             end_line_signature=end_line,
             origin=self.origin,
-            chunk_type=ChunkType.CLASS,
+            chunk_type=CodeChunkType.CLASS,
         )
 
     def build_preamble_metadata(
@@ -131,7 +131,7 @@ def build_preamble_metadata(
             start_line=start_line,
             end_line=end_line,
             origin=self.origin,
-            chunk_type=ChunkType.PREAMBLE,
+            chunk_type=CodeChunkType.PREAMBLE,
         )
 
     def calculate_line_numbers(
diff --git a/docling_core/transforms/chunker/code_chunking_strategy.py b/docling_core/transforms/chunker/code_chunking_strategy.py
@@ -1,18 +1,18 @@
 from typing import Any, Dict, Iterator, Optional
 
-from docling_core.transforms.chunker.base_code_chunker import CodeChunker
+from docling_core.transforms.chunker.base_code_chunker import _CodeChunker
 from docling_core.transforms.chunker.code_chunk_utils.utils import Language
 from docling_core.transforms.chunker.hierarchical_chunker import (
-    ChunkType,
     CodeChunk,
+    CodeChunkType,
     CodeDocMeta,
 )
 from docling_core.transforms.chunker.language_code_chunkers import (
-    CFunctionChunker,
-    JavaFunctionChunker,
-    JavaScriptFunctionChunker,
-    PythonFunctionChunker,
-    TypeScriptFunctionChunker,
+    _CFunctionChunker,
+    _JavaFunctionChunker,
+    _JavaScriptFunctionChunker,
+    _PythonFunctionChunker,
+    _TypeScriptFunctionChunker,
 )
 from docling_core.types.doc.base import Size
 from docling_core.types.doc.document import (
@@ -183,15 +183,15 @@ class CodeChunkingStrategyFactory:
     """Factory for creating language-specific code chunking strategies."""
 
     @staticmethod
-    def create_chunker(language: Language, **kwargs: Any) -> CodeChunker:
+    def create_chunker(language: Language, **kwargs: Any) -> _CodeChunker:
         """Create a language-specific code chunker."""
 
         chunker_map = {
-            Language.PYTHON: PythonFunctionChunker,
-            Language.TYPESCRIPT: TypeScriptFunctionChunker,
-            Language.JAVASCRIPT: JavaScriptFunctionChunker,
-            Language.C: CFunctionChunker,
-            Language.JAVA: JavaFunctionChunker,
+            Language.PYTHON: _PythonFunctionChunker,
+            Language.TYPESCRIPT: _TypeScriptFunctionChunker,
+            Language.JAVASCRIPT: _JavaScriptFunctionChunker,
+            Language.C: _CFunctionChunker,
+            Language.JAVA: _JavaFunctionChunker,
         }
 
         chunker_class = chunker_map.get(language)
@@ -208,9 +208,9 @@ def __init__(self, **chunker_kwargs: Any):
         """Initialize the strategy with optional chunker parameters."""
 
         self.chunker_kwargs = chunker_kwargs
-        self._chunker_cache: Dict[Language, CodeChunker] = {}
+        self._chunker_cache: Dict[Language, _CodeChunker] = {}
 
-    def _get_chunker(self, language: Language) -> CodeChunker:
+    def _get_chunker(self, language: Language) -> _CodeChunker:
         """Get or create a chunker for the given language."""
 
         if language not in self._chunker_cache:
@@ -238,10 +238,12 @@ def chunk_code_item(
             filename = original_doc.origin.filename or "code_chunk"
             mimetype = original_doc.origin.mimetype or "text/plain"
             binary_hash = _create_hash(code_text)
+            uri = getattr(original_doc.origin, "uri", None)
         else:
             filename = "code_chunk"
             mimetype = "text/plain"
             binary_hash = _create_hash(code_text)
+            uri = None
 
         if original_item and hasattr(original_item, "self_ref"):
             self_ref = original_item.self_ref
@@ -255,7 +257,7 @@ def chunk_code_item(
             texts=[code_item],
             pages={0: PageItem(page_no=0, size=Size(width=612.0, height=792.0))},
             origin=DocumentOrigin(
-                filename=filename, mimetype=mimetype, binary_hash=binary_hash
+                filename=filename, mimetype=mimetype, binary_hash=binary_hash, uri=uri
             ),
         )
 
@@ -279,7 +281,7 @@ def chunk_code_item(
             return
 
         meta = CodeDocMeta(
-            chunk_type=ChunkType.CODE_BLOCK,
+            chunk_type=CodeChunkType.CODE_BLOCK,
             start_line=1,
             end_line=len(code_text.splitlines()),
         )
diff --git a/docling_core/transforms/chunker/hierarchical_chunker.py b/docling_core/transforms/chunker/hierarchical_chunker.py
@@ -10,7 +10,16 @@
 import logging
 import re
 from enum import Enum
-from typing import Any, ClassVar, Final, Iterator, Literal, Optional, Protocol
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    ClassVar,
+    Final,
+    Iterator,
+    Literal,
+    Optional,
+    Protocol,
+)
 
 from pydantic import ConfigDict, Field, StringConstraints, field_validator
 from typing_extensions import Annotated, override
@@ -119,7 +128,7 @@ def check_version_is_compatible(cls, v: str) -> str:
 
 
 class CodeDocMeta(DocMeta):
-    """Data model for CodeChunker metadata."""
+    """Data model for code chunk metadata."""
 
     doc_items: Optional[list[DocItem]] = Field(default=None, alias=_KEY_DOC_ITEMS)
     part_name: Optional[str] = Field(default=None)
@@ -137,7 +146,7 @@ class CodeChunk(BaseChunk):
     meta: CodeDocMeta
 
 
-class ChunkType(str, Enum):
+class CodeChunkType(str, Enum):
     """Chunk type"""
 
     FUNCTION = "function"
@@ -157,6 +166,12 @@ def chunk_code_item(
         ...
 
 
+if TYPE_CHECKING:
+    CodeChunkingStrategyType = CodeChunkingStrategy
+else:
+    CodeChunkingStrategyType = Any
+
+
 class DocChunk(BaseChunk):
     """Data model for document chunks."""
 
@@ -248,7 +263,7 @@ class HierarchicalChunker(BaseChunker):
     model_config = ConfigDict(arbitrary_types_allowed=True)
 
     serializer_provider: BaseSerializerProvider = ChunkingSerializerProvider()
-    code_chunking_strategy: Optional[Any] = Field(default=None)
+    code_chunking_strategy: Optional[CodeChunkingStrategyType] = Field(default=None)
 
     # deprecated:
     merge_list_items: Annotated[bool, Field(deprecated=True)] = True
@@ -308,14 +323,21 @@ def chunk(
                     )
 
                     if language:
-                        for code_chunk in self.code_chunking_strategy.chunk_code_item(
-                            item.text,
-                            language,
-                            original_doc=dl_doc,
-                            original_item=item,
-                            **kwargs,
-                        ):
-                            yield code_chunk
+                        ser_res = my_doc_ser.serialize(item=item, visited=visited)
+                        if ser_res.text:
+                            code_text = self._strip_markdown_code_formatting(
+                                ser_res.text
+                            )
+                            for (
+                                code_chunk
+                            ) in self.code_chunking_strategy.chunk_code_item(
+                                code_text,
+                                language,
+                                original_doc=dl_doc,
+                                original_item=item,
+                                **kwargs,
+                            ):
+                                yield code_chunk
                         continue
 
                 ser_res = my_doc_ser.serialize(item=item, visited=visited)
@@ -335,3 +357,14 @@ def chunk(
                     ),
                 )
                 yield c
+
+    def _strip_markdown_code_formatting(self, text: str) -> str:
+        """Strip markdown code block formatting from text."""
+        if not text.startswith("```") or not text.endswith("```"):
+            return text
+
+        lines = text.split("\n")
+        if len(lines) >= 3 and lines[0].startswith("```") and lines[-1] == "```":
+            return "\n".join(lines[1:-1])
+
+        return text
diff --git a/docling_core/transforms/chunker/language_code_chunkers.py b/docling_core/transforms/chunker/language_code_chunkers.py
@@ -4,7 +4,7 @@
 from tree_sitter import Node, Tree
 from typing_extensions import override
 
-from docling_core.transforms.chunker.base_code_chunker import CodeChunker
+from docling_core.transforms.chunker.base_code_chunker import _CodeChunker
 from docling_core.transforms.chunker.code_chunk_utils.utils import (
     Language,
     _get_default_tokenizer,
@@ -16,7 +16,7 @@
 from docling_core.transforms.chunker.tokenizer.base import BaseTokenizer
 
 
-class PythonFunctionChunker(CodeChunker):
+class _PythonFunctionChunker(_CodeChunker):
 
     language: Language = Language.PYTHON
     ts_language: Any = Field(default=None)
@@ -151,7 +151,7 @@ def _is_local_assignment(self, identifier_node: Node) -> bool:
         return False
 
 
-class TypeScriptFunctionChunker(CodeChunker):
+class _TypeScriptFunctionChunker(_CodeChunker):
     language: Language = Language.TYPESCRIPT
     ts_language: Any = Field(default=None)
     parser: Any = Field(default=None)
@@ -232,12 +232,12 @@ def _get_module_variables(self, tree: Tree) -> Dict[str, Node]:
         return {}
 
 
-class JavaScriptFunctionChunker(TypeScriptFunctionChunker):
+class _JavaScriptFunctionChunker(_TypeScriptFunctionChunker):
     def __init__(self, **data):
         super().__init__(language=Language.JAVASCRIPT)
 
 
-class CFunctionChunker(CodeChunker):
+class _CFunctionChunker(_CodeChunker):
     language: Language = Language.C
     ts_language: Any = Field(default=None)
     parser: Any = Field(default=None)
@@ -389,7 +389,7 @@ def collect_identifiers(node, depth=0):
         return used_macros
 
 
-class JavaFunctionChunker(CodeChunker):
+class _JavaFunctionChunker(_CodeChunker):
 
     language: Language = Language.JAVA
     ts_language: Any = Field(default=None)
diff --git a/test/test_code_chunker.py b/test/test_code_chunker.py
@@ -9,13 +9,10 @@
 
 from docling_core.transforms.chunker.base_code_chunker import CodeChunk
 from docling_core.transforms.chunker.code_chunk_utils.utils import Language
-from docling_core.transforms.chunker.language_code_chunkers import (
-    CFunctionChunker,
-    JavaFunctionChunker,
-    JavaScriptFunctionChunker,
-    PythonFunctionChunker,
-    TypeScriptFunctionChunker,
+from docling_core.transforms.chunker.code_chunking_strategy import (
+    DefaultCodeChunkingStrategy,
 )
+from docling_core.transforms.chunker.hierarchical_chunker import HierarchicalChunker
 from docling_core.types.doc import DoclingDocument, DocumentOrigin
 from docling_core.types.doc.labels import DocItemLabel
 from docling_core.utils.legacy import _create_hash
@@ -96,31 +93,41 @@ def create_documents_from_repository(
         "Java",
         "/test/data/chunker_repo/repos/acmeair",
         "https://github.com/acmeair/acmeair",
-        lambda: JavaFunctionChunker(max_tokens=5000),
+        lambda: HierarchicalChunker(
+            code_chunking_strategy=DefaultCodeChunkingStrategy(max_tokens=5000)
+        ),
     ),
     (
         "TypeScript",
         "/test/data/chunker_repo/repos/outline",
         "https://github.com/outline/outline",
-        lambda: TypeScriptFunctionChunker(max_tokens=5000),
+        lambda: HierarchicalChunker(
+            code_chunking_strategy=DefaultCodeChunkingStrategy(max_tokens=5000)
+        ),
     ),
     (
         "JavaScript",
         "/test/data/chunker_repo/repos/jquery",
         "https://github.com/jquery/jquery",
-        lambda: JavaScriptFunctionChunker(max_tokens=5000),
+        lambda: HierarchicalChunker(
+            code_chunking_strategy=DefaultCodeChunkingStrategy(max_tokens=5000)
+        ),
     ),
     (
         "Python",
         "/test/data/chunker_repo/repos/docling",
         "https://github.com/docling-project/docling",
-        lambda: PythonFunctionChunker(max_tokens=5000),
+        lambda: HierarchicalChunker(
+            code_chunking_strategy=DefaultCodeChunkingStrategy(max_tokens=5000)
+        ),
     ),
     (
         "C",
         "/test/data/chunker_repo/repos/json-c",
         "https://github.com/json-c/json-c",
-        lambda: CFunctionChunker(max_tokens=5000),
+        lambda: HierarchicalChunker(
+            code_chunking_strategy=DefaultCodeChunkingStrategy(max_tokens=5000)
+        ),
     ),
 ]
 
