blasttrace 1.0.2__py3-none-any.whl

blasttrace/parsers/python_ast_parser.py

@@ -0,0 +1,340 @@
+"""
+BlastTrace Python AST Parser
+Extracts function/method signatures, parameter defaults, return types, class fields,
+ORM/SQL references, and cross-module call edges for the Deterministic Dependency Reachability Engine.
+"""
+
+import ast
+import os
+import re
+from typing import Dict, Any, List, Optional
+
+
+def _unparse_safe(node: Optional[ast.AST]) -> str:
+    if node is None:
+        return ""
+    try:
+        return ast.unparse(node)
+    except Exception:
+        return ""
+
+
+def _extract_function_signature(fn_node: ast.AST) -> Dict[str, Any]:
+    args_obj = getattr(fn_node, "args", None)
+    params: List[Dict[str, Any]] = []
+    if args_obj:
+        pos_args = list(getattr(args_obj, "posonlyargs", [])) + list(args_obj.args)
+        defaults = list(args_obj.defaults)
+        num_no_default = len(pos_args) - len(defaults)
+
+        for idx, arg in enumerate(pos_args):
+            if arg.arg in ("self", "cls"):
+                continue
+            has_def = idx >= num_no_default
+            def_node = defaults[idx - num_no_default] if has_def else None
+            params.append({
+                "name": arg.arg,
+                "annotation": _unparse_safe(arg.annotation),
+                "required": not has_def,
+                "default": _unparse_safe(def_node) if has_def else None,
+                "kind": "positional_or_keyword",
+            })
+
+        if args_obj.vararg:
+            params.append({
+                "name": f"*{args_obj.vararg.arg}",
+                "annotation": _unparse_safe(args_obj.vararg.annotation),
+                "required": False,
+                "default": None,
+                "kind": "vararg",
+            })
+
+        kw_defaults = list(args_obj.kw_defaults)
+        for idx, kwarg in enumerate(args_obj.kwonlyargs):
+            def_node = kw_defaults[idx] if idx < len(kw_defaults) else None
+            has_def = def_node is not None
+            params.append({
+                "name": kwarg.arg,
+                "annotation": _unparse_safe(kwarg.annotation),
+                "required": not has_def,
+                "default": _unparse_safe(def_node) if has_def else None,
+                "kind": "keyword_only",
+            })
+
+        if args_obj.kwarg:
+            params.append({
+                "name": f"**{args_obj.kwarg.arg}",
+                "annotation": _unparse_safe(args_obj.kwarg.annotation),
+                "required": False,
+                "default": None,
+                "kind": "kwarg",
+            })
+
+    return_type = _unparse_safe(getattr(fn_node, "returns", None))
+    decorators = [_unparse_safe(d) for d in getattr(fn_node, "decorator_list", [])]
+
+    sig_parts = []
+    for p in params:
+        part = p["name"]
+        if p["annotation"]:
+            part += f": {p['annotation']}"
+        if not p["required"] and p["default"] is not None:
+            part += f" = {p['default']}"
+        sig_parts.append(part)
+    sig_str = f"def {getattr(fn_node, 'name', 'unknown')}({', '.join(sig_parts)})"
+    if return_type:
+        sig_str += f" -> {return_type}"
+
+    return {
+        "name": getattr(fn_node, "name", ""),
+        "is_async": isinstance(fn_node, ast.AsyncFunctionDef),
+        "params": params,
+        "return_type": return_type,
+        "decorators": decorators,
+        "signature_string": sig_str,
+        "lineno": getattr(fn_node, "lineno", 1),
+    }
+
+
+def _inspect_function_body(fn_node: ast.AST) -> Dict[str, Any]:
+    calls: List[str] = []
+    attributes: List[str] = []
+    returned_keys: List[str] = []
+    string_literals: List[str] = []
+
+    for sub in ast.walk(fn_node):
+        if isinstance(sub, ast.Call):
+            if isinstance(sub.func, ast.Name):
+                calls.append(sub.func.id)
+            elif isinstance(sub.func, ast.Attribute):
+                owner = _unparse_safe(sub.func.value)
+                calls.append(f"{owner}.{sub.func.attr}" if owner else sub.func.attr)
+                calls.append(sub.func.attr)
+        elif isinstance(sub, ast.Attribute):
+            owner = _unparse_safe(sub.value)
+            attributes.append(f"{owner}.{sub.attr}" if owner else sub.attr)
+            attributes.append(sub.attr)
+        elif isinstance(sub, ast.Dict):
+            for k in sub.keys:
+                if isinstance(k, ast.Constant) and isinstance(k.value, str):
+                    returned_keys.append(k.value)
+        elif isinstance(sub, ast.Constant) and isinstance(sub.value, str):
+            if len(sub.value) > 2:
+                string_literals.append(sub.value)
+
+    return {
+        "calls": list(dict.fromkeys(calls)),
+        "attributes": list(dict.fromkeys(attributes)),
+        "dict_keys": list(dict.fromkeys(returned_keys)),
+        "string_literals": string_literals,
+    }
+
+
+def parse_python_content(content: str, rel_file: str) -> Dict[str, Any]:
+    functions: Dict[str, Dict[str, Any]] = {}
+    classes: Dict[str, Dict[str, Any]] = {}
+    imports: List[Dict[str, str]] = []
+    nodes: List[Dict[str, Any]] = []
+    edges: List[Dict[str, str]] = []
+
+    module_node_id = f"py:{rel_file}"
+    nodes.append({
+        "id": module_node_id,
+        "type": "python_module",
+        "name": rel_file,
+        "file": rel_file,
+        "layer": "backend",
+    })
+
+    try:
+        tree = ast.parse(content)
+    except SyntaxError as exc:
+        return {
+            "error": f"SyntaxError in {rel_file}: {exc}",
+            "functions": {},
+            "classes": {},
+            "imports": [],
+            "nodes": nodes,
+            "edges": edges,
+        }
+
+    for node in tree.body:
+        if isinstance(node, ast.Import):
+            for alias in node.names:
+                imports.append({
+                    "module": alias.name,
+                    "symbol": alias.asname or alias.name,
+                    "raw": f"import {alias.name}",
+                })
+        elif isinstance(node, ast.ImportFrom):
+            mod = node.module or ""
+            for alias in node.names:
+                imports.append({
+                    "module": mod,
+                    "symbol": alias.name,
+                    "alias": alias.asname or alias.name,
+                    "raw": f"from {mod} import {alias.name}",
+                })
+        elif isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
+            sig = _extract_function_signature(node)
+            body_info = _inspect_function_body(node)
+            sig.update(body_info)
+            sig["file"] = rel_file
+            fn_id = f"py:{rel_file}:{node.name}"
+            functions[node.name] = sig
+            nodes.append({
+                "id": fn_id,
+                "type": "python_function",
+                "name": node.name,
+                "qualified_name": f"{rel_file}:{node.name}",
+                "signature": sig["signature_string"],
+                "params": sig["params"],
+                "return_type": sig["return_type"],
+                "file": rel_file,
+                "lineno": sig["lineno"],
+                "layer": "backend",
+            })
+            edges.append({
+                "source": module_node_id,
+                "target": fn_id,
+                "relation": "defines_function",
+            })
+
+        elif isinstance(node, ast.ClassDef):
+            bases = [_unparse_safe(b) for b in node.bases]
+            cls_fields: Dict[str, Dict[str, Any]] = {}
+            cls_methods: Dict[str, Dict[str, Any]] = {}
+            tablename = None
+
+            for item in node.body:
+                if isinstance(item, ast.Assign):
+                    for target in item.targets:
+                        if isinstance(target, ast.Name):
+                            val_str = _unparse_safe(item.value)
+                            if target.id == "__tablename__":
+                                tablename = val_str.strip("\"'")
+                            elif not target.id.startswith("__"):
+                                cls_fields[target.id] = {
+                                    "name": target.id,
+                                    "annotation": "",
+                                    "default": val_str,
+                                    "lineno": item.lineno,
+                                }
+                elif isinstance(item, ast.AnnAssign) and isinstance(item.target, ast.Name):
+                    fname = item.target.id
+                    cls_fields[fname] = {
+                        "name": fname,
+                        "annotation": _unparse_safe(item.annotation),
+                        "default": _unparse_safe(item.value) if item.value else None,
+                        "lineno": item.lineno,
+                    }
+                elif isinstance(item, (ast.FunctionDef, ast.AsyncFunctionDef)):
+                    msig = _extract_function_body_and_sig(item, rel_file, node.name)
+                    cls_methods[item.name] = msig
+                    functions[f"{node.name}.{item.name}"] = msig
+
+            cls_id = f"py:{rel_file}:{node.name}"
+            classes[node.name] = {
+                "name": node.name,
+                "bases": bases,
+                "tablename": tablename,
+                "fields": cls_fields,
+                "methods": cls_methods,
+                "file": rel_file,
+                "lineno": node.lineno,
+            }
+            nodes.append({
+                "id": cls_id,
+                "type": "python_class",
+                "name": node.name,
+                "bases": bases,
+                "fields": list(cls_fields.keys()),
+                "file": rel_file,
+                "lineno": node.lineno,
+                "layer": "backend",
+            })
+            edges.append({
+                "source": module_node_id,
+                "target": cls_id,
+                "relation": "defines_class",
+            })
+
+            for fname, fmeta in cls_fields.items():
+                field_id = f"py:{rel_file}:{node.name}.{fname}"
+                nodes.append({
+                    "id": field_id,
+                    "type": "class_field",
+                    "name": f"{node.name}.{fname}",
+                    "class_name": node.name,
+                    "field": fname,
+                    "annotation": fmeta["annotation"],
+                    "file": rel_file,
+                    "layer": "backend",
+                })
+                edges.append({
+                    "source": cls_id,
+                    "target": field_id,
+                    "relation": "has_field",
+                })
+
+            for mname, mmeta in cls_methods.items():
+                method_id = f"py:{rel_file}:{node.name}.{mname}"
+                nodes.append({
+                    "id": method_id,
+                    "type": "python_method",
+                    "name": f"{node.name}.{mname}",
+                    "signature": mmeta["signature_string"],
+                    "params": mmeta["params"],
+                    "return_type": mmeta["return_type"],
+                    "file": rel_file,
+                    "lineno": mmeta["lineno"],
+                    "layer": "backend",
+                })
+                edges.append({
+                    "source": cls_id,
+                    "target": method_id,
+                    "relation": "has_method",
+                })
+
+    # Link intra-module function calls (callee -> caller impact direction)
+    for fn_name, fn_meta in functions.items():
+        caller_id = f"py:{rel_file}:{fn_name}"
+        for called in fn_meta.get("calls", []):
+            if called in functions and called != fn_name:
+                callee_id = f"py:{rel_file}:{called}"
+                edges.append({
+                    "source": callee_id,
+                    "target": caller_id,
+                    "relation": "called_by",
+                })
+
+    return {
+        "file": rel_file,
+        "functions": functions,
+        "classes": classes,
+        "imports": imports,
+        "nodes": nodes,
+        "edges": edges,
+    }
+
+
+def _extract_function_body_and_sig(fn_node: ast.AST, rel_file: str, cls_name: str) -> Dict[str, Any]:
+    sig = _extract_function_signature(fn_node)
+    body_info = _inspect_function_body(fn_node)
+    sig.update(body_info)
+    sig["file"] = rel_file
+    sig["class_name"] = cls_name
+    return sig
+
+
+def parse_python_file(file_path: str, repo_root: str = "") -> Dict[str, Any]:
+    if not os.path.exists(file_path):
+        return {"error": f"Python file not found: {file_path}", "functions": {}, "classes": {}, "imports": [], "nodes": [], "edges": []}
+    rel_file = os.path.relpath(file_path, repo_root).replace("\\", "/") if repo_root else os.path.basename(file_path)
+    try:
+        with open(file_path, "r", encoding="utf-8", errors="ignore") as f:
+            content = f.read()
+    except Exception as exc:
+        return {"error": str(exc), "functions": {}, "classes": {}, "imports": [], "nodes": [], "edges": []}
+    return parse_python_content(content, rel_file)
+