headroom-ai 0.2.13__py3-none-any.whl

headroom/compression/handlers/code_handler.py

@@ -0,0 +1,506 @@
+"""Code structure handler using AST parsing.
+
+Extracts structural elements from source code:
+- Import statements
+- Function/method signatures
+- Class definitions
+- Type annotations
+- Decorators
+
+Function bodies are marked as compressible while preserving signatures.
+This enables the LLM to see all available functions/methods while body
+implementations are compressed.
+
+Uses tree-sitter for parsing when available, falls back to regex patterns.
+"""
+
+from __future__ import annotations
+
+import logging
+import re
+import threading
+from dataclasses import dataclass
+from enum import Enum
+from typing import Any
+
+from headroom.compression.handlers.base import BaseStructureHandler, HandlerResult
+from headroom.compression.masks import StructureMask
+
+logger = logging.getLogger(__name__)
+
+# Lazy-loaded tree-sitter
+_tree_sitter_available: bool | None = None
+_tree_sitter_parsers: dict[str, Any] = {}
+_tree_sitter_lock = threading.Lock()
+
+
+def _check_tree_sitter() -> bool:
+    """Check if tree-sitter is available."""
+    global _tree_sitter_available
+    if _tree_sitter_available is None:
+        try:
+            import tree_sitter_language_pack  # noqa: F401
+
+            _tree_sitter_available = True
+        except ImportError:
+            _tree_sitter_available = False
+    return _tree_sitter_available
+
+
+def _get_parser(language: str) -> Any:
+    """Get tree-sitter parser for language."""
+    global _tree_sitter_parsers
+
+    if not _check_tree_sitter():
+        raise ImportError("tree-sitter-language-pack not installed")
+
+    with _tree_sitter_lock:
+        if language not in _tree_sitter_parsers:
+            from tree_sitter_language_pack import get_parser
+
+            _tree_sitter_parsers[language] = get_parser(language)  # type: ignore[arg-type]
+
+        return _tree_sitter_parsers[language]
+
+
+class CodeLanguage(Enum):
+    """Supported programming languages."""
+
+    PYTHON = "python"
+    JAVASCRIPT = "javascript"
+    TYPESCRIPT = "typescript"
+    GO = "go"
+    RUST = "rust"
+    JAVA = "java"
+    C = "c"
+    CPP = "cpp"
+
+
+@dataclass
+class CodeSpan:
+    """A span of code with its structural role."""
+
+    start: int
+    end: int
+    role: str  # "import", "signature", "body", "decorator", etc.
+    is_structural: bool
+
+
+# Language-specific AST node types that are structural
+_STRUCTURAL_NODE_TYPES: dict[str, set[str]] = {
+    "python": {
+        "import_statement",
+        "import_from_statement",
+        "function_definition",  # Just the signature part
+        "class_definition",
+        "decorated_definition",
+        "type_alias_statement",
+    },
+    "javascript": {
+        "import_statement",
+        "export_statement",
+        "function_declaration",
+        "class_declaration",
+        "method_definition",
+        "arrow_function",  # Signature only
+    },
+    "typescript": {
+        "import_statement",
+        "export_statement",
+        "function_declaration",
+        "class_declaration",
+        "method_definition",
+        "interface_declaration",
+        "type_alias_declaration",
+    },
+    "go": {
+        "import_declaration",
+        "function_declaration",
+        "method_declaration",
+        "type_declaration",
+        "interface_type",
+    },
+    "rust": {
+        "use_declaration",
+        "function_item",
+        "impl_item",
+        "struct_item",
+        "enum_item",
+        "trait_item",
+    },
+    "java": {
+        "import_declaration",
+        "class_declaration",
+        "method_declaration",
+        "interface_declaration",
+        "annotation",
+    },
+}
+
+# Regex patterns for fallback detection
+_SIGNATURE_PATTERNS: dict[str, list[re.Pattern[str]]] = {
+    "python": [
+        re.compile(r"^\s*(async\s+)?def\s+\w+\s*\([^)]*\)\s*(->\s*[^:]+)?:", re.MULTILINE),
+        re.compile(r"^\s*class\s+\w+(\([^)]*\))?:", re.MULTILINE),
+        re.compile(r"^\s*@\w+(\([^)]*\))?\s*$", re.MULTILINE),
+    ],
+    "javascript": [
+        re.compile(r"^\s*(async\s+)?function\s+\w+\s*\([^)]*\)", re.MULTILINE),
+        re.compile(r"^\s*class\s+\w+(\s+extends\s+\w+)?", re.MULTILINE),
+        re.compile(r"^\s*(const|let|var)\s+\w+\s*=\s*(async\s+)?\([^)]*\)\s*=>", re.MULTILINE),
+    ],
+    "typescript": [
+        re.compile(r"^\s*(async\s+)?function\s+\w+\s*(<[^>]+>)?\s*\([^)]*\)", re.MULTILINE),
+        re.compile(r"^\s*class\s+\w+(<[^>]+>)?(\s+extends\s+\w+)?", re.MULTILINE),
+        re.compile(r"^\s*interface\s+\w+(<[^>]+>)?", re.MULTILINE),
+        re.compile(r"^\s*type\s+\w+(<[^>]+>)?\s*=", re.MULTILINE),
+    ],
+    "go": [
+        re.compile(r"^\s*func\s+(\([^)]+\)\s+)?\w+\s*\([^)]*\)", re.MULTILINE),
+        re.compile(r"^\s*type\s+\w+\s+(struct|interface)", re.MULTILINE),
+    ],
+    "rust": [
+        re.compile(r"^\s*(pub\s+)?(async\s+)?fn\s+\w+\s*(<[^>]+>)?\s*\([^)]*\)", re.MULTILINE),
+        re.compile(r"^\s*(pub\s+)?struct\s+\w+", re.MULTILINE),
+        re.compile(r"^\s*(pub\s+)?enum\s+\w+", re.MULTILINE),
+        re.compile(r"^\s*(pub\s+)?trait\s+\w+", re.MULTILINE),
+        re.compile(r"^\s*impl(<[^>]+>)?\s+\w+", re.MULTILINE),
+    ],
+    "java": [
+        re.compile(
+            r"^\s*(public|private|protected)?\s*(static\s+)?\w+\s+\w+\s*\([^)]*\)", re.MULTILINE
+        ),
+        re.compile(r"^\s*(public\s+)?(class|interface|enum)\s+\w+", re.MULTILINE),
+        re.compile(r"^\s*@\w+(\([^)]*\))?\s*$", re.MULTILINE),
+    ],
+}
+
+# Import patterns for fallback
+_IMPORT_PATTERNS: dict[str, re.Pattern[str]] = {
+    "python": re.compile(r"^\s*(import\s+\w+|from\s+\w+\s+import)", re.MULTILINE),
+    "javascript": re.compile(r"^\s*(import\s+.*from|require\s*\()", re.MULTILINE),
+    "typescript": re.compile(r"^\s*(import\s+.*from|require\s*\()", re.MULTILINE),
+    "go": re.compile(r'^\s*import\s+(\(|")', re.MULTILINE),
+    "rust": re.compile(r"^\s*use\s+\w+", re.MULTILINE),
+    "java": re.compile(r"^\s*import\s+[\w.]+;", re.MULTILINE),
+}
+
+
+class CodeStructureHandler(BaseStructureHandler):
+    """Handler for source code.
+
+    Preserves:
+    - Import/use statements
+    - Function/method signatures (not bodies)
+    - Class/struct/interface definitions
+    - Type declarations
+    - Decorators/annotations
+
+    Marks as compressible:
+    - Function/method bodies
+    - Comments (optionally preserved)
+    - Whitespace
+
+    Example:
+        >>> handler = CodeStructureHandler()
+        >>> code = '''
+        ... def hello(name: str) -> str:
+        ...     message = f"Hello, {name}!"
+        ...     return message
+        ... '''
+        >>> result = handler.get_mask(code, language="python")
+        >>> # Signature "def hello(name: str) -> str:" preserved
+        >>> # Body content compressed
+    """
+
+    def __init__(
+        self,
+        preserve_comments: bool = False,
+        use_tree_sitter: bool = True,
+        default_language: str = "python",
+    ):
+        """Initialize the code handler.
+
+        Args:
+            preserve_comments: Whether to preserve comments as structural.
+            use_tree_sitter: Whether to use tree-sitter for parsing.
+                Falls back to regex if False or unavailable.
+            default_language: Default language when detection fails.
+        """
+        super().__init__(name="code")
+        self.preserve_comments = preserve_comments
+        self.use_tree_sitter = use_tree_sitter
+        self.default_language = default_language
+
+    def can_handle(self, content: str) -> bool:
+        """Check if content looks like source code."""
+        # Quick heuristic checks
+        code_indicators = [
+            "def ",
+            "class ",
+            "function ",
+            "import ",
+            "const ",
+            "let ",
+            "var ",
+            "func ",
+            "fn ",
+            "pub ",
+            "package ",
+            "struct ",
+            "interface ",
+        ]
+        return any(indicator in content for indicator in code_indicators)
+
+    def _extract_mask(
+        self,
+        content: str,
+        tokens: list[str],
+        language: str | None = None,
+        **kwargs: Any,
+    ) -> HandlerResult:
+        """Extract structure mask from code.
+
+        Args:
+            content: Source code content.
+            tokens: Character-level tokens.
+            language: Programming language (auto-detected if None).
+            **kwargs: Additional options.
+
+        Returns:
+            HandlerResult with mask marking structural elements.
+        """
+        # Detect language if not provided
+        if language is None:
+            language = self._detect_language(content)
+
+        # Try tree-sitter first
+        if self.use_tree_sitter and _check_tree_sitter():
+            try:
+                return self._extract_with_tree_sitter(content, tokens, language)
+            except Exception as e:
+                logger.debug("Tree-sitter parsing failed, using fallback: %s", e)
+
+        # Fallback to regex
+        return self._extract_with_regex(content, tokens, language)
+
+    def _extract_with_tree_sitter(
+        self,
+        content: str,
+        tokens: list[str],
+        language: str,
+    ) -> HandlerResult:
+        """Extract structure using tree-sitter AST.
+
+        Args:
+            content: Source code.
+            tokens: Character tokens.
+            language: Language name.
+
+        Returns:
+            HandlerResult with mask.
+        """
+        parser = _get_parser(language)
+        tree = parser.parse(content.encode("utf-8"))
+
+        # Collect structural spans
+        spans: list[CodeSpan] = []
+
+        def visit_node(node: Any, depth: int = 0) -> None:
+            """Visit AST node and collect structural spans."""
+            node_type = node.type
+            structural_types = _STRUCTURAL_NODE_TYPES.get(language, set())
+
+            # Check if this is a structural node type
+            if node_type in structural_types:
+                # For functions, only the signature is structural
+                if "function" in node_type or "method" in node_type:
+                    # Find the body node and exclude it
+                    body_node = None
+                    for child in node.children:
+                        if child.type in ("block", "statement_block", "compound_statement"):
+                            body_node = child
+                            break
+
+                    if body_node:
+                        # Signature is from start to body start
+                        spans.append(
+                            CodeSpan(
+                                start=node.start_byte,
+                                end=body_node.start_byte,
+                                role="signature",
+                                is_structural=True,
+                            )
+                        )
+                        # Body is compressible
+                        spans.append(
+                            CodeSpan(
+                                start=body_node.start_byte,
+                                end=body_node.end_byte,
+                                role="body",
+                                is_structural=False,
+                            )
+                        )
+                    else:
+                        # No body found, preserve whole thing
+                        spans.append(
+                            CodeSpan(
+                                start=node.start_byte,
+                                end=node.end_byte,
+                                role=node_type,
+                                is_structural=True,
+                            )
+                        )
+                else:
+                    # Non-function structural nodes
+                    spans.append(
+                        CodeSpan(
+                            start=node.start_byte,
+                            end=node.end_byte,
+                            role=node_type,
+                            is_structural=True,
+                        )
+                    )
+            elif node_type == "comment" and self.preserve_comments:
+                spans.append(
+                    CodeSpan(
+                        start=node.start_byte,
+                        end=node.end_byte,
+                        role="comment",
+                        is_structural=True,
+                    )
+                )
+
+            # Recurse into children
+            for child in node.children:
+                visit_node(child, depth + 1)
+
+        visit_node(tree.root_node)
+
+        # Build mask from spans
+        mask = self._spans_to_mask(spans, len(content))
+
+        return HandlerResult(
+            mask=StructureMask(tokens=tokens, mask=mask),
+            handler_name=self.name,
+            confidence=0.95,
+            metadata={
+                "language": language,
+                "parser": "tree-sitter",
+                "structural_spans": len([s for s in spans if s.is_structural]),
+            },
+        )
+
+    def _extract_with_regex(
+        self,
+        content: str,
+        tokens: list[str],
+        language: str,
+    ) -> HandlerResult:
+        """Extract structure using regex patterns (fallback).
+
+        Args:
+            content: Source code.
+            tokens: Character tokens.
+            language: Language name.
+
+        Returns:
+            HandlerResult with mask.
+        """
+        spans: list[CodeSpan] = []
+
+        # Match imports
+        import_pattern = _IMPORT_PATTERNS.get(language)
+        if import_pattern:
+            for match in import_pattern.finditer(content):
+                # Find end of import line
+                end = content.find("\n", match.end())
+                if end == -1:
+                    end = len(content)
+                spans.append(
+                    CodeSpan(
+                        start=match.start(),
+                        end=end,
+                        role="import",
+                        is_structural=True,
+                    )
+                )
+
+        # Match signatures
+        signature_patterns = _SIGNATURE_PATTERNS.get(language, [])
+        for pattern in signature_patterns:
+            for match in pattern.finditer(content):
+                spans.append(
+                    CodeSpan(
+                        start=match.start(),
+                        end=match.end(),
+                        role="signature",
+                        is_structural=True,
+                    )
+                )
+
+        # Build mask from spans
+        mask = self._spans_to_mask(spans, len(content))
+
+        return HandlerResult(
+            mask=StructureMask(tokens=tokens, mask=mask),
+            handler_name=self.name,
+            confidence=0.7,  # Lower confidence for regex
+            metadata={
+                "language": language,
+                "parser": "regex",
+                "structural_spans": len(spans),
+            },
+        )
+
+    def _spans_to_mask(self, spans: list[CodeSpan], length: int) -> list[bool]:
+        """Convert spans to character-level mask.
+
+        Args:
+            spans: List of code spans.
+            length: Total content length.
+
+        Returns:
+            Boolean mask aligned to characters.
+        """
+        mask = [False] * length
+
+        for span in spans:
+            if span.is_structural:
+                for i in range(span.start, min(span.end, length)):
+                    mask[i] = True
+
+        return mask
+
+    def _detect_language(self, content: str) -> str:
+        """Detect programming language from content.
+
+        Args:
+            content: Source code content.
+
+        Returns:
+            Language name (lowercase).
+        """
+        # Check for language-specific markers
+        markers = {
+            "python": ["def ", "import ", "from ", "class ", "async def"],
+            "javascript": ["function ", "const ", "let ", "var ", "=>"],
+            "typescript": ["interface ", "type ", ": string", ": number"],
+            "go": ["func ", "package ", "import (", "type "],
+            "rust": ["fn ", "let mut", "impl ", "pub fn", "use "],
+            "java": ["public class", "private ", "protected ", "void "],
+        }
+
+        scores: dict[str, int] = {}
+        for lang, patterns in markers.items():
+            scores[lang] = sum(1 for p in patterns if p in content)
+
+        if not scores or max(scores.values()) == 0:
+            return self.default_language
+
+        return max(scores, key=lambda k: scores[k])
+
+
+def is_tree_sitter_available() -> bool:
+    """Check if tree-sitter is available."""
+    return _check_tree_sitter()