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

1from __future__ import annotations 

2 

3import ast 

4 

5from qdrant_loader.core.chunking.strategy.code.parser.common import ( 

6 CodeElement, 

7 CodeElementType, 

8) 

9 

10 

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 [] 

20 

21 elements: list[CodeElement] = [] 

22 

23 class Visitor(ast.NodeVisitor): 

24 def __init__(self): 

25 self.level = 0 

26 

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 

37 

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 

44 

45 snippet_lines = content.split("\n")[start_line - 1 : end_line] 

46 snippet = "\n".join(snippet_lines) 

47 if not snippet.strip(): 

48 return 

49 

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) 

62 

63 Visitor().visit(tree) 

64 return elements