codegraph-brain 0.6.0__py3-none-any.whl

cgis/extractors/_python_functions.py

@@ -0,0 +1,312 @@
+"""Function/method-definition handling for the Python extractor.
+
+Emits FUNCTION/METHOD nodes plus their CALLS, DECLARES/CONTAINS, DEPENDS_ON
+edges, and collects local type information (assignments and typed parameters)
+used downstream for call resolution.
+"""
+
+from collections.abc import Callable
+from typing import Any
+
+from tree_sitter import Node as BaseNode
+
+from cgis.core.models import Edge, EdgeType, Node, NodeType
+from cgis.extractors._python_ast import (
+    PYTHON_LANG,
+    extract_node_name,
+    get_id,
+    get_identifier,
+    is_method,
+)
+from cgis.extractors._python_types import TypeResolver
+
+
+class FunctionHandler:
+    """Extracts function/method nodes, call edges and local type metadata."""
+
+    _DI_CALL_NAMES: frozenset[str] = frozenset({"Depends", "Security"})
+
+    def __init__(
+        self,
+        pick_source_root: Callable[[str], str | None],
+        type_resolver: TypeResolver,
+    ) -> None:
+        """Store the source-root picker and the shared type resolver."""
+        self._pick_source_root = pick_source_root
+        self._types = type_resolver
+
+    def process_function_node(
+        self,
+        node: BaseNode,
+        code_bytes: bytes,
+        file_path: str,
+        nodes: list[Node],
+        edges: list[Edge],
+        module_fqn: str,
+        decorators: list[str] | None = None,
+    ) -> Node:
+        """Process function or method definition node."""
+        child = node.child_by_field_name("name")
+        node_id = get_id(node, code_bytes, file_path, self._pick_source_root(file_path))
+        node_name = extract_node_name(child, code_bytes)
+        node_type = NodeType.METHOD if is_method(node) else NodeType.FUNCTION
+
+        metadata: dict[str, Any] = {}
+        if decorators:
+            metadata["decorators"] = decorators
+        if decorators and any(
+            d == "abstractmethod" or d.endswith(".abstractmethod") for d in decorators
+        ):
+            metadata["is_abstract"] = True
+
+        func_node = Node(
+            id=node_id,
+            type=node_type,
+            name=node_name,
+            file_path=file_path,
+            start_line=node.start_point.row + 1,
+            end_line=node.end_point.row + 1,
+            language=PYTHON_LANG,
+            metadata=metadata,
+        )
+        nodes.append(func_node)
+
+        edges.extend(
+            Edge(
+                id=f"{node_id}:decorator:{i}:{deco_name}",
+                type=EdgeType.CALLS,
+                source=node_id,
+                target=f"raw_call:{deco_name}",
+                confidence=0.5,
+                file_path=file_path,
+            )
+            for i, deco_name in enumerate(decorators or [])
+        )
+
+        parts = node_id.rsplit(".", maxsplit=1)
+        parent_fqn = parts[0] if len(parts) > 1 else module_fqn
+        edge_type = EdgeType.DECLARES if node_type == NodeType.METHOD else EdgeType.CONTAINS
+        edges.append(
+            Edge(
+                id=f"{parent_fqn}:structural:{node_id}",
+                type=edge_type,
+                source=parent_fqn,
+                target=node_id,
+                confidence=1.0,
+                file_path=file_path,
+            )
+        )
+        return func_node
+
+    def process_call_node(
+        self, node: BaseNode, code_bytes: bytes, file_path: str, source_id: str, edges: list[Edge]
+    ) -> None:
+        """
+        Finds call expressions node.
+        """
+        child = node.child_by_field_name("function")
+        edge_id = f"{file_path}:edge_{node.start_byte}_{node.end_byte}"
+        if child:
+            call_name = get_identifier(child, code_bytes)
+            if call_name == "unknown":
+                return
+            target_id = f"raw_call:{call_name}"
+
+            edges.append(
+                Edge(
+                    id=edge_id,
+                    type=EdgeType.CALLS,
+                    source=source_id,
+                    target=target_id,
+                    confidence=0.5,
+                    context=f"Call to {call_name}",
+                    file_path=file_path,
+                    line_number=node.start_point.row + 1,
+                )
+            )
+            if call_name in self._DI_CALL_NAMES:
+                provider = self._di_provider_name(node, code_bytes)
+                if provider:
+                    edges.append(
+                        Edge(
+                            id=f"{file_path}:dep_{node.start_byte}_{node.end_byte}",
+                            type=EdgeType.DEPENDS_ON,
+                            source=source_id,
+                            target=f"raw_call:{provider}",
+                            confidence=0.5,
+                            context=f"DI dependency on {provider}",
+                            file_path=file_path,
+                            line_number=node.start_point.row + 1,
+                        )
+                    )
+
+    def _di_provider_name(self, call_node: BaseNode, code_bytes: bytes) -> str | None:
+        """Return the first positional argument's identifier/dotted name, or None.
+
+        None for argless calls, keyword-only calls, and non-name arguments
+        (lambdas, calls, subscripts) — those emit no DEPENDS_ON edge (spec §3.2a/b).
+
+        Note: ``child_by_field_name("arguments")`` returns a truthy
+        ``argument_list`` node even for ``Depends()`` (argless), so the
+        ``if not args`` guard here only covers a *missing* arguments field
+        (should not occur in practice).  The argless case — where
+        ``args.named_children`` is empty — is handled by the loop falling
+        through to the final ``return None``.
+        """
+        args = call_node.child_by_field_name("arguments")
+        if not args:
+            return None
+        for child in args.named_children:
+            if child.type == "keyword_argument":
+                continue
+            if child.type in ("identifier", "attribute"):
+                name = get_identifier(child, code_bytes)
+                return name if name != "unknown" else None
+            return None
+        return None
+
+    def _find_di_calls(self, node: BaseNode, code_bytes: bytes) -> list[BaseNode]:
+        """Return all call nodes in the subtree whose callee is a DI name (spec §3.2a)."""
+        found: list[BaseNode] = []
+        stack = [node]
+        while stack:
+            curr = stack.pop()
+            if curr.type == "call":
+                fn = curr.child_by_field_name("function")
+                if fn and get_identifier(fn, code_bytes) in self._DI_CALL_NAMES:
+                    found.append(curr)
+            stack.extend(curr.children)
+        return found
+
+    def process_module_assignment(
+        self,
+        node: BaseNode,
+        code_bytes: bytes,
+        file_path: str,
+        nodes: list[Node],
+        edges: list[Edge],
+        module_fqn: str,
+    ) -> None:
+        """Emit a VARIABLE alias node + DEPENDS_ON edges for module-level DI assignments.
+
+        Only fires for `Name = <RHS containing Depends/Security>` with a plain
+        identifier LHS at true module level (class bodies excluded by the
+        caller). Plain constants never reach the node list (spec §3.2a).
+        """
+        left = node.child_by_field_name("left")
+        right = node.child_by_field_name("right")
+        if not left or not right or left.type != "identifier":
+            return
+        di_calls = self._find_di_calls(right, code_bytes)
+        if not di_calls:
+            return
+        name = get_identifier(left, code_bytes)
+        if name == "unknown":
+            return
+        alias_id = f"{module_fqn}.{name}" if module_fqn else name
+        nodes.append(
+            Node(
+                id=alias_id,
+                type=NodeType.VARIABLE,
+                name=name,
+                file_path=file_path,
+                start_line=node.start_point.row + 1,
+                end_line=node.end_point.row + 1,
+                language=PYTHON_LANG,
+            )
+        )
+        for call in di_calls:
+            provider = self._di_provider_name(call, code_bytes)
+            if not provider:
+                continue
+            edges.append(
+                Edge(
+                    id=f"{file_path}:dep_{call.start_byte}_{call.end_byte}",
+                    type=EdgeType.DEPENDS_ON,
+                    source=alias_id,
+                    target=f"raw_call:{provider}",
+                    confidence=0.5,
+                    context=f"DI alias for {provider}",
+                    file_path=file_path,
+                    line_number=node.start_point.row + 1,
+                )
+            )
+
+    def collect_assignment_type(
+        self,
+        node: BaseNode,
+        code_bytes: bytes,
+        import_map: dict[str, str] | None,
+        func_node: Node,
+        acc: dict[str, dict[str, str]],
+    ) -> None:
+        """Populate acc with var→FQN for `var = ClassName(...)` assignments."""
+        left_node = node.child_by_field_name("left")
+        right_node = node.child_by_field_name("right")
+        if not left_node or not right_node or right_node.type != "call":
+            return
+        var_name = get_identifier(left_node, code_bytes)
+        if var_name == "unknown":
+            return
+        func_call_node = right_node.child_by_field_name("function")
+        if not func_call_node:
+            return
+        class_name = get_identifier(func_call_node, code_bytes)
+        if class_name == "unknown":
+            return
+        if "." in class_name:
+            # Module-qualified constructor (e.g. models.Store()): keep only if prefix
+            # is a known import alias. Otherwise it's a method-call result — skip.
+            module_part, _, _ = class_name.partition(".")
+            if not import_map or module_part not in import_map:
+                return
+        acc.setdefault(func_node.id, {})[var_name] = self._types.resolve_type_fqn(
+            class_name, import_map, func_node.file_path
+        )
+
+    def collect_param_type(
+        self,
+        node: BaseNode,
+        code_bytes: bytes,
+        import_map: dict[str, str] | None,
+        func_node: Node,
+        acc: dict[str, dict[str, str]],
+        edges: list[Edge],
+    ) -> None:
+        """Populate acc with param→FQN for typed parameter annotations.
+
+        Also emits a speculative `raw_dep:` DEPENDS_ON candidate per typed
+        parameter; the resolver keeps it only when it resolves to a DI alias
+        (VARIABLE node) and drops it otherwise (spec §3.2c).
+        """
+        if not node.named_children:
+            return
+        name_node = node.named_children[0]
+        type_node = node.child_by_field_name("type")
+        if not type_node:
+            return
+        var_name = get_identifier(name_node, code_bytes)
+        if var_name == "unknown":
+            return
+        # Slice raw bytes to capture union/generic types like `A | None` or `list[X]`
+        raw_type = (
+            code_bytes[type_node.start_byte : type_node.end_byte].decode("utf-8").strip("\"'")
+        )
+        clean_type = self._types.clean_python_type_string(raw_type)
+        if not clean_type:
+            return
+        acc.setdefault(func_node.id, {})[var_name] = self._types.resolve_type_fqn(
+            clean_type, import_map, func_node.file_path
+        )
+        edges.append(
+            Edge(
+                id=f"{func_node.file_path}:rawdep_{node.start_byte}_{node.end_byte}",
+                type=EdgeType.DEPENDS_ON,
+                source=func_node.id,
+                target=f"raw_dep:{clean_type}",
+                confidence=0.1,
+                context=f"Annotation candidate {clean_type}",
+                file_path=func_node.file_path,
+                line_number=node.start_point.row + 1,
+            )
+        )

