codebrain 0.1.0__py3-none-any.whl

codebrain/parser/python_parser.py

@@ -0,0 +1,658 @@
+"""Python AST visitor that extracts structural nodes and edges."""
+
+from __future__ import annotations
+
+import ast
+import json
+from pathlib import Path
+
+from codebrain.parser.base import BaseParser
+from codebrain.parser.models import ParsedEdge, ParsedFile, ParsedNode
+from codebrain.utils import content_hash
+
+
+def _get_docstring(node: ast.AST) -> str:
+    """Extract first-line docstring from a function/class/module, or return ''."""
+    ds = ast.get_docstring(node)
+    if ds:
+        first_line = ds.split("\n", 1)[0].strip()
+        return first_line[:200]
+    return ""
+
+
+def _get_decorator_names(node: ast.FunctionDef | ast.AsyncFunctionDef | ast.ClassDef) -> list[str]:
+    """Return a list of decorator name strings."""
+    names: list[str] = []
+    for dec in node.decorator_list:
+        if isinstance(dec, ast.Name):
+            names.append(dec.id)
+        elif isinstance(dec, ast.Attribute):
+            names.append(ast.unparse(dec))
+        elif isinstance(dec, ast.Call):
+            if isinstance(dec.func, ast.Name):
+                names.append(dec.func.id)
+            elif isinstance(dec.func, ast.Attribute):
+                names.append(ast.unparse(dec.func))
+            else:
+                names.append(ast.unparse(dec.func))
+        else:
+            names.append(ast.unparse(dec))
+    return names
+
+
+def _build_signature(node: ast.FunctionDef | ast.AsyncFunctionDef) -> str:
+    """Build a human-readable function signature string."""
+    params: list[str] = []
+    args = node.args
+
+    # positional-only
+    for i, arg in enumerate(args.posonlyargs):
+        p = arg.arg
+        if arg.annotation:
+            p += f": {ast.unparse(arg.annotation)}"
+        params.append(p)
+    if args.posonlyargs:
+        params.append("/")
+
+    # normal positional/keyword
+    num_defaults = len(args.defaults)
+    num_args = len(args.args)
+    for i, arg in enumerate(args.args):
+        p = arg.arg
+        if arg.annotation:
+            p += f": {ast.unparse(arg.annotation)}"
+        default_idx = i - (num_args - num_defaults)
+        if default_idx >= 0:
+            p += f" = {ast.unparse(args.defaults[default_idx])}"
+        params.append(p)
+
+    # *args
+    if args.vararg:
+        p = f"*{args.vararg.arg}"
+        if args.vararg.annotation:
+            p += f": {ast.unparse(args.vararg.annotation)}"
+        params.append(p)
+    elif args.kwonlyargs:
+        params.append("*")
+
+    # keyword-only
+    for i, arg in enumerate(args.kwonlyargs):
+        p = arg.arg
+        if arg.annotation:
+            p += f": {ast.unparse(arg.annotation)}"
+        if args.kw_defaults[i] is not None:
+            p += f" = {ast.unparse(args.kw_defaults[i])}"
+        params.append(p)
+
+    # **kwargs
+    if args.kwarg:
+        p = f"**{args.kwarg.arg}"
+        if args.kwarg.annotation:
+            p += f": {ast.unparse(args.kwarg.annotation)}"
+        params.append(p)
+
+    sig = f"({', '.join(params)})"
+    if node.returns:
+        sig += f" -> {ast.unparse(node.returns)}"
+    return sig
+
+
+def _call_target_name(node: ast.Call) -> str | None:
+    """Best-effort extraction of the callable name from a Call node."""
+    func = node.func
+    if isinstance(func, ast.Name):
+        return func.id
+    if isinstance(func, ast.Attribute):
+        return ast.unparse(func)
+    return None
+
+
+# Route decorator patterns for API endpoint extraction
+_ROUTE_METHODS = frozenset({"get", "post", "put", "delete", "patch", "head", "options", "route"})
+
+
+def _extract_route_info(decorators: list[ast.expr]) -> str | None:
+    """Extract API route path from decorators like @app.get('/path') or @router.post('/path')."""
+    for dec in decorators:
+        if isinstance(dec, ast.Call) and isinstance(dec.func, ast.Attribute):
+            method = dec.func.attr
+            if method in _ROUTE_METHODS and dec.args:
+                arg = dec.args[0]
+                if isinstance(arg, ast.Constant) and isinstance(arg.value, str):
+                    return f"{method.upper()} {arg.value}"
+    return None
+
+
+# MongoDB collection access patterns
+_MONGO_PATTERNS = (
+    # db["collection_name"] or db['collection_name']
+    r"""(?:db|database|mongo_db|self\.db|self\.database)\s*\[\s*['"](\w+)['"]\s*\]""",
+    # db.collection_name (attribute access)
+    r"""(?:db|database|mongo_db|self\.db|self\.database)\.(\w+)\b(?!\s*\()""",
+    # get_collection("name") or collection("name")
+    r"""(?:get_collection|collection)\s*\(\s*['"](\w+)['"]\s*\)""",
+)
+
+import re as _re
+_MONGO_RE = [_re.compile(p) for p in _MONGO_PATTERNS]
+
+# Significant comment patterns
+_SIGNIFICANT_COMMENT = _re.compile(
+    r"""#\s*(TODO|FIXME|HACK|NOTE|WARNING|BUG|XXX|IMPORTANT|REFACTOR)\b[:\s]*(.+)""",
+    _re.IGNORECASE,
+)
+
+
+def _extract_significant_comments(source: str, max_comments: int = 20) -> list[str]:
+    """Extract TODO, FIXME, HACK, NOTE and other significant comments."""
+    comments: list[str] = []
+    for m in _SIGNIFICANT_COMMENT.finditer(source):
+        tag = m.group(1).upper()
+        text = m.group(2).strip()[:100]
+        line = source.count("\n", 0, m.start()) + 1
+        comments.append(f"L{line} {tag}: {text}")
+        if len(comments) >= max_comments:
+            break
+    return comments
+
+# Common MongoDB method names to filter out
+_MONGO_METHODS = frozenset({
+    "find", "find_one", "insert_one", "insert_many", "update_one", "update_many",
+    "delete_one", "delete_many", "aggregate", "count_documents", "create_index",
+    "drop", "distinct", "bulk_write", "watch", "list_collection_names",
+    "get_database", "get_collection", "client", "close", "command",
+})
+
+
+def _extract_mongo_collections(source: str) -> list[str]:
+    """Extract MongoDB collection names from source code."""
+    collections: set[str] = set()
+    for pattern in _MONGO_RE:
+        for m in pattern.finditer(source):
+            name = m.group(1)
+            if name and name not in _MONGO_METHODS and not name.startswith("_"):
+                collections.add(name)
+    return sorted(collections)
+
+
+class PythonVisitor(ast.NodeVisitor):
+    """Walk a Python AST and collect structural nodes and edges."""
+
+    def __init__(self, file_path: str, module_name: str) -> None:
+        self.file_path = file_path
+        self.module_name = module_name
+        self.nodes: list[ParsedNode] = []
+        self.edges: list[ParsedEdge] = []
+        self._scope_stack: list[str] = []  # qualified name parts
+        self._all_names: set[str] | None = None  # populated if __all__ found
+        self._current_class: str | None = None  # current class node id
+
+    def _make_id(self, qualified_name: str) -> str:
+        return f"{self.file_path}::{qualified_name}"
+
+    def _resolve_call_target(self, name: str, class_node_id: str | None) -> str:
+        """Resolve ``self.X`` to ``ClassName.X`` using the enclosing class context."""
+        if class_node_id and name.startswith("self."):
+            class_qname = class_node_id.split("::", 1)[1] if "::" in class_node_id else class_node_id
+            return f"{class_qname}.{name[5:]}"
+        return name
+
+    def _qualified_name(self, name: str) -> str:
+        if self._scope_stack:
+            return ".".join(self._scope_stack) + "." + name
+        return name
+
+    def _is_exported(self, name: str) -> bool:
+        if self._all_names is not None:
+            return name in self._all_names
+        # If no __all__, top-level non-underscore names are considered exported
+        return len(self._scope_stack) == 0 and not name.startswith("_")
+
+    # ------------------------------------------------------------------
+    # Pre-scan for __all__
+    # ------------------------------------------------------------------
+    def _scan_all(self, tree: ast.Module) -> None:
+        for node in ast.iter_child_nodes(tree):
+            if isinstance(node, ast.Assign):
+                for target in node.targets:
+                    if isinstance(target, ast.Name) and target.id == "__all__":
+                        if isinstance(node.value, (ast.List, ast.Tuple)):
+                            self._all_names = set()
+                            for elt in node.value.elts:
+                                if isinstance(elt, ast.Constant) and isinstance(elt.value, str):
+                                    self._all_names.add(elt.value)
+
+    # ------------------------------------------------------------------
+    # Visitors
+    # ------------------------------------------------------------------
+    def visit_Module(self, node: ast.Module) -> None:
+        self._scan_all(node)
+        self.generic_visit(node)
+
+    def visit_FunctionDef(self, node: ast.FunctionDef) -> None:
+        self._handle_function(node)
+
+    def visit_AsyncFunctionDef(self, node: ast.AsyncFunctionDef) -> None:
+        self._handle_function(node)
+
+    def _handle_function(self, node: ast.FunctionDef | ast.AsyncFunctionDef) -> None:
+        qname = self._qualified_name(node.name)
+        node_id = self._make_id(qname)
+        node_type = "method" if self._current_class else "function"
+
+        decorators = _get_decorator_names(node)
+
+        # Check for API route decorators
+        route_info = _extract_route_info(node.decorator_list)
+        if route_info:
+            decorators.append(f"endpoint:{route_info}")
+
+        pnode = ParsedNode(
+            id=node_id,
+            name=node.name,
+            qualified_name=qname,
+            type=node_type,
+            file_path=self.file_path,
+            line_start=node.lineno,
+            line_end=node.end_lineno or node.lineno,
+            signature=_build_signature(node),
+            decorators=decorators,
+            docstring=_get_docstring(node),
+            is_exported=self._is_exported(node.name),
+        )
+        self.nodes.append(pnode)
+
+        # CONTAINS edge from parent
+        if self._current_class:
+            self.edges.append(ParsedEdge(
+                source=self._current_class,
+                target=node_id,
+                type="CONTAINS",
+                file_path=self.file_path,
+                line=node.lineno,
+            ))
+        else:
+            # file contains this function
+            file_node_id = self._make_id(self.module_name)
+            self.edges.append(ParsedEdge(
+                source=file_node_id,
+                target=node_id,
+                type="CONTAINS",
+                file_path=self.file_path,
+                line=node.lineno,
+            ))
+
+        # Walk body for calls, nested defs, etc.
+        old_class = self._current_class
+        self._current_class = None  # nested funcs are not methods
+        self._scope_stack.append(node.name)
+        self._visit_body_for_calls(node, class_node_id=old_class)
+        self.generic_visit(node)
+        self._scope_stack.pop()
+        self._current_class = old_class
+
+    def visit_ClassDef(self, node: ast.ClassDef) -> None:
+        qname = self._qualified_name(node.name)
+        node_id = self._make_id(qname)
+
+        pnode = ParsedNode(
+            id=node_id,
+            name=node.name,
+            qualified_name=qname,
+            type="class",
+            file_path=self.file_path,
+            line_start=node.lineno,
+            line_end=node.end_lineno or node.lineno,
+            decorators=_get_decorator_names(node),
+            docstring=_get_docstring(node),
+            is_exported=self._is_exported(node.name),
+        )
+        self.nodes.append(pnode)
+
+        # CONTAINS edge from file
+        file_node_id = self._make_id(self.module_name)
+        self.edges.append(ParsedEdge(
+            source=file_node_id,
+            target=node_id,
+            type="CONTAINS",
+            file_path=self.file_path,
+            line=node.lineno,
+        ))
+
+        # EXTENDS edges for base classes
+        for base in node.bases:
+            base_name = ast.unparse(base)
+            self.edges.append(ParsedEdge(
+                source=node_id,
+                target=base_name,
+                type="EXTENDS",
+                file_path=self.file_path,
+                line=node.lineno,
+            ))
+
+        # Visit children
+        old_class = self._current_class
+        self._current_class = node_id
+        self._scope_stack.append(node.name)
+        self.generic_visit(node)
+        self._scope_stack.pop()
+        self._current_class = old_class
+
+    def visit_Import(self, node: ast.Import) -> None:
+        container_id = self._make_id(
+            ".".join(self._scope_stack) if self._scope_stack else self.module_name
+        )
+        for alias in node.names:
+            self.edges.append(ParsedEdge(
+                source=container_id,
+                target=alias.name,
+                type="IMPORTS",
+                file_path=self.file_path,
+                line=node.lineno,
+            ))
+
+    def visit_ImportFrom(self, node: ast.ImportFrom) -> None:
+        container_id = self._make_id(
+            ".".join(self._scope_stack) if self._scope_stack else self.module_name
+        )
+        module = node.module or ""
+        if node.names:
+            for alias in node.names:
+                if alias.name == "*":
+                    # Star import: create edge to the module itself so dependency
+                    # tracking knows this file depends on the entire module.
+                    if module:
+                        self.edges.append(ParsedEdge(
+                            source=container_id,
+                            target=module,
+                            type="IMPORTS",
+                            file_path=self.file_path,
+                            line=node.lineno,
+                        ))
+                    continue
+                target = f"{module}.{alias.name}" if module else alias.name
+                self.edges.append(ParsedEdge(
+                    source=container_id,
+                    target=target,
+                    type="IMPORTS",
+                    file_path=self.file_path,
+                    line=node.lineno,
+                ))
+
+    def _visit_body_for_calls(self, parent: ast.AST, class_node_id: str | None = None) -> None:
+        """Walk all Call nodes inside *parent* and emit CALLS and DATAFLOW edges.
+
+        DATAFLOW edges represent data dependencies between symbols — cases where
+        the return value of one callable flows into another.  They use the same
+        ``source=consumer, target=producer`` convention as CALLS edges, except
+        for nested-call patterns where ``source=inner, target=outer`` to show
+        the inner call's result feeding into the outer call.
+
+        DATAFLOW edges are created for the following patterns:
+
+        1. **Assignment from call**: ``x = foo()`` — the enclosing function
+           (container) has a DATAFLOW edge targeting ``foo`` because the
+           container consumes the value produced by ``foo``.
+        2. **Return of call**: ``return foo()`` — same as (1); the container's
+           return value depends on ``foo``'s output.
+        3. **Nested calls**: ``outer(inner())`` — ``inner``'s output flows into
+           ``outer``, so ``source=inner, target=outer``.
+        4. **Keyword-arg calls**: ``outer(key=inner())`` — same as (3).
+        5. **Yield of call**: ``yield foo()`` — container yields ``foo``'s
+           output, so ``source=container, target=foo``.
+        6. **Await of call**: ``await foo()`` — container awaits ``foo``'s
+           output, so ``source=container, target=foo``.
+
+        Note: ``impact_of_change`` in ``graph/query.py`` currently traverses
+        only CALLS and IMPORTS edges, so DATAFLOW edges do not inflate impact
+        results.  They are available for finer-grained data-dependency analysis.
+        """
+        container_qname = ".".join(self._scope_stack) if self._scope_stack else self.module_name
+        container_id = self._make_id(container_qname)
+        for child in ast.walk(parent):
+            if isinstance(child, ast.Call):
+                name = _call_target_name(child)
+                if name:
+                    name = self._resolve_call_target(name, class_node_id)
+                    self.edges.append(ParsedEdge(
+                        source=container_id,
+                        target=name,
+                        type="CALLS",
+                        file_path=self.file_path,
+                        line=getattr(child, "lineno", 0),
+                    ))
+            # Data flow: variable = function_call()
+            if isinstance(child, ast.Assign):
+                if isinstance(child.value, ast.Call):
+                    call_name = _call_target_name(child.value)
+                    if call_name:
+                        call_name = self._resolve_call_target(call_name, class_node_id)
+                        for target in child.targets:
+                            if isinstance(target, ast.Name):
+                                self.edges.append(ParsedEdge(
+                                    source=container_id,
+                                    target=call_name,
+                                    type="DATAFLOW",
+                                    file_path=self.file_path,
+                                    line=getattr(child, "lineno", 0),
+                                ))
+            # Data flow: return function_call()
+            if isinstance(child, ast.Return) and child.value:
+                if isinstance(child.value, ast.Call):
+                    call_name = _call_target_name(child.value)
+                    if call_name:
+                        call_name = self._resolve_call_target(call_name, class_node_id)
+                        self.edges.append(ParsedEdge(
+                            source=container_id,
+                            target=call_name,
+                            type="DATAFLOW",
+                            file_path=self.file_path,
+                            line=getattr(child, "lineno", 0),
+                        ))
+            # Data flow: function_call(other_call()) — nested calls
+            if isinstance(child, ast.Call):
+                for arg in child.args:
+                    if isinstance(arg, ast.Call):
+                        inner = _call_target_name(arg)
+                        outer = _call_target_name(child)
+                        if inner and outer:
+                            inner = self._resolve_call_target(inner, class_node_id)
+                            outer = self._resolve_call_target(outer, class_node_id)
+                            self.edges.append(ParsedEdge(
+                                source=inner,
+                                target=outer,
+                                type="DATAFLOW",
+                                file_path=self.file_path,
+                                line=getattr(child, "lineno", 0),
+                            ))
+                # keyword args: func(x=other_call())
+                for kw in child.keywords:
+                    if isinstance(kw.value, ast.Call):
+                        inner = _call_target_name(kw.value)
+                        outer = _call_target_name(child)
+                        if inner and outer:
+                            inner = self._resolve_call_target(inner, class_node_id)
+                            outer = self._resolve_call_target(outer, class_node_id)
+                            self.edges.append(ParsedEdge(
+                                source=inner,
+                                target=outer,
+                                type="DATAFLOW",
+                                file_path=self.file_path,
+                                line=getattr(child, "lineno", 0),
+                            ))
+            # Data flow: yield call()
+            if isinstance(child, ast.Yield) and child.value and isinstance(child.value, ast.Call):
+                call_name = _call_target_name(child.value)
+                if call_name:
+                    call_name = self._resolve_call_target(call_name, class_node_id)
+                    self.edges.append(ParsedEdge(
+                        source=container_id,
+                        target=call_name,
+                        type="DATAFLOW",
+                        file_path=self.file_path,
+                        line=getattr(child, "lineno", 0),
+                    ))
+            # Data flow: await call()
+            if isinstance(child, ast.Await) and isinstance(child.value, ast.Call):
+                call_name = _call_target_name(child.value)
+                if call_name:
+                    call_name = self._resolve_call_target(call_name, class_node_id)
+                    self.edges.append(ParsedEdge(
+                        source=container_id,
+                        target=call_name,
+                        type="DATAFLOW",
+                        file_path=self.file_path,
+                        line=getattr(child, "lineno", 0),
+                    ))
+
+    def visit_Assign(self, node: ast.Assign) -> None:
+        # Only capture module-level variable assignments
+        if self._scope_stack:
+            self.generic_visit(node)
+            return
+        for target in node.targets:
+            if isinstance(target, ast.Name):
+                if target.id == "__all__":
+                    continue
+                qname = self._qualified_name(target.id)
+                node_id = self._make_id(qname)
+                pnode = ParsedNode(
+                    id=node_id,
+                    name=target.id,
+                    qualified_name=qname,
+                    type="variable",
+                    file_path=self.file_path,
+                    line_start=node.lineno,
+                    line_end=node.end_lineno or node.lineno,
+                    is_exported=self._is_exported(target.id),
+                )
+                self.nodes.append(pnode)
+                file_node_id = self._make_id(self.module_name)
+                self.edges.append(ParsedEdge(
+                    source=file_node_id,
+                    target=node_id,
+                    type="CONTAINS",
+                    file_path=self.file_path,
+                    line=node.lineno,
+                ))
+        self.generic_visit(node)
+
+    def visit_AnnAssign(self, node: ast.AnnAssign) -> None:
+        if self._scope_stack:
+            self.generic_visit(node)
+            return
+        if isinstance(node.target, ast.Name):
+            name = node.target.id
+            qname = self._qualified_name(name)
+            node_id = self._make_id(qname)
+            pnode = ParsedNode(
+                id=node_id,
+                name=name,
+                qualified_name=qname,
+                type="variable",
+                file_path=self.file_path,
+                line_start=node.lineno,
+                line_end=node.end_lineno or node.lineno,
+                is_exported=self._is_exported(name),
+            )
+            self.nodes.append(pnode)
+            file_node_id = self._make_id(self.module_name)
+            self.edges.append(ParsedEdge(
+                source=file_node_id,
+                target=node_id,
+                type="CONTAINS",
+                file_path=self.file_path,
+                line=node.lineno,
+            ))
+        self.generic_visit(node)
+
+
+def parse_file(path: Path, repo_root: Path) -> ParsedFile:
+    """Parse a Python file and return a ParsedFile with nodes and edges."""
+    rel_path = path.relative_to(repo_root).as_posix()
+    source = path.read_bytes()
+    hash_val = content_hash(source)
+
+    try:
+        tree = ast.parse(source, filename=str(path))
+    except (SyntaxError, ValueError):
+        # ValueError: null bytes in source. SyntaxError: bad syntax.
+        # Retry with explicit UTF-8 decoding and null-byte stripping.
+        try:
+            text = source.decode("utf-8", errors="replace").replace("\x00", "")
+            tree = ast.parse(text, filename=str(path))
+        except (SyntaxError, ValueError):
+            return ParsedFile(path=rel_path, content_hash=hash_val, line_count=source.count(b"\n") + 1)
+
+    line_count = source.count(b"\n") + 1
+
+    # Derive a module-like name from the relative path
+    module_name = rel_path.replace("/", ".").removesuffix(".py")
+    if module_name.endswith(".__init__"):
+        module_name = module_name.removesuffix(".__init__")
+
+    # Create a file-level node
+    try:
+        source_text_for_comments = source.decode("utf-8", errors="replace")
+    except Exception:
+        source_text_for_comments = ""
+    significant_comments = _extract_significant_comments(source_text_for_comments)
+    file_docstring = _get_docstring(tree)
+    if significant_comments:
+        comment_block = " | ".join(significant_comments)
+        if file_docstring:
+            file_docstring += f" [Comments: {comment_block}]"
+        else:
+            file_docstring = f"[Comments: {comment_block}]"
+
+    file_node = ParsedNode(
+        id=f"{rel_path}::{module_name}",
+        name=module_name.rsplit(".", 1)[-1],
+        qualified_name=module_name,
+        type="file",
+        file_path=rel_path,
+        line_start=1,
+        line_end=line_count,
+        docstring=file_docstring,
+        is_exported=True,
+    )
+
+    visitor = PythonVisitor(file_path=rel_path, module_name=module_name)
+    visitor.visit(tree)
+
+    nodes = [file_node] + visitor.nodes
+    edges = visitor.edges
+
+    # Extract MongoDB collection references
+    try:
+        source_text = source.decode("utf-8", errors="replace")
+    except Exception:
+        source_text = ""
+    collections = _extract_mongo_collections(source_text)
+    file_node_id = f"{rel_path}::{module_name}"
+    for coll_name in collections:
+        edges.append(ParsedEdge(
+            source=file_node_id,
+            target=f"mongodb:{coll_name}",
+            type="DATAFLOW",
+            file_path=rel_path,
+            line=0,
+        ))
+
+    return ParsedFile(
+        path=rel_path,
+        content_hash=hash_val,
+        nodes=nodes,
+        edges=edges,
+        line_count=line_count,
+    )
+
+
+class PythonParser(BaseParser):
+    """BaseParser wrapper around the Python AST parser."""
+
+    def extensions(self) -> frozenset[str]:
+        return frozenset({".py"})
+
+    def parse(self, path: Path, repo_root: Path) -> ParsedFile:
+        return parse_file(path, repo_root)