Coverage for haystack/components/preprocessors/python_code_splitter.py: 95%
291 statements
« prev ^ index » next coverage.py v7.16.1, created at 2026-10-09 14:34 +0000
1# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
2#
3# SPDX-License-Identifier: Apache-2.0
5import ast
6import math
7from copy import deepcopy
8from dataclasses import dataclass, field
9from typing import Any
11from haystack import Document, component, logging
12from haystack.components.preprocessors.document_splitter import DocumentSplitter
14logger = logging.getLogger(__name__)
17@dataclass
18class _CodeUnit:
19 """One syntactic split unit (function, class header, method, imports block, statement, ...)."""
21 source: str
22 start_line: int
23 end_line: int
24 kind: str
25 name: str | None = None
26 class_name: str | None = None
27 class_signature: str | None = None
28 decorators: list[str] = field(default_factory=list)
29 docstring: str | None = None
32@component
33class PythonCodeSplitter:
34 """
35 Split Python source code into syntax-aware chunks.
37 The component parses each source with :mod:`ast` into *units* (module docstring,
38 consecutive ``import`` blocks, top-level functions, class headers, methods, nested
39 classes, and remaining statements) and merges them greedily in source order toward
40 ``max_effective_lines`` per chunk, where effective lines are
41 ``ceil(len(source) / expected_chars_per_line)``. Functions and methods are kept
42 whole; the resulting chunks read top-to-bottom like the original file with comments
43 and blank lines preserved.
45 A function whose effective length exceeds ``oversized_factor * max_effective_lines``
46 is the only case where chunks may overlap: it is broken down with a line-based
47 secondary split (:class:`DocumentSplitter`, ``split_by="line"``) and the resulting
48 pieces carry ``secondary_split=True`` along with the originating function's metadata.
49 The primary split never adds overlap.
51 Per-chunk metadata: ``source_id``, ``split_id``, ``start_line``, ``end_line``,
52 ``unit_kinds``; plus ``include_classes``, ``decorators``, and ``docstrings`` (when
53 ``strip_docstrings=True``) where applicable. ``file_name`` and any other parent
54 document meta are propagated.
56 Usage example:
58 ```python
59 from haystack import Document
60 from haystack.components.preprocessors import PythonCodeSplitter
62 source = '''
63 \"\"\"Example module.\"\"\"
64 from math import sqrt
67 class Circle:
68 def __init__(self, r: float) -> None:
69 self.r = r
71 def area(self) -> float:
72 return 3.14159 * self.r * self.r
73 '''
75 splitter = PythonCodeSplitter(min_effective_lines=4, max_effective_lines=6)
76 result = splitter.run(documents=[Document(content=source, meta={"file_name": "circle.py"})])
77 for chunk in result["documents"]:
78 print(chunk.meta["start_line"], chunk.meta["end_line"], chunk.meta.get("include_classes"))
79 ```
81 Pass ``strip_docstrings=True`` to move docstrings out of the chunk content and into
82 each chunk's ``meta["docstrings"]`` list. This is useful for RAG when docstrings are
83 large: stripping shrinks the stored content while the docstring text can still
84 influence retrieval via ``meta_fields_to_embed=["docstrings"]`` on the embedder.
85 """
87 def __init__(
88 self,
89 *,
90 min_effective_lines: int = 20,
91 max_effective_lines: int = 100,
92 expected_chars_per_line: int = 45,
93 oversized_factor: int = 3,
94 strip_docstrings: bool = False,
95 preserve_class_definition: bool = True,
96 secondary_split_overlap: int = 5,
97 secondary_split_length: int | None = None,
98 ) -> None:
99 """
100 Initialize the PythonCodeSplitter.
102 :param min_effective_lines: Minimum effective lines per chunk. While the running
103 chunk is below this threshold the splitter keeps merging in the next unit.
104 :param max_effective_lines: Target effective lines per chunk. Units are merged
105 greedily while doing so brings the running total closer to this target.
106 :param expected_chars_per_line: Used to convert characters into effective lines as
107 ``ceil(len(source) / expected_chars_per_line)``; long lines count as more than one.
108 :param oversized_factor: A function whose effective length exceeds
109 ``oversized_factor * max_effective_lines`` triggers the line-based secondary
110 split with overlap.
111 :param strip_docstrings: If ``True``, function/method/class docstrings are moved
112 from the chunk content into ``meta["docstrings"]`` (source order). The
113 module-level docstring is kept in place since it is itself a top-level unit.
114 :param preserve_class_definition: If ``True`` (default), chunks that contain class
115 members but not the class header are prefixed with the bare class signature
116 (decorators plus the ``class Foo(...):`` lines) in source order.
117 :param secondary_split_overlap: Line overlap for the secondary splitter; only used
118 in the oversized fallback. The primary AST split never adds overlap.
119 :param secondary_split_length: Lines per chunk for the secondary splitter.
120 Defaults to ``max_effective_lines`` when ``None``.
121 :raises ValueError: If any parameter is invalid (negative, zero where positive is
122 required, or ``min_effective_lines > max_effective_lines``).
123 """
124 if min_effective_lines < 1:
125 raise ValueError("min_effective_lines must be at least 1.")
126 if max_effective_lines < 1:
127 raise ValueError("max_effective_lines must be at least 1.")
128 if min_effective_lines > max_effective_lines:
129 raise ValueError("min_effective_lines must not be greater than max_effective_lines.")
130 if expected_chars_per_line < 1:
131 raise ValueError("expected_chars_per_line must be at least 1.")
132 if oversized_factor < 1:
133 raise ValueError("oversized_factor must be at least 1.")
134 if secondary_split_overlap < 0:
135 raise ValueError("secondary_split_overlap must be non-negative.")
136 if secondary_split_length is not None and secondary_split_length < 1:
137 raise ValueError("secondary_split_length must be at least 1.")
139 self.min_effective_lines = min_effective_lines
140 self.max_effective_lines = max_effective_lines
141 self.expected_chars_per_line = expected_chars_per_line
142 self.oversized_factor = oversized_factor
143 self.strip_docstrings = strip_docstrings
144 self.preserve_class_definition = preserve_class_definition
145 self.secondary_split_overlap = secondary_split_overlap
146 self.secondary_split_length = secondary_split_length
148 def _effective_lines(self, text: str) -> int:
149 """Return the number of *effective lines* for ``text`` (see class docstring)."""
150 if not text:
151 return 0
152 return max(1, math.ceil(len(text) / self.expected_chars_per_line))
154 def _is_oversized(self, unit: "_CodeUnit") -> bool:
155 """Return ``True`` if ``unit`` should trigger the secondary line-based split."""
156 return self._effective_lines(unit.source) > self.oversized_factor * self.max_effective_lines
158 @staticmethod
159 def _slice_lines(source_lines: list[str], start: int, end: int) -> str:
160 """Slice ``source_lines`` between the 1-indexed ``start`` and ``end`` (inclusive)."""
161 start = max(start, 1)
162 if end < start:
163 return ""
164 return "".join(source_lines[start - 1 : end])
166 @staticmethod
167 def _safe_unparse(node: ast.AST) -> str:
168 """Return ``ast.unparse(node)`` but tolerate exotic nodes by falling back to ``repr``."""
169 try:
170 return ast.unparse(node)
171 except Exception: # pragma: no cover - defensive guard
172 return repr(node)
174 def _strip_docstring(
175 self,
176 node: ast.FunctionDef | ast.AsyncFunctionDef | ast.ClassDef,
177 source_lines: list[str],
178 unit_start: int,
179 unit_end: int,
180 ) -> tuple[str, str | None]:
181 """Strip ``node``'s docstring from ``source_lines[unit_start..unit_end]`` if safely possible."""
182 docstring = ast.get_docstring(node)
183 body = node.body
184 if not docstring or not body:
185 return self._slice_lines(source_lines, unit_start, unit_end), None
187 first = body[0]
188 if not (
189 isinstance(first, ast.Expr) and isinstance(first.value, ast.Constant) and isinstance(first.value.value, str)
190 ):
191 return self._slice_lines(source_lines, unit_start, unit_end), None
193 # Skip stripping when the docstring shares a line with the def/class (would
194 # leave broken syntax) or extends past the caller's slice (e.g. class_header).
195 ds_start = first.lineno
196 ds_end = first.end_lineno or first.lineno
197 if ds_start <= node.lineno or ds_end > unit_end:
198 return self._slice_lines(source_lines, unit_start, unit_end), None
200 before = source_lines[unit_start - 1 : ds_start - 1]
201 after = source_lines[ds_end:unit_end]
202 return "".join(before + after), docstring
204 def _emit_class_units(self, cls: ast.ClassDef, source_lines: list[str], cursor: int, units: list[_CodeUnit]) -> int:
205 """Emit class header and per-method units for ``cls``; return the next cursor (1-indexed)."""
206 class_start = cls.decorator_list[0].lineno if cls.decorator_list else cls.lineno
207 class_end = cls.end_lineno or cls.lineno
208 class_name = cls.name
209 class_decorators = [self._safe_unparse(d) for d in cls.decorator_list]
211 # Methods, async methods, and nested classes become their own units so a
212 # method is never split mid-statement.
213 split_children_idx = [
214 k
215 for k, child in enumerate(cls.body)
216 if isinstance(child, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef))
217 ]
219 # Bare class signature (decorators + ``class Foo(...):`` lines) used by
220 # ``preserve_class_definition`` to prefix later chunks of the same class.
221 class_signature: str | None = None
222 if cls.body:
223 body_start = cls.body[0].lineno
224 if body_start > class_start:
225 class_signature = self._slice_lines(source_lines, class_start, body_start - 1)
227 # Whole class fits in one unit when there are no inner split points.
228 if not split_children_idx:
229 unit_slice = self._slice_lines(source_lines, cursor, class_end)
230 stripped_docstring: str | None = None
231 if self.strip_docstrings:
232 unit_slice, stripped_docstring = self._strip_docstring(cls, source_lines, cursor, class_end)
233 units.append(
234 _CodeUnit(
235 source=unit_slice,
236 start_line=class_start,
237 end_line=class_end,
238 kind="class",
239 name=class_name,
240 class_name=class_name,
241 class_signature=class_signature,
242 decorators=class_decorators,
243 docstring=stripped_docstring,
244 )
245 )
246 return class_end + 1
248 # Class header: from outer cursor up to (but excluding) the first split child.
249 first_child = cls.body[split_children_idx[0]]
250 if (
251 isinstance(first_child, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef))
252 and first_child.decorator_list
253 ):
254 first_child_start = first_child.decorator_list[0].lineno
255 else:
256 first_child_start = first_child.lineno
257 header_end = first_child_start - 1
258 header_slice = self._slice_lines(source_lines, cursor, header_end)
259 header_docstring: str | None = None
260 if self.strip_docstrings:
261 header_slice, header_docstring = self._strip_docstring(cls, source_lines, cursor, header_end)
263 units.append(
264 _CodeUnit(
265 source=header_slice,
266 start_line=class_start,
267 end_line=header_end,
268 kind="class_header",
269 name=class_name,
270 class_name=class_name,
271 class_signature=class_signature,
272 decorators=class_decorators,
273 docstring=header_docstring,
274 )
275 )
276 inner_cursor = header_end + 1
278 for idx in split_children_idx:
279 child = cls.body[idx]
280 if not isinstance(child, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)):
281 continue # narrowed above; kept for the type checker
282 child_start = child.decorator_list[0].lineno if child.decorator_list else child.lineno
283 child_end = child.end_lineno or child.lineno
284 decorators = [self._safe_unparse(d) for d in child.decorator_list]
286 unit_slice = self._slice_lines(source_lines, inner_cursor, child_end)
287 stripped_docstring = None
288 if self.strip_docstrings:
289 unit_slice, stripped_docstring = self._strip_docstring(child, source_lines, inner_cursor, child_end)
291 kind = "method" if isinstance(child, (ast.FunctionDef, ast.AsyncFunctionDef)) else "nested_class"
292 units.append(
293 _CodeUnit(
294 source=unit_slice,
295 start_line=child_start,
296 end_line=child_end,
297 kind=kind,
298 name=child.name,
299 class_name=class_name,
300 class_signature=class_signature,
301 decorators=decorators,
302 docstring=stripped_docstring,
303 )
304 )
305 inner_cursor = child_end + 1
307 # Append trailing class-body lines (comments / blanks after the last method).
308 if inner_cursor <= class_end and units:
309 trailing = self._slice_lines(source_lines, inner_cursor, class_end)
310 units[-1].source += trailing
311 units[-1].end_line = class_end
313 return class_end + 1
315 def _extract_units(self, source: str) -> list[_CodeUnit]:
316 """Parse ``source`` and produce the ordered list of syntactic split units."""
317 tree = ast.parse(source)
318 source_lines = source.splitlines(keepends=True)
319 total_lines = len(source_lines)
321 units: list[_CodeUnit] = []
322 cursor = 1
323 body = tree.body
324 node_idx = 0
325 node_count = len(body)
327 while node_idx < node_count:
328 node = body[node_idx]
330 # Module docstring (only valid as the very first statement).
331 if (
332 node_idx == 0
333 and isinstance(node, ast.Expr)
334 and isinstance(node.value, ast.Constant)
335 and isinstance(node.value.value, str)
336 ):
337 end = node.end_lineno or node.lineno
338 units.append(
339 _CodeUnit(
340 source=self._slice_lines(source_lines, cursor, end),
341 start_line=node.lineno,
342 end_line=end,
343 kind="module_docstring",
344 )
345 )
346 cursor = end + 1
347 node_idx += 1
348 continue
350 # Group consecutive imports into one unit.
351 if isinstance(node, (ast.Import, ast.ImportFrom)):
352 import_end_idx = node_idx
353 while import_end_idx < node_count and isinstance(body[import_end_idx], (ast.Import, ast.ImportFrom)):
354 import_end_idx += 1
355 last = body[import_end_idx - 1]
356 end = last.end_lineno or last.lineno
357 units.append(
358 _CodeUnit(
359 source=self._slice_lines(source_lines, cursor, end),
360 start_line=node.lineno,
361 end_line=end,
362 kind="imports",
363 )
364 )
365 cursor = end + 1
366 node_idx = import_end_idx
367 continue
369 if isinstance(node, ast.ClassDef):
370 cursor = self._emit_class_units(node, source_lines, cursor, units)
371 node_idx += 1
372 continue
374 if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
375 start = node.decorator_list[0].lineno if node.decorator_list else node.lineno
376 end = node.end_lineno or node.lineno
377 decorators = [self._safe_unparse(d) for d in node.decorator_list]
378 unit_slice = self._slice_lines(source_lines, cursor, end)
379 stripped_docstring: str | None = None
380 if self.strip_docstrings:
381 unit_slice, stripped_docstring = self._strip_docstring(node, source_lines, cursor, end)
382 units.append(
383 _CodeUnit(
384 source=unit_slice,
385 start_line=start,
386 end_line=end,
387 kind="function",
388 name=node.name,
389 decorators=decorators,
390 docstring=stripped_docstring,
391 )
392 )
393 cursor = end + 1
394 node_idx += 1
395 continue
397 # Catch-all for top-level statements (assignments, conditionals, etc.).
398 end = node.end_lineno or node.lineno
399 units.append(
400 _CodeUnit(
401 source=self._slice_lines(source_lines, cursor, end),
402 start_line=node.lineno,
403 end_line=end,
404 kind="statement",
405 )
406 )
407 cursor = end + 1
408 node_idx += 1
410 # Append trailing content (comments after the last node) so the split is loss-less.
411 if cursor <= total_lines and units:
412 trailing = self._slice_lines(source_lines, cursor, total_lines)
413 units[-1].source += trailing
414 units[-1].end_line = total_lines
415 elif cursor <= total_lines and not units:
416 units.append(
417 _CodeUnit(
418 source=self._slice_lines(source_lines, cursor, total_lines),
419 start_line=cursor,
420 end_line=total_lines,
421 kind="statement",
422 )
423 )
425 return units
427 def _merge_units(self, units: list[_CodeUnit]) -> list[list[_CodeUnit]]:
428 """Greedily merge units toward ``max_effective_lines``; oversized units become solo chunks."""
429 chunks: list[list[_CodeUnit]] = []
430 current: list[_CodeUnit] = []
431 current_lines = 0
432 target = self.max_effective_lines
434 def flush() -> None:
435 nonlocal current, current_lines
436 if current:
437 chunks.append(current)
438 current = []
439 current_lines = 0
441 for unit in units:
442 if self._is_oversized(unit):
443 flush()
444 chunks.append([unit])
445 continue
447 unit_eff = self._effective_lines(unit.source)
449 if not current:
450 current = [unit]
451 current_lines = unit_eff
452 continue
454 # Keep merging while below the minimum or while adding moves us closer to the target.
455 new_total = current_lines + unit_eff
456 if current_lines < self.min_effective_lines or abs(new_total - target) < abs(current_lines - target):
457 current.append(unit)
458 current_lines = new_total
459 else:
460 flush()
461 current = [unit]
462 current_lines = unit_eff
464 flush()
465 return chunks
467 @staticmethod
468 def _ordered_unique(items: list[str]) -> list[str]:
469 """Return the list of unique items in their first-seen order."""
470 return list(dict.fromkeys(items))
472 def _build_chunk_meta(self, chunk: list[_CodeUnit], parent_doc: Document) -> dict[str, Any]:
473 """Construct the output meta dict for a chunk of merged units."""
474 meta: dict[str, Any] = {}
475 if parent_doc.meta:
476 meta.update(deepcopy({k: v for k, v in parent_doc.meta.items() if k not in {"split_id"}}))
477 meta["source_id"] = parent_doc.id
479 # Units are emitted in source order, so chunk[0]/chunk[-1] give the extremes.
480 meta["start_line"] = chunk[0].start_line
481 meta["end_line"] = chunk[-1].end_line
482 meta["unit_kinds"] = [u.kind for u in chunk]
484 include_classes = self._ordered_unique([u.class_name for u in chunk if u.class_name])
485 if include_classes:
486 meta["include_classes"] = include_classes
488 decorators: list[str] = []
489 for u in chunk:
490 decorators.extend(u.decorators)
491 decorators = self._ordered_unique(decorators)
492 if decorators:
493 meta["decorators"] = decorators
495 if self.strip_docstrings:
496 docstrings = [u.docstring for u in chunk if u.docstring]
497 if docstrings:
498 meta["docstrings"] = docstrings
500 return meta
502 def _render_chunk_content(self, chunk: list[_CodeUnit]) -> str:
503 """Render chunk content, optionally prefixing class signatures for orphan members."""
504 body = "".join(u.source for u in chunk)
505 if not self.preserve_class_definition:
506 return body
508 classes_with_header = {u.class_name for u in chunk if u.kind in {"class", "class_header"} and u.class_name}
509 prepended: list[str] = []
510 seen: set[str] = set()
511 for u in chunk:
512 if (
513 u.class_name
514 and u.class_name not in classes_with_header
515 and u.class_name not in seen
516 and u.class_signature
517 ):
518 prepended.append(u.class_signature)
519 seen.add(u.class_name)
521 if not prepended:
522 return body
523 return "".join(prepended) + body
525 def _secondary_split(self, unit: _CodeUnit, parent_doc: Document) -> list[Document]:
526 """Apply a line-based fallback split with overlap to a single oversized unit."""
527 qualified_name = unit.name or unit.kind
528 if unit.class_name and unit.name:
529 qualified_name = f"{unit.class_name}.{unit.name}"
530 logger.warning(
531 "Oversized {kind} '{func_name}' at lines {start}-{end} ({eff} effective lines) exceeds "
532 "{factor}x max_effective_lines={max_effective_lines}; falling back to line-based secondary split "
533 "with overlap={overlap}.",
534 kind=unit.kind,
535 func_name=qualified_name,
536 start=unit.start_line,
537 end=unit.end_line,
538 eff=self._effective_lines(unit.source),
539 factor=self.oversized_factor,
540 max_effective_lines=self.max_effective_lines,
541 overlap=self.secondary_split_overlap,
542 )
544 # DocumentSplitter measures in physical lines; this approximates effective lines.
545 split_length = (
546 self.secondary_split_length if self.secondary_split_length is not None else self.max_effective_lines
547 )
548 overlap = min(self.secondary_split_overlap, max(0, split_length - 1))
550 splitter = DocumentSplitter(split_by="line", split_length=split_length, split_overlap=overlap)
551 intermediate = splitter.run(documents=[Document(content=unit.source)])["documents"]
553 base_meta = self._build_chunk_meta([unit], parent_doc)
554 results: list[Document] = []
555 for idx, piece in enumerate(intermediate):
556 meta = deepcopy(base_meta)
557 meta["secondary_split"] = True
558 meta["secondary_split_index"] = idx
559 meta["secondary_split_total"] = len(intermediate)
560 meta["qualified_name"] = qualified_name
561 results.append(Document(content=piece.content or "", meta=meta))
563 return results
565 @component.output_types(documents=list[Document])
566 def run(self, documents: list[Document]) -> dict[str, list[Document]]:
567 """
568 Split each Python source ``Document`` into syntax-aware chunks.
570 :param documents: Documents whose ``content`` is Python source code. Each
571 document's ``meta`` is propagated onto its chunks.
572 :returns: ``{"documents": [...]}`` where each chunk's meta additionally carries
573 ``source_id``, ``split_id``, ``start_line``, ``end_line``, ``unit_kinds`` and
574 - where applicable - ``include_classes``, ``decorators``, ``docstrings``,
575 ``secondary_split``.
576 :raises ValueError: If any document's content is ``None``.
577 :raises TypeError: If any document's content is not a string.
578 :raises SyntaxError: If a document's content is not valid Python.
579 """
580 for doc in documents:
581 if doc.content is None:
582 raise ValueError(
583 f"PythonCodeSplitter only works with text documents but content for document ID {doc.id} is None."
584 )
585 if not isinstance(doc.content, str):
586 raise TypeError("PythonCodeSplitter only works with text documents (str content).")
588 final_docs: list[Document] = []
589 for doc in documents:
590 assert doc.content is not None # narrowed by the loop above
591 if not doc.content.strip():
592 logger.warning("Document ID {doc_id} has empty content. Skipping this document.", doc_id=doc.id)
593 continue
595 units = self._extract_units(doc.content)
596 if not units:
597 continue
599 chunks = self._merge_units(units)
600 split_id = 0
601 for chunk in chunks:
602 if len(chunk) == 1 and self._is_oversized(chunk[0]):
603 for piece in self._secondary_split(chunk[0], doc):
604 piece.meta["split_id"] = split_id
605 split_id += 1
606 final_docs.append(piece)
607 continue
609 content = self._render_chunk_content(chunk)
610 meta = self._build_chunk_meta(chunk, doc)
611 meta["split_id"] = split_id
612 split_id += 1
613 final_docs.append(Document(content=content, meta=meta))
615 return {"documents": final_docs}