shotgun-sh 0.2.29.dev2__py3-none-any.whl → 0.6.1.dev1__py3-none-any.whl

shotgun/codebase/core/extractors/protocol.py

@@ -0,0 +1,109 @@
+"""Protocol definition for language extractors."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Protocol, runtime_checkable
+
+if TYPE_CHECKING:
+    from tree_sitter import Node
+
+    from .types import SupportedLanguage
+
+
+@runtime_checkable
+class LanguageExtractor(Protocol):
+    """Protocol for language-specific AST extraction.
+
+    Each language implementation provides methods to extract:
+    - Decorators/attributes from function/class nodes
+    - Docstrings/documentation
+    - Inheritance information from class definitions
+    - Call expression parsing
+    - Node type information for the language
+    """
+
+    @property
+    def language(self) -> SupportedLanguage:
+        """The language this extractor handles."""
+        ...
+
+    def extract_decorators(self, node: Node) -> list[str]:
+        """Extract decorators/attributes from a function or class node.
+
+        Args:
+            node: The AST node (function or class definition)
+
+        Returns:
+            List of decorator/attribute names
+        """
+        ...
+
+    def extract_docstring(self, node: Node) -> str | None:
+        """Extract documentation string from a function or class node.
+
+        Args:
+            node: The AST node (function or class definition)
+
+        Returns:
+            The docstring content, or None if not present
+        """
+        ...
+
+    def extract_inheritance(self, class_node: Node) -> list[str]:
+        """Extract parent class names from a class definition.
+
+        Args:
+            class_node: The class definition AST node
+
+        Returns:
+            List of parent class names (simple names, may need resolution)
+        """
+        ...
+
+    def parse_call_node(self, call_node: Node) -> tuple[str | None, str | None]:
+        """Parse a call expression node to extract callee information.
+
+        Args:
+            call_node: The call expression AST node
+
+        Returns:
+            Tuple of (callee_name, object_name) where:
+            - callee_name: The name of the function/method being called
+            - object_name: The object the method is called on (for method calls)
+        """
+        ...
+
+    def find_parent_class(self, func_node: Node, module_qn: str) -> str | None:
+        """Find the parent class of a function node.
+
+        Args:
+            func_node: The function definition AST node
+            module_qn: Module qualified name
+
+        Returns:
+            Qualified name of parent class, or None if not in a class
+        """
+        ...
+
+    def find_containing_function(self, node: Node, module_qn: str) -> str | None:
+        """Find the containing function/method of a node.
+
+        Args:
+            node: The AST node
+            module_qn: Module qualified name
+
+        Returns:
+            Qualified name of containing function, or None
+        """
+        ...
+
+    def count_ast_nodes(self, node: Node) -> int:
+        """Count total AST nodes for metrics.
+
+        Args:
+            node: Root AST node
+
+        Returns:
+            Total node count
+        """
+        ...

shotgun/codebase/core/extractors/python/__init__.py

@@ -0,0 +1,7 @@
+"""Python language extractor."""
+
+from __future__ import annotations
+
+from .extractor import PythonExtractor
+
+__all__ = ["PythonExtractor"]

shotgun/codebase/core/extractors/python/extractor.py

@@ -0,0 +1,141 @@
+"""Python language extractor implementation."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from shotgun.codebase.core.extractors.base import BaseExtractor
+from shotgun.codebase.core.extractors.types import SupportedLanguage
+
+if TYPE_CHECKING:
+    from tree_sitter import Node
+
+
+class PythonExtractor(BaseExtractor):
+    """Extractor for Python source code."""
+
+    @property
+    def language(self) -> SupportedLanguage:
+        """The language this extractor handles."""
+        return SupportedLanguage.PYTHON
+
+    def _class_definition_types(self) -> list[str]:
+        """Return node types that represent class definitions."""
+        return ["class_definition"]
+
+    def _function_definition_types(self) -> list[str]:
+        """Return node types that represent function definitions."""
+        return ["function_definition"]
+
+    def extract_decorators(self, node: Node) -> list[str]:
+        """Extract decorators from a function or class node.
+
+        Args:
+            node: The AST node (function or class definition)
+
+        Returns:
+            List of decorator names
+        """
+        decorators: list[str] = []
+
+        for child in node.children:
+            if child.type == "decorator":
+                for grandchild in child.children:
+                    if grandchild.type == "identifier" and grandchild.text:
+                        decorators.append(grandchild.text.decode("utf-8"))
+                        break
+                    elif grandchild.type == "attribute":
+                        attr_node = grandchild.child_by_field_name("attribute")
+                        if attr_node and attr_node.text:
+                            decorators.append(attr_node.text.decode("utf-8"))
+                            break
+
+        return decorators
+
+    def extract_docstring(self, node: Node) -> str | None:
+        """Extract docstring from a function or class node.
+
+        Args:
+            node: The AST node (function or class definition)
+
+        Returns:
+            The docstring content, or None if not present
+        """
+        body_node = node.child_by_field_name("body")
+        if not body_node or not body_node.children:
+            return None
+
+        first_statement = body_node.children[0]
+        if first_statement.type == "expression_statement":
+            for child in first_statement.children:
+                if child.type == "string" and child.text:
+                    docstring = child.text.decode("utf-8")
+                    docstring = docstring.strip()
+                    if (
+                        docstring.startswith('"""')
+                        and docstring.endswith('"""')
+                        or docstring.startswith("'''")
+                        and docstring.endswith("'''")
+                    ):
+                        docstring = docstring[3:-3]
+                    elif (
+                        docstring.startswith('"')
+                        and docstring.endswith('"')
+                        or docstring.startswith("'")
+                        and docstring.endswith("'")
+                    ):
+                        docstring = docstring[1:-1]
+                    return docstring.strip()
+
+        return None
+
+    def extract_inheritance(self, class_node: Node) -> list[str]:
+        """Extract parent class names from a class definition.
+
+        Args:
+            class_node: The class definition AST node
+
+        Returns:
+            List of parent class names (simple names, may need resolution)
+        """
+        parent_names: list[str] = []
+
+        for child in class_node.children:
+            if child.type == "argument_list":
+                for arg in child.children:
+                    if arg.type == "identifier" and arg.text:
+                        parent_names.append(arg.text.decode("utf-8"))
+                    elif arg.type == "attribute":
+                        full_name_parts: list[str] = []
+                        self._extract_full_name(arg, full_name_parts)
+                        if full_name_parts:
+                            parent_names.append(".".join(full_name_parts))
+
+        return parent_names
+
+    def parse_call_node(self, call_node: Node) -> tuple[str | None, str | None]:
+        """Parse a call expression node to extract callee information.
+
+        Args:
+            call_node: The call expression AST node
+
+        Returns:
+            Tuple of (callee_name, object_name)
+        """
+        callee_name = None
+        object_name = None
+
+        for child in call_node.children:
+            if child.type == "identifier" and child.text:
+                callee_name = child.text.decode("utf-8")
+                break
+            elif child.type == "attribute":
+                obj_node = child.child_by_field_name("object")
+                attr_node = child.child_by_field_name("attribute")
+                if obj_node and obj_node.text:
+                    object_name = obj_node.text.decode("utf-8")
+                if attr_node and attr_node.text:
+                    callee_name = attr_node.text.decode("utf-8")
+                    break
+
+        return callee_name, object_name

