Coverage for src/qdrant_loader/core/chunking/strategy/code/parser/python_ast.py: 89%
38 statements
« prev ^ index » next coverage.py v7.10.6, created at 2025-09-08 06:05 +0000
« prev ^ index » next coverage.py v7.10.6, created at 2025-09-08 06:05 +0000
1from __future__ import annotations
3import ast
5from qdrant_loader.core.chunking.strategy.code.parser.common import (
6 CodeElement,
7 CodeElementType,
8)
11def parse_python_ast(
12 content: str,
13 *,
14 max_elements_to_process: int,
15) -> list[CodeElement]:
16 try:
17 tree = ast.parse(content)
18 except Exception:
19 return []
21 elements: list[CodeElement] = []
23 class Visitor(ast.NodeVisitor):
24 def __init__(self):
25 self.level = 0
27 def generic_visit(self, node):
28 if len(elements) >= max_elements_to_process:
29 return
30 node_type = type(node).__name__
31 if isinstance(node, ast.FunctionDef | ast.AsyncFunctionDef):
32 elem_type = CodeElementType.FUNCTION
33 elif isinstance(node, ast.ClassDef):
34 elem_type = CodeElementType.CLASS
35 else:
36 elem_type = CodeElementType.MODULE
38 try:
39 start_line = getattr(node, "lineno", 1)
40 end_line = getattr(node, "end_lineno", start_line)
41 except Exception:
42 start_line = 1
43 end_line = start_line
45 snippet_lines = content.split("\n")[start_line - 1 : end_line]
46 snippet = "\n".join(snippet_lines)
47 if not snippet.strip():
48 return
50 element = CodeElement(
51 name=getattr(node, "name", node_type),
52 element_type=elem_type,
53 content=snippet,
54 start_line=start_line,
55 end_line=end_line,
56 level=self.level,
57 )
58 elements.append(element)
59 self.level += 1
60 super().generic_visit(node)
61 self.level = max(0, self.level - 1)
63 Visitor().visit(tree)
64 return elements