cgis/extractors/_python_imports.py

@@ -0,0 +1,188 @@
+"""Import-statement handling for the Python extractor.
+
+Walks ``import`` / ``import_from`` AST nodes, builds the per-file ``import_map``
+(local name -> target FQN) and emits IMPORTS / IMPORTS_SYMBOL edges.
+"""
+
+from tree_sitter import Node as BaseNode
+
+from cgis.core.models import Edge, EdgeType
+from cgis.extractors._python_ast import resolve_relative_module
+
+
+class ImportHandler:
+    """Extracts import edges and populates the import map from import AST nodes."""
+
+    def handle(
+        self,
+        node: BaseNode,
+        code_bytes: bytes,
+        file_path: str,
+        import_map: dict[str, str] | None,
+        module_fqn: str | None,
+        edges: list[Edge],
+    ) -> None:
+        """Dispatch import or import-from AST nodes to the appropriate handler."""
+        if import_map is None or module_fqn is None:
+            return
+        if node.type == "import_statement":
+            self._process_import_statement(
+                node, code_bytes, module_fqn, file_path, import_map, edges
+            )
+        else:
+            self._process_import_from_statement(
+                node, code_bytes, module_fqn, file_path, import_map, edges
+            )
+
+    def _process_import_statement(
+        self,
+        node: BaseNode,
+        code_bytes: bytes,
+        module_fqn: str,
+        file_path: str,
+        import_map: dict[str, str],
+        edges: list[Edge],
+    ) -> None:
+        """Handle `import X` and `import X as Y` statements."""
+        for child in node.children:
+            if child.type == "dotted_name":
+                module_str = code_bytes[child.start_byte : child.end_byte].decode("utf-8")
+                local_name = module_str.split(".")[0]
+                import_map[local_name] = local_name
+                edges.append(
+                    Edge(
+                        id=f"{module_fqn}:imports:{module_str}",
+                        type=EdgeType.IMPORTS,
+                        source=module_fqn,
+                        target=module_str,
+                        confidence=1.0,
+                        file_path=file_path,
+                    )
+                )
+            elif child.type == "aliased_import":
+                name_node = child.child_by_field_name("name")
+                alias_node = child.child_by_field_name("alias")
+                if name_node and alias_node:
+                    module_str = code_bytes[name_node.start_byte : name_node.end_byte].decode(
+                        "utf-8"
+                    )
+                    alias = code_bytes[alias_node.start_byte : alias_node.end_byte].decode("utf-8")
+                    import_map[alias] = module_str
+                    edges.append(
+                        Edge(
+                            id=f"{module_fqn}:imports:{module_str}",
+                            type=EdgeType.IMPORTS,
+                            source=module_fqn,
+                            target=module_str,
+                            confidence=1.0,
+                            file_path=file_path,
+                        )
+                    )
+
+    def _parse_relative_import(self, node: BaseNode, code_bytes: bytes) -> tuple[int, str]:
+        """Return (leading_dots, raw_module_str) from a relative_import node."""
+        leading_dots = 0
+        raw_module_str = ""
+        for sub in node.children:
+            if sub.type == "import_prefix":
+                leading_dots = sub.end_byte - sub.start_byte
+            elif sub.type == "dotted_name":
+                raw_module_str = code_bytes[sub.start_byte : sub.end_byte].decode("utf-8")
+        return leading_dots, raw_module_str
+
+    def _collect_imported_symbols(
+        self, children: list[BaseNode], code_bytes: bytes
+    ) -> list[tuple[str, str]]:
+        """Collect (local_name, qualified_symbol) pairs from a list of sibling nodes.
+
+        Flattens parenthesized imports (`import_list` node) transparently.
+        """
+        items: list[BaseNode] = []
+        for child in children:
+            if child.type == "import_list":
+                items.extend(child.children)
+            else:
+                items.append(child)
+        symbols: list[tuple[str, str]] = []
+        for child in items:
+            if child.type in ("dotted_name", "identifier"):
+                sym = code_bytes[child.start_byte : child.end_byte].decode("utf-8")
+                symbols.append((sym.split(".")[-1], sym))
+            elif child.type == "aliased_import":
+                name_node = child.child_by_field_name("name")
+                alias_node = child.child_by_field_name("alias")
+                if name_node and alias_node:
+                    sym = code_bytes[name_node.start_byte : name_node.end_byte].decode("utf-8")
+                    alias = code_bytes[alias_node.start_byte : alias_node.end_byte].decode("utf-8")
+                    symbols.append((alias, sym))
+            # wildcard_import, punctuation → skip
+        return symbols
+
+    def _process_import_from_statement(
+        self,
+        node: BaseNode,
+        code_bytes: bytes,
+        module_fqn: str,
+        file_path: str,
+        import_map: dict[str, str],
+        edges: list[Edge],
+    ) -> None:
+        """Handle `from X import Y [as Z]` and relative imports."""
+        leading_dots = 0
+        raw_module_str = ""
+        past_import_kw = False
+        import_symbol_node: BaseNode | None = None
+
+        for child in node.children:
+            if child.type == "relative_import":
+                leading_dots, raw_module_str = self._parse_relative_import(child, code_bytes)
+            elif child.type == "dotted_name" and not past_import_kw:
+                raw_module_str = code_bytes[child.start_byte : child.end_byte].decode("utf-8")
+            elif child.type == "import":
+                past_import_kw = True
+                import_symbol_node = child
+            # imported symbols are siblings after the 'import' keyword — collected below
+
+        # Collect all sibling nodes that come after the 'import' keyword
+        symbols: list[tuple[str, str]] = []
+        if import_symbol_node is not None:
+            idx = node.children.index(import_symbol_node)
+            symbols = self._collect_imported_symbols(node.children[idx + 1 :], code_bytes)
+
+        norm_path = file_path.replace("\\", "/")
+        is_package = norm_path == "__init__.py" or norm_path.endswith("/__init__.py")
+        base_module = (
+            resolve_relative_module(module_fqn, leading_dots, raw_module_str, is_package=is_package)
+            if leading_dots > 0
+            else raw_module_str
+        )
+
+        for local_name, sym in symbols:
+            target_fqn = f"{base_module}.{sym}" if base_module else sym
+            import_map[local_name] = target_fqn
+            # Symbol-level import edge (#161 slice 2): raw_import: candidates are
+            # resolved to an existing node by the ResolverEngine or DROPPED —
+            # they never leak into output (spec §2.2). Literal prefix mirrors
+            # the raw_dep:/raw_call: convention used elsewhere in this file.
+            edges.append(
+                Edge(
+                    id=f"{module_fqn}:imports_symbol:{target_fqn}",
+                    type=EdgeType.IMPORTS_SYMBOL,
+                    source=module_fqn,
+                    target=f"raw_import:{target_fqn}",
+                    confidence=0.1,
+                    file_path=file_path,
+                )
+            )
+
+        if base_module:
+            edges.append(
+                Edge(
+                    id=f"{module_fqn}:imports:{base_module}",
+                    type=EdgeType.IMPORTS,
+                    source=module_fqn,
+                    target=base_module,
+                    confidence=1.0,
+                    file_path=file_path,
+                )
+            )

