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