Coverage for src / mcp_server_langgraph / builder / importer / ast_parser.py: 67%
92 statements
« prev ^ index » next coverage.py v7.12.0, created at 2025-12-03 00:43 +0000
1"""
2AST Parser for Python Code Import
4Parses Python source code into Abstract Syntax Tree for analysis.
6Key capabilities:
7- Parse Python code safely
8- Extract function definitions
9- Find class definitions
10- Identify LangGraph API calls
11- Extract configuration values
13Example:
14 from mcp_server_langgraph.builder.importer import PythonCodeParser
16 parser = PythonCodeParser()
17 ast_tree = parser.parse_file("agent.py")
19 # Extract all function calls
20 calls = parser.find_function_calls(ast_tree, "add_node")
21"""
23import ast
24from typing import Any
27class PythonCodeParser:
28 """
29 Parser for Python source code using AST.
31 Safely analyzes Python code to extract structure without execution.
32 """
34 def __init__(self) -> None:
35 """Initialize parser."""
36 self.ast_tree: ast.Module | None = None
38 def parse_code(self, code: str) -> ast.Module:
39 """
40 Parse Python code string into AST.
42 Args:
43 code: Python source code
45 Returns:
46 AST Module
48 Raises:
49 SyntaxError: If code is invalid Python
51 Example:
52 >>> parser = PythonCodeParser()
53 >>> tree = parser.parse_code("def foo(): pass")
54 >>> isinstance(tree, ast.Module)
55 True
56 """
57 self.ast_tree = ast.parse(code)
58 return self.ast_tree
60 def parse_file(self, file_path: str) -> ast.Module:
61 """
62 Parse Python file into AST.
64 Args:
65 file_path: Path to Python file
67 Returns:
68 AST Module
70 Example:
71 >>> parser = PythonCodeParser()
72 >>> tree = parser.parse_file("agent.py")
73 """
74 with open(file_path) as f:
75 code = f.read()
77 return self.parse_code(code)
79 def find_function_calls(self, tree: ast.Module | None = None, function_name: str | None = None) -> list[dict[str, Any]]:
80 """
81 Find all function calls in AST.
83 Args:
84 tree: AST tree (uses self.ast_tree if None)
85 function_name: Optional filter for specific function name
87 Returns:
88 List of function call information
90 Example:
91 >>> calls = parser.find_function_calls(tree, "add_node")
92 >>> len(calls)
93 3
94 """
95 tree = tree or self.ast_tree
96 if not tree: 96 ↛ 97line 96 didn't jump to line 97 because the condition on line 96 was never true
97 return []
99 calls = []
101 for node in ast.walk(tree):
102 if isinstance(node, ast.Call):
103 # Get function name
104 if isinstance(node.func, ast.Name):
105 func_name = node.func.id
106 elif isinstance(node.func, ast.Attribute): 106 ↛ 109line 106 didn't jump to line 109 because the condition on line 106 was always true
107 func_name = node.func.attr
108 else:
109 continue
111 # Filter if function_name specified
112 if function_name and func_name != function_name:
113 continue
115 # Extract arguments
116 args = []
117 for arg in node.args:
118 args.append(self._extract_value(arg))
120 # Extract keyword arguments
121 kwargs = {}
122 for keyword in node.keywords: 122 ↛ 123line 122 didn't jump to line 123 because the loop on line 122 never started
123 if keyword.arg:
124 kwargs[keyword.arg] = self._extract_value(keyword.value)
126 calls.append({"function": func_name, "args": args, "kwargs": kwargs, "lineno": node.lineno})
128 return calls
130 def find_class_definitions(self, tree: ast.Module | None = None) -> list[dict[str, Any]]:
131 """
132 Find all class definitions.
134 Args:
135 tree: AST tree
137 Returns:
138 List of class information
140 Example:
141 >>> classes = parser.find_class_definitions(tree)
142 >>> [c["name"] for c in classes]
143 ['AgentState', 'MyClass']
144 """
145 tree = tree or self.ast_tree
146 if not tree: 146 ↛ 147line 146 didn't jump to line 147 because the condition on line 146 was never true
147 return []
149 classes = []
151 for node in ast.walk(tree):
152 if isinstance(node, ast.ClassDef):
153 # Extract base classes
154 bases = []
155 for base in node.bases:
156 if isinstance(base, ast.Name): 156 ↛ 155line 156 didn't jump to line 155 because the condition on line 156 was always true
157 bases.append(base.id)
159 # Extract class docstring
160 docstring = ast.get_docstring(node)
162 classes.append({"name": node.name, "bases": bases, "docstring": docstring, "lineno": node.lineno})
164 return classes
166 def find_variable_assignments(
167 self, tree: ast.Module | None = None, variable_name: str | None = None
168 ) -> list[dict[str, Any]]:
169 """
170 Find variable assignments.
172 Args:
173 tree: AST tree
174 variable_name: Optional filter for specific variable
176 Returns:
177 List of assignments
179 Example:
180 >>> assignments = parser.find_variable_assignments(tree, "graph")
181 >>> assignments[0]["value"]
182 'StateGraph(...)'
183 """
184 tree = tree or self.ast_tree
185 if not tree: 185 ↛ 186line 185 didn't jump to line 186 because the condition on line 185 was never true
186 return []
188 assignments = []
190 for node in ast.walk(tree):
191 if isinstance(node, ast.Assign):
192 for target in node.targets:
193 if isinstance(target, ast.Name): 193 ↛ 192line 193 didn't jump to line 192 because the condition on line 193 was always true
194 var_name = target.id
196 # Filter if variable_name specified
197 if variable_name and var_name != variable_name: 197 ↛ 198line 197 didn't jump to line 198 because the condition on line 197 was never true
198 continue
200 value = self._extract_value(node.value)
202 assignments.append({"variable": var_name, "value": value, "lineno": node.lineno})
204 return assignments
206 def _extract_value(self, node: ast.AST) -> Any:
207 """
208 Extract value from AST node.
210 Handles constants, names, lists, dicts, calls, etc.
212 Args:
213 node: AST node
215 Returns:
216 Extracted value (str representation for complex types)
217 """
218 if isinstance(node, ast.Constant):
219 return node.value
220 elif isinstance(node, ast.Name):
221 return node.id
222 elif isinstance(node, ast.List): 222 ↛ 223line 222 didn't jump to line 223 because the condition on line 222 was never true
223 return [self._extract_value(elt) for elt in node.elts]
224 elif isinstance(node, ast.Dict): 224 ↛ 225line 224 didn't jump to line 225 because the condition on line 224 was never true
225 return {
226 self._extract_value(k): self._extract_value(v)
227 for k, v in zip(node.keys, node.values, strict=False)
228 if k is not None
229 }
230 elif isinstance(node, ast.Call):
231 # Return string representation of call
232 if isinstance(node.func, ast.Name): 232 ↛ 234line 232 didn't jump to line 234 because the condition on line 232 was always true
233 return f"{node.func.id}(...)"
234 elif isinstance(node.func, ast.Attribute):
235 return f"{node.func.attr}(...)"
236 return "call(...)"
237 else:
238 # Return string representation for complex types
239 return ast.unparse(node) if hasattr(ast, "unparse") else "..."
241 def get_imports(self, tree: ast.Module | None = None) -> list[str]:
242 """
243 Extract all imports from code.
245 Args:
246 tree: AST tree
248 Returns:
249 List of imported modules
251 Example:
252 >>> imports = parser.get_imports(tree)
253 >>> "langgraph" in imports
254 True
255 """
256 tree = tree or self.ast_tree
257 if not tree:
258 return []
260 imports = []
262 for node in ast.walk(tree):
263 if isinstance(node, ast.Import):
264 for alias in node.names:
265 imports.append(alias.name)
266 elif isinstance(node, ast.ImportFrom) and node.module:
267 imports.append(node.module)
269 return imports
272# ==============================================================================
273# Example Usage
274# ==============================================================================
276if __name__ == "__main__":
277 # Example Python code
278 sample_code = """
279from langgraph.graph import StateGraph
281class MyState(TypedDict):
282 query: str
283 result: str
285graph = StateGraph(MyState)
286graph.add_node("search", search_function)
287graph.add_node("summarize", summarize_function)
288graph.add_edge("search", "summarize")
289graph.set_entry_point("search")
291app = graph.compile()
292"""
294 # Parse code
295 parser = PythonCodeParser()
296 tree = parser.parse_code(sample_code)
298 print("=" * 80)
299 print("AST PARSER - TEST RUN")
300 print("=" * 80)
302 # Find imports
303 imports = parser.get_imports(tree)
304 print(f"\nImports found: {imports}")
306 # Find add_node calls
307 add_node_calls = parser.find_function_calls(tree, "add_node")
308 print(f"\nadd_node calls: {len(add_node_calls)}")
309 for call in add_node_calls:
310 print(f" - Line {call['lineno']}: add_node({call['args']})")
312 # Find add_edge calls
313 add_edge_calls = parser.find_function_calls(tree, "add_edge")
314 print(f"\nadd_edge calls: {len(add_edge_calls)}")
315 for call in add_edge_calls:
316 print(f" - Line {call['lineno']}: add_edge({call['args']})")
318 # Find class definitions
319 classes = parser.find_class_definitions(tree)
320 print(f"\nClasses found: {len(classes)}")
321 for cls in classes:
322 print(f" - {cls['name']} (bases: {cls['bases']})")
324 print("\n" + "=" * 80)