cgis/extractors/_python_types.py

@@ -0,0 +1,84 @@
+"""Type-annotation resolution for the Python extractor."""
+
+from collections.abc import Callable
+
+from cgis.extractors._python_ast import file_path_to_module_fqn
+
+
+class TypeResolver:
+    """Resolves Python type annotations to fully qualified names.
+
+    Owns the type-string cleanup (stripping Optionals/Unions/generics) and the
+    import-map lookup that turns a bare or module-prefixed type name into a FQN.
+    """
+
+    # `Annotated[T, ...]` unwraps to its first arg T (the type); the remaining
+    # args are metadata (e.g. FastAPI `Depends(...)`, captured separately as a
+    # DEPENDS_ON edge via the call-node path). See #194.
+    _GENERIC_WRAPPERS: frozenset[str] = frozenset({"Optional", "Union", "Annotated"})
+
+    def __init__(self, pick_source_root: Callable[[str], str | None]) -> None:
+        """Store the per-file source-root picker used for FQN construction."""
+        self._pick_source_root = pick_source_root
+
+    def resolve_type_fqn(
+        self,
+        type_name: str,
+        import_map: dict[str, str] | None,
+        file_path: str,
+    ) -> str:
+        """Resolve a type name (possibly module-prefixed) to a FQN."""
+        if import_map:
+            if type_name in import_map:
+                return import_map[type_name]
+            if "." in type_name:
+                module_part, _, rest = type_name.partition(".")
+                if module_part in import_map:
+                    return f"{import_map[module_part]}.{rest}"
+        module = file_path_to_module_fqn(file_path, self._pick_source_root(file_path))
+        return f"{module}.{type_name}"
+
+    def clean_python_type_string(self, type_str: str) -> str:
+        """
+        Extract the main class name from type strings with Optionals/Unions.
+
+        Examples:
+            "SQLiteStore | None" -> "SQLiteStore"
+            "Optional[SQLiteStore]" -> "SQLiteStore"
+            "Union[SQLiteStore, None]" -> "SQLiteStore"
+            "list[Node]" -> "list"
+        """
+        # Handle union types (PEP 604): "A | None" -> "A"
+        if " | " in type_str:
+            type_str = type_str.split(" | ", maxsplit=1)[0]
+        # Handle generic wrappers: Optional[X] -> X, Union[X, Y] -> X, list[X] -> list
+        if "[" in type_str:
+            outer, _, inner = type_str.partition("[")
+            # Support both `Optional[X]` and `typing.Optional[X]`
+            outer_base = outer.strip().split(".")[-1]
+            if outer_base in self._GENERIC_WRAPPERS:
+                # Strip only the single outermost `]` to avoid mangling nested generics
+                if inner.endswith("]"):
+                    inner = inner[:-1]
+                first_arg = self._first_top_level_arg(inner)
+                return self.clean_python_type_string(first_arg)
+            type_str = outer
+        return type_str.strip()
+
+    @staticmethod
+    def _first_top_level_arg(args: str) -> str:
+        """Return the first argument of a generic's arg list.
+
+        Splits on the first comma at bracket depth 0 so nested generics keep
+        their inner commas intact, e.g. "dict[str, int], Meta" -> "dict[str, int]"
+        (a naive ``split(",")[0]`` would truncate to "dict[str").
+        """
+        depth = 0
+        for idx, char in enumerate(args):
+            if char == "[":
+                depth += 1
+            elif char == "]":
+                depth -= 1
+            elif char == "," and depth == 0:
+                return args[:idx].strip()
+        return args.strip()