shotgun/codebase/core/extractors/rust/__init__.py

@@ -0,0 +1,7 @@
+"""Rust language extractor."""
+
+from __future__ import annotations
+
+from .extractor import RustExtractor
+
+__all__ = ["RustExtractor"]

shotgun/codebase/core/extractors/rust/extractor.py

@@ -0,0 +1,139 @@
+"""Rust language extractor implementation."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from shotgun.codebase.core.extractors.base import BaseExtractor
+from shotgun.codebase.core.extractors.types import SupportedLanguage
+
+if TYPE_CHECKING:
+    from tree_sitter import Node
+
+
+class RustExtractor(BaseExtractor):
+    """Extractor for Rust source code."""
+
+    @property
+    def language(self) -> SupportedLanguage:
+        """The language this extractor handles."""
+        return SupportedLanguage.RUST
+
+    def _class_definition_types(self) -> list[str]:
+        """Return node types that represent type definitions.
+
+        Rust has structs, enums, traits, and type aliases.
+        """
+        return ["struct_item", "enum_item", "trait_item", "type_item"]
+
+    def _function_definition_types(self) -> list[str]:
+        """Return node types that represent function definitions."""
+        return ["function_item", "function_signature_item"]
+
+    def extract_decorators(self, node: Node) -> list[str]:
+        """Extract attributes from a function or struct node.
+
+        Rust uses #[attribute] syntax.
+
+        Args:
+            node: The AST node
+
+        Returns:
+            List of attribute names
+        """
+        attributes: list[str] = []
+
+        for child in node.children:
+            if child.type == "attribute_item":
+                for attr_child in child.children:
+                    if attr_child.type == "attribute":
+                        path = attr_child.child_by_field_name("path")
+                        if path and path.text:
+                            attributes.append(path.text.decode("utf-8"))
+
+        return attributes
+
+    def extract_docstring(self, node: Node) -> str | None:
+        """Extract doc comment from a function or type node.
+
+        Rust uses /// for outer doc comments and //! for inner doc comments.
+
+        Args:
+            node: The AST node
+
+        Returns:
+            The doc comment, or None if not present
+        """
+        doc_lines: list[str] = []
+        prev_sibling = node.prev_named_sibling
+
+        while prev_sibling and prev_sibling.type == "line_comment":
+            comment_text = prev_sibling.text
+            if comment_text:
+                text = comment_text.decode("utf-8")
+                if text.startswith("///"):
+                    doc_lines.insert(0, text[3:].strip())
+                else:
+                    break
+            prev_sibling = prev_sibling.prev_named_sibling
+
+        if doc_lines:
+            return "\n".join(doc_lines)
+        return None
+
+    def extract_inheritance(self, class_node: Node) -> list[str]:
+        """Extract trait bounds or supertraits from a type definition.
+
+        For structs, this returns nothing (Rust doesn't have struct inheritance).
+        For traits, this returns supertraits.
+
+        Args:
+            class_node: The type definition AST node
+
+        Returns:
+            List of supertrait names
+        """
+        supertraits: list[str] = []
+
+        if class_node.type == "trait_item":
+            for child in class_node.children:
+                if child.type == "trait_bounds":
+                    for bound in child.children:
+                        if bound.type == "type_identifier" and bound.text:
+                            supertraits.append(bound.text.decode("utf-8"))
+                        elif bound.type == "generic_type":
+                            type_node = bound.child_by_field_name("type")
+                            if type_node and type_node.text:
+                                supertraits.append(type_node.text.decode("utf-8"))
+
+        return supertraits
+
+    def parse_call_node(self, call_node: Node) -> tuple[str | None, str | None]:
+        """Parse a call expression node to extract callee information.
+
+        Args:
+            call_node: The call expression AST node
+
+        Returns:
+            Tuple of (callee_name, object_name)
+        """
+        callee_name = None
+        object_name = None
+
+        func_node = call_node.child_by_field_name("function")
+        if func_node:
+            if func_node.type == "identifier" and func_node.text:
+                callee_name = func_node.text.decode("utf-8")
+            elif func_node.type == "field_expression":
+                value = func_node.child_by_field_name("value")
+                field = func_node.child_by_field_name("field")
+                if value and value.text:
+                    object_name = value.text.decode("utf-8")
+                if field and field.text:
+                    callee_name = field.text.decode("utf-8")
+            elif func_node.type == "scoped_identifier":
+                name = func_node.child_by_field_name("name")
+                if name and name.text:
+                    callee_name = name.text.decode("utf-8")
+
+        return callee_name, object_name

shotgun/codebase/core/extractors/types.py

@@ -0,0 +1,15 @@
+"""Type definitions for the extractors module."""
+
+from __future__ import annotations
+
+from enum import StrEnum
+
+
+class SupportedLanguage(StrEnum):
+    """Supported programming languages for AST extraction."""
+
+    PYTHON = "python"
+    JAVASCRIPT = "javascript"
+    TYPESCRIPT = "typescript"
+    GO = "go"
+    RUST = "rust"

shotgun/codebase/core/extractors/typescript/__init__.py

@@ -0,0 +1,7 @@
+"""TypeScript language extractor."""
+
+from __future__ import annotations
+
+from .extractor import TypeScriptExtractor
+
+__all__ = ["TypeScriptExtractor"]

shotgun/codebase/core/extractors/typescript/extractor.py

@@ -0,0 +1,92 @@
+"""TypeScript language extractor implementation."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from shotgun.codebase.core.extractors.javascript.extractor import JavaScriptExtractor
+from shotgun.codebase.core.extractors.types import SupportedLanguage
+
+if TYPE_CHECKING:
+    from tree_sitter import Node
+
+
+class TypeScriptExtractor(JavaScriptExtractor):
+    """Extractor for TypeScript source code.
+
+    TypeScript is a superset of JavaScript, so this extends
+    JavaScriptExtractor with TypeScript-specific features.
+    """
+
+    @property
+    def language(self) -> SupportedLanguage:
+        """The language this extractor handles."""
+        return SupportedLanguage.TYPESCRIPT
+
+    def _class_definition_types(self) -> list[str]:
+        """Return node types that represent class definitions.
+
+        TypeScript adds interface_declaration and type_alias_declaration.
+        """
+        return [
+            "class_declaration",
+            "class",
+            "interface_declaration",
+            "type_alias_declaration",
+        ]
+
+    def extract_decorators(self, node: Node) -> list[str]:
+        """Extract decorators from a function or class node.
+
+        TypeScript supports decorators (experimental feature).
+
+        Args:
+            node: The AST node
+
+        Returns:
+            List of decorator names
+        """
+        decorators: list[str] = []
+
+        for child in node.children:
+            if child.type == "decorator":
+                for grandchild in child.children:
+                    if grandchild.type == "identifier" and grandchild.text:
+                        decorators.append(grandchild.text.decode("utf-8"))
+                        break
+                    elif grandchild.type == "call_expression":
+                        for call_child in grandchild.children:
+                            if call_child.type == "identifier" and call_child.text:
+                                decorators.append(call_child.text.decode("utf-8"))
+                                break
+
+        return decorators
+
+    def extract_inheritance(self, class_node: Node) -> list[str]:
+        """Extract parent class/interface names from a class or interface.
+
+        TypeScript classes can extend one class and implement multiple interfaces.
+        Interfaces can extend multiple interfaces.
+
+        Args:
+            class_node: The class/interface definition AST node
+
+        Returns:
+            List of parent names
+        """
+        parent_names: list[str] = []
+
+        for child in class_node.children:
+            if child.type in ["extends_clause", "implements_clause"]:
+                for type_node in child.children:
+                    if type_node.type == "type_identifier" and type_node.text:
+                        parent_names.append(type_node.text.decode("utf-8"))
+                    elif type_node.type == "generic_type":
+                        name_node = type_node.child_by_field_name("name")
+                        if name_node and name_node.text:
+                            parent_names.append(name_node.text.decode("utf-8"))
+
+        if not parent_names:
+            parent_names = super().extract_inheritance(class_node)
+
+        return parent_names