cgis/extractors/base.py

@@ -0,0 +1,42 @@
+"""Base Extractor Interface"""
+
+from abc import ABC, abstractmethod
+
+from cgis.core.models import Edge, Node
+
+
+class BaseExtractor(ABC):
+    """
+    Abstract Base Class for all language-specific extractors.
+
+    Every extractor must implement the 'parse' method, which converts
+    raw source code into our standardized Node and Edge atoms.
+    """
+
+    def __init__(self, source_roots: list[str] | None = None) -> None:
+        """Store optional source roots used to strip prefixes from FQNs."""
+        # Pre-normalize once at init; _pick_source_root is called per node during parsing.
+        self._source_roots: list[tuple[str, str]] = [
+            (sr, sr.replace("\\", "/").strip("/") + "/") for sr in (source_roots or [])
+        ]
+
+    def _pick_source_root(self, file_path: str) -> str | None:
+        """Return the first source root that matches file_path as a prefix, or None."""
+        clean = file_path.replace("\\", "/").lstrip("/")
+        for sr, sr_norm in self._source_roots:
+            if clean.startswith(sr_norm):
+                return sr
+        return None
+
+    @abstractmethod
+    def parse(self, code: str, file_path: str) -> tuple[list[Node], list[Edge]]:
+        """
+        Parses the provided code and returns a list of Nodes and Edges.
+
+        Args:
+            code: The raw source code as a string.
+            file_path: The path to the file (used for FQN construction).
+
+        Returns:
+            A tuple containing (list_of_nodes, list_of_edges).
+        """