"""Code-aware compressor using AST parsing for syntax-preserving compression. This module provides AST-based compression for source code that guarantees valid syntax output. Unlike token-level compression, this preserves structural elements while compressing function bodies. Key Features: - Syntax validity guaranteed (output always parses) - Preserves imports, signatures, type annotations, error handlers - Compresses function bodies while maintaining structure - Multi-language support via tree-sitter - Data-driven language config (no per-language method duplication) - Thread-safe (thread-local tree-sitter parsers, no shared mutable state) Supported Languages (Tier 1): - Python, JavaScript, TypeScript Supported Languages (Tier 2): - Go, Rust, Java, C, C++ Compression Strategy: 1. Parse code into AST using tree-sitter 2. Extract and preserve critical structures (imports, signatures, types) 3. Rank functions by importance (using semantic analysis) 4. Compress function bodies while preserving signatures 5. Reassemble into valid code Installation: pip install headroom-ai[code] Usage: >>> from headroom.transforms import CodeAwareCompressor >>> compressor = CodeAwareCompressor() >>> result = compressor.compress(python_code) >>> print(result.compressed) # Valid Python code >>> print(result.syntax_valid) # True Reference: LongCodeZip: Compress Long Context for Code Language Models https://arxiv.org/abs/2510.00446 """ from __future__ import annotations import logging import re import threading from dataclasses import dataclass, field from enum import Enum from typing import Any from ..config import TransformResult from ..tokenizer import Tokenizer from .base import Transform logger = logging.getLogger(__name__) # Lazy import for optional dependency _tree_sitter_available: bool | None = None _tree_sitter_local = threading.local() def _check_tree_sitter_available() -> bool: """Check if tree-sitter is available *and actually parses*. The mere presence of ``tree_sitter_language_pack`` is not enough: prior versions of this code green-lit a code path that raised ``TypeError`` at parse time and silently fell back to a lossy stripper. To stop misleading callers, we now verify an end-to-end parse of a tiny snippet and only return ``True`` if it yields a real AST. """ global _tree_sitter_available if _tree_sitter_available is None: try: parser = _get_parser("python") tree = parser.parse(b"def _probe():\n return 1\n") root = tree.root_node # A real parse yields a non-error root with children. _tree_sitter_available = ( root is not None and root.type == "module" and root.child_count > 0 and not _has_syntax_issues(root) ) except Exception: _tree_sitter_available = False return _tree_sitter_available def _tree_sitter_importable() -> bool: """Return True if the tree-sitter grammar pack can be imported. This only checks importability (cheap, no parse). Use :func:`_check_tree_sitter_available` for the stronger "parsing actually works" guarantee. """ try: import tree_sitter_language_pack # noqa: F401 return True except ImportError: return False def _get_parser(language: str) -> Any: """Get a tree-sitter parser for the given language. Returns a **thread-local** ``tree_sitter.Parser`` instance. tree-sitter ≥ 0.23 wraps the C ``TSParser`` in a PyO3 ``#[pyclass(unsendable)]`` which hard-panics if the object is accessed from any thread other than its creator. Because Headroom runs compression inside a ``ThreadPoolExecutor``, a single shared parser would be touched from arbitrary pool threads → instant crash. We use the stock ``tree_sitter.Parser`` (which returns standard ``tree_sitter.Node`` / ``tree_sitter.Tree`` with property access) and set its language via ``tree_sitter_language_pack.get_language()``. Storing one parser per (thread, language) satisfies the ``unsendable`` contract with negligible extra memory. Args: language: Language name (e.g., 'python', 'javascript'). Returns: Configured ``tree_sitter.Parser`` bound to the current thread. Raises: ImportError: If tree-sitter is not installed. ValueError: If language is not supported. """ # NOTE: guard on importability (not _check_tree_sitter_available), because # _check_tree_sitter_available now performs a real end-to-end parse via # _get_parser; guarding on it here would recurse. if not _tree_sitter_importable(): raise ImportError( "tree-sitter is not installed. Install with: pip install headroom-ai[code]\n" "This adds ~50MB for tree-sitter grammars." ) parsers: dict[str, Any] | None = getattr(_tree_sitter_local, "parsers", None) if parsers is None: parsers = {} _tree_sitter_local.parsers = parsers if language not in parsers: try: from tree_sitter import Parser from tree_sitter_language_pack import get_language parser = Parser() # `language` is a validated runtime str; get_language types its arg # as a Literal of supported names, which a dynamic str can't satisfy. parser.language = get_language(language) # type: ignore[arg-type] parsers[language] = parser logger.debug( "Loaded tree-sitter parser for %s (thread %s)", language, threading.current_thread().name, ) except Exception as e: raise ValueError( f"Language '{language}' is not supported by tree-sitter. " f"Supported: python, javascript, typescript, go, rust, java, c, cpp, csharp, perl. " f"Error: {e}" ) from e return parsers[language] def is_tree_sitter_available() -> bool: """Check if tree-sitter is installed and available. Returns: True if tree-sitter-languages package is installed. """ return _check_tree_sitter_available() def is_tree_sitter_loaded() -> bool: """Check if any tree-sitter parsers are loaded on the current thread. Returns: True if parsers are loaded in this thread's local storage. """ parsers: dict[str, Any] | None = getattr(_tree_sitter_local, "parsers", None) return bool(parsers) def unload_tree_sitter() -> bool: """Unload tree-sitter parsers on the current thread to free memory. Returns: True if parsers were unloaded, False if none were loaded. """ parsers: dict[str, Any] | None = getattr(_tree_sitter_local, "parsers", None) if parsers: count = len(parsers) parsers.clear() logger.info( "Unloaded %d tree-sitter parsers (thread %s)", count, threading.current_thread().name ) return True return False class CodeLanguage(Enum): """Supported programming languages.""" PYTHON = "python" JAVASCRIPT = "javascript" TYPESCRIPT = "typescript" GO = "go" RUST = "rust" JAVA = "java" C = "c" CPP = "cpp" PERL = "perl" CSHARP = "csharp" UNKNOWN = "unknown" # Common language hints and markdown fence tags that are not the canonical # ``CodeLanguage`` value. Mapping them here keeps ``` ```js ``` / ``` ```ts ``` # / ``` ```py ``` fenced blocks (and callers that pass an alias) on the # code-aware path instead of raising ValueError. _LANGUAGE_ALIASES: dict[str, CodeLanguage] = { "js": CodeLanguage.JAVASCRIPT, "jsx": CodeLanguage.JAVASCRIPT, "mjs": CodeLanguage.JAVASCRIPT, "cjs": CodeLanguage.JAVASCRIPT, "node": CodeLanguage.JAVASCRIPT, "ts": CodeLanguage.TYPESCRIPT, "tsx": CodeLanguage.TYPESCRIPT, "py": CodeLanguage.PYTHON, "python3": CodeLanguage.PYTHON, "golang": CodeLanguage.GO, "rs": CodeLanguage.RUST, "c++": CodeLanguage.CPP, "cxx": CodeLanguage.CPP, "cc": CodeLanguage.CPP, "hpp": CodeLanguage.CPP, "pl": CodeLanguage.PERL, } def coerce_language(value: str) -> CodeLanguage: """Map a language hint or markdown fence tag to a ``CodeLanguage``. Accepts the canonical enum values and common aliases/fence tags (``js``/``ts``/``py``/...). Unknown strings return ``CodeLanguage.UNKNOWN`` instead of raising ``ValueError`` from ``CodeLanguage(value)``, so an unrecognized fence tag falls back to content-based detection rather than crashing the caller (or, inside the router, silently skipping code-aware compression because the ValueError is swallowed). """ key = (value or "").strip().lower() if not key: return CodeLanguage.UNKNOWN try: return CodeLanguage(key) except ValueError: return _LANGUAGE_ALIASES.get(key, CodeLanguage.UNKNOWN) class DocstringMode(Enum): """How to handle docstrings.""" FULL = "full" # Keep entire docstring FIRST_LINE = "first_line" # Keep only first line REMOVE = "remove" # Remove docstrings completely NONE = "none" # Alias for REMOVE (deprecated) # ========================================================================= # Data-driven language configuration # ========================================================================= @dataclass(frozen=True) class LangConfig: """Data-driven configuration for a programming language. Instead of per-language methods, each language declares its AST node types and syntactic conventions. The compressor uses these tables to drive extraction and compression generically. """ # AST node types for structural extraction import_nodes: frozenset[str] function_nodes: frozenset[str] class_nodes: frozenset[str] type_nodes: frozenset[str] body_node_types: frozenset[str] # Node types that represent function/method bodies decorator_node: str | None # e.g. "decorated_definition" for Python # Syntax conventions comment_prefix: str # "#" for Python, "//" for C-family uses_colon_after_signature: bool # Python: True, C-family: False package_node: str | None = None # e.g. "package_clause" for Go # Quick pre-filter hints for language detection (substrings to check) detection_hints: tuple[str, ...] = () # Optional override for node types that contain class/impl members. class_body_node_types: frozenset[str] | None = None # Optional container nodes (e.g. C# block-scoped namespaces) that wrap # type declarations in a member list. Routed through class compression so # their members compress and the wrapper isn't re-emitted verbatim. container_node_types: frozenset[str] | None = None # Optional opaque nodes (e.g. C# `#if`/`#endif` conditionals) preserved # verbatim without recursing into them. Recursing would capture their # descendants individually while the top-level pass re-emits the whole # wrapper verbatim, duplicating content. Prefer the false negative. opaque_node_types: frozenset[str] | None = None _LANG_CONFIGS: dict[CodeLanguage, LangConfig] = { CodeLanguage.PYTHON: LangConfig( import_nodes=frozenset( {"future_import_statement", "import_statement", "import_from_statement"} ), function_nodes=frozenset({"function_definition"}), class_nodes=frozenset({"class_definition"}), type_nodes=frozenset({"type_alias_statement"}), body_node_types=frozenset({"block"}), decorator_node="decorated_definition", comment_prefix="#", uses_colon_after_signature=True, detection_hints=("def ", "import ", "from ", "class ", "async def"), ), CodeLanguage.JAVASCRIPT: LangConfig( import_nodes=frozenset({"import_statement", "import_declaration"}), function_nodes=frozenset({"function_declaration", "method_definition"}), class_nodes=frozenset({"class_declaration"}), type_nodes=frozenset(), body_node_types=frozenset({"statement_block"}), decorator_node=None, comment_prefix="//", uses_colon_after_signature=False, detection_hints=("function ", "const ", "let ", "var ", "export ", "require("), class_body_node_types=frozenset({"class_body"}), ), CodeLanguage.TYPESCRIPT: LangConfig( import_nodes=frozenset({"import_statement", "import_declaration"}), function_nodes=frozenset({"function_declaration", "method_definition"}), class_nodes=frozenset({"class_declaration"}), type_nodes=frozenset({"interface_declaration", "type_alias_declaration"}), body_node_types=frozenset({"statement_block"}), decorator_node=None, comment_prefix="//", uses_colon_after_signature=False, detection_hints=("interface ", "type ", ": string", ": number", ": boolean"), class_body_node_types=frozenset({"class_body"}), ), CodeLanguage.GO: LangConfig( import_nodes=frozenset({"import_declaration"}), function_nodes=frozenset({"function_declaration", "method_declaration"}), class_nodes=frozenset(), type_nodes=frozenset({"type_declaration"}), body_node_types=frozenset({"block"}), decorator_node=None, comment_prefix="//", uses_colon_after_signature=False, package_node="package_clause", detection_hints=("func ", "package ", "struct {"), ), CodeLanguage.RUST: LangConfig( import_nodes=frozenset({"use_declaration"}), function_nodes=frozenset({"function_item"}), class_nodes=frozenset({"impl_item"}), type_nodes=frozenset({"struct_item", "enum_item", "type_item", "trait_item"}), body_node_types=frozenset({"block"}), decorator_node=None, comment_prefix="//", uses_colon_after_signature=False, detection_hints=("fn ", "struct ", "impl ", "mod ", "use "), class_body_node_types=frozenset({"declaration_list"}), ), CodeLanguage.JAVA: LangConfig( import_nodes=frozenset({"import_declaration"}), function_nodes=frozenset({"method_declaration", "constructor_declaration"}), class_nodes=frozenset({"class_declaration", "interface_declaration"}), type_nodes=frozenset({"enum_declaration"}), body_node_types=frozenset({"block"}), decorator_node=None, comment_prefix="//", uses_colon_after_signature=False, package_node="package_declaration", detection_hints=("public ", "private ", "protected ", "class ", "interface "), class_body_node_types=frozenset({"class_body"}), ), CodeLanguage.C: LangConfig( import_nodes=frozenset({"preproc_include"}), function_nodes=frozenset({"function_definition"}), class_nodes=frozenset(), type_nodes=frozenset({"struct_specifier", "enum_specifier", "type_definition"}), body_node_types=frozenset({"compound_statement"}), decorator_node=None, comment_prefix="//", uses_colon_after_signature=False, detection_hints=("#include", "typedef ", "int main("), ), CodeLanguage.CPP: LangConfig( import_nodes=frozenset({"preproc_include"}), function_nodes=frozenset({"function_definition"}), class_nodes=frozenset({"class_specifier"}), type_nodes=frozenset({"struct_specifier", "enum_specifier", "type_definition"}), body_node_types=frozenset({"compound_statement"}), decorator_node=None, comment_prefix="//", uses_colon_after_signature=False, detection_hints=("#include", "namespace ", "class ", "::"), class_body_node_types=frozenset({"field_declaration_list"}), ), CodeLanguage.PERL: LangConfig( import_nodes=frozenset({"use_statement", "use_version_statement"}), function_nodes=frozenset( {"subroutine_declaration_statement", "method_declaration_statement"} ), class_nodes=frozenset({"package_statement", "class_statement", "role_statement"}), type_nodes=frozenset(), body_node_types=frozenset({"block"}), decorator_node=None, comment_prefix="#", uses_colon_after_signature=False, package_node="package_statement", detection_hints=("sub ", "my ", "our ", "use ", "package "), ), CodeLanguage.CSHARP: LangConfig( import_nodes=frozenset({"using_directive", "file_scoped_namespace_declaration"}), function_nodes=frozenset( { "method_declaration", "constructor_declaration", "destructor_declaration", "operator_declaration", "local_function_statement", } ), class_nodes=frozenset( { "class_declaration", "struct_declaration", "record_declaration", "interface_declaration", } ), type_nodes=frozenset({"enum_declaration", "delegate_declaration"}), body_node_types=frozenset({"block"}), decorator_node=None, comment_prefix="//", uses_colon_after_signature=False, detection_hints=("using ", "namespace ", "public ", "private ", "void "), class_body_node_types=frozenset({"declaration_list"}), container_node_types=frozenset({"namespace_declaration"}), opaque_node_types=frozenset({"preproc_if"}), ), } @dataclass class CodeStructure: """Extracted structure from parsed code.""" header_code: list[str] = field(default_factory=list) imports: list[str] = field(default_factory=list) type_definitions: list[str] = field(default_factory=list) class_definitions: list[str] = field(default_factory=list) function_signatures: list[str] = field(default_factory=list) function_bodies: list[tuple[str, str, int]] = field( default_factory=list ) # (signature, body, line) decorators: list[str] = field(default_factory=list) comments: list[str] = field(default_factory=list) top_level_code: list[str] = field(default_factory=list) other: list[str] = field(default_factory=list) @dataclass class CodeCompressorConfig: """Configuration for code-aware compression. Attributes: preserve_imports: Always keep import statements. preserve_signatures: Always keep function/method signatures. preserve_type_annotations: Keep type hints and annotations. preserve_decorators: Keep decorators on functions/classes. docstring_mode: How to handle docstrings. target_compression_rate: Target compression ratio (0.2 = keep 20%). max_body_lines: Maximum lines to keep per function body. compress_comments: Remove non-docstring comments. min_tokens_for_compression: Minimum tokens to trigger compression. language_hint: Explicit language (None = auto-detect). fallback_to_kompress: Use Kompress for unknown languages. enable_ccr: Store originals for retrieval. ccr_ttl: TTL for CCR entries in seconds. """ # Preservation settings preserve_imports: bool = True preserve_signatures: bool = True preserve_type_annotations: bool = True preserve_decorators: bool = True docstring_mode: DocstringMode = DocstringMode.FIRST_LINE # Compression settings target_compression_rate: float = 0.2 max_body_lines: int = 5 compress_comments: bool = True # Thresholds min_tokens_for_compression: int = 100 # Language handling language_hint: str | None = None fallback_to_kompress: bool = True # Semantic analysis (symbol importance scoring) semantic_analysis: bool = True # CCR integration enable_ccr: bool = True ccr_ttl: int = 300 # 5 minutes @dataclass class CodeCompressionResult: """Result of code-aware compression. Attributes: compressed: The compressed code (guaranteed valid syntax). original: Original code before compression. original_tokens: Token count before compression. compressed_tokens: Token count after compression. compression_ratio: Actual compression ratio achieved. language: Detected or specified language. language_confidence: Confidence in language detection. preserved_imports: Number of import statements preserved. preserved_signatures: Number of function signatures preserved. compressed_bodies: Number of function bodies compressed. syntax_valid: Whether output is syntactically valid. cache_key: CCR cache key if stored. """ compressed: str original: str original_tokens: int compressed_tokens: int compression_ratio: float # Code-specific metadata language: CodeLanguage = CodeLanguage.UNKNOWN language_confidence: float = 0.0 # Structure analysis preserved_imports: int = 0 preserved_signatures: int = 0 compressed_bodies: int = 0 # Validation syntax_valid: bool = True # CCR cache_key: str | None = None # Semantic analysis symbol_scores: dict[str, float] = field(default_factory=dict) @property def tokens_saved(self) -> int: """Number of tokens saved by compression.""" return max(0, self.original_tokens - self.compressed_tokens) @property def savings_percentage(self) -> float: """Percentage of tokens saved.""" if self.original_tokens == 0: return 0.0 return (self.tokens_saved / self.original_tokens) * 100 @property def summary(self) -> str: """Human-readable summary of compression.""" analysis_note = "" if self.symbol_scores: high = sum(1 for s in self.symbol_scores.values() if s >= 0.7) low = sum(1 for s in self.symbol_scores.values() if s < 0.1) if high or low: analysis_note = f" Semantic: {high} high-importance, {low} low-importance." return ( f"Compressed {self.language.value} code: " f"{self.original_tokens:,}→{self.compressed_tokens:,} tokens " f"({self.savings_percentage:.0f}% saved). " f"Kept {self.preserved_imports} imports, " f"{self.preserved_signatures} signatures, " f"compressed {self.compressed_bodies} bodies." f"{analysis_note}" ) # ========================================================================= # Language detection # ========================================================================= # Lightweight pre-filter patterns for language detection. # These are ONLY used as a quick check to avoid parsing with every language. # Actual detection is done by tree-sitter (fewest parse errors wins). _LANGUAGE_PREFILTER: dict[CodeLanguage, list[re.Pattern[str]]] = { CodeLanguage.PYTHON: [ re.compile(r"^\s*(def|class|import|from|async def)\s+\w+", re.MULTILINE), re.compile(r"^\s*@\w+", re.MULTILINE), re.compile(r'^\s*"""', re.MULTILINE), re.compile(r"^\s*if __name__\s*==", re.MULTILINE), ], CodeLanguage.JAVASCRIPT: [ re.compile(r"^\s*(function|const|let|var|class|export)\s+\w+", re.MULTILINE), re.compile(r"^\s*async\s+(function|=>)", re.MULTILINE), re.compile(r"^\s*module\.exports", re.MULTILINE), re.compile(r"^\s*(import|export)\s+.*\s+from\s+['\"]", re.MULTILINE), ], CodeLanguage.TYPESCRIPT: [ re.compile(r"^\s*(interface|type|enum|namespace)\s+\w+", re.MULTILINE), re.compile(r":\s*(string|number|boolean|any|void|Promise)\b", re.MULTILINE), ], CodeLanguage.GO: [ re.compile(r"^\s*(func|type|package|import)\s+", re.MULTILINE), re.compile(r"^\s*func\s+\([^)]+\)\s+\w+", re.MULTILINE), re.compile(r"\bstruct\s*\{", re.MULTILINE), ], CodeLanguage.RUST: [ re.compile(r"^\s*(fn|struct|enum|impl|mod|use|pub)\s+", re.MULTILINE), re.compile(r"^\s*#\[", re.MULTILINE), ], CodeLanguage.JAVA: [ re.compile(r"^\s*(public|private|protected)\s+(class|interface|enum)", re.MULTILINE), re.compile(r"^\s*package\s+[\w.]+;", re.MULTILINE), ], CodeLanguage.C: [ re.compile(r"^\s*#include\s*[<\"]", re.MULTILINE), re.compile(r"^\s*(int|void|char|float|double)\s+\w+\s*\(", re.MULTILINE), re.compile(r"^\s*typedef\s+", re.MULTILINE), ], CodeLanguage.CPP: [ re.compile(r"^\s*#include\s*[<\"]", re.MULTILINE), re.compile(r"\bnamespace\s+\w+", re.MULTILINE), re.compile(r"::\w+", re.MULTILINE), ], CodeLanguage.PERL: [ re.compile(r"^\s*(sub|package|use|require)\s+[\w:]+", re.MULTILINE), re.compile(r"^\s*(my|our|local)\s+[\$@%]", re.MULTILINE), re.compile(r"[\$@%]\w+", re.MULTILINE), ], CodeLanguage.CSHARP: [ re.compile(r"^\s*using\s+[\w.]+\s*;", re.MULTILINE), re.compile(r"^\s*namespace\s+[\w.]+", re.MULTILINE), re.compile( r"^\s*(public|private|protected|internal|sealed|static|abstract|partial)\s+" r"(class|struct|record|interface|enum)\b", re.MULTILINE, ), re.compile(r"\bget;\s*set;", re.MULTILINE), ], } def _count_error_nodes(node: Any) -> int: """Count ERROR and MISSING nodes in a tree-sitter AST.""" count = 0 if node.type == "ERROR" or node.is_missing: count += 1 for child in node.children: count += _count_error_nodes(child) return count def detect_language(code: str) -> tuple[CodeLanguage, float]: """Detect the programming language of code. Uses tree-sitter AST parsing when available (most accurate), with a regex pre-filter to avoid parsing with all languages. Falls back to regex-only scoring when tree-sitter is unavailable. Args: code: Source code to analyze. Returns: Tuple of (detected language, confidence score 0.0-1.0). """ if not code or not code.strip(): return CodeLanguage.UNKNOWN, 0.0 sample = code[:5000] # Phase 1: Pre-filter — find candidate languages using quick regex candidates: dict[CodeLanguage, int] = {} for lang, patterns in _LANGUAGE_PREFILTER.items(): score = 0 for pattern in patterns: matches = len(pattern.findall(sample)) score += matches if score > 0: candidates[lang] = score if not candidates: return CodeLanguage.UNKNOWN, 0.0 # Disambiguation: TypeScript superset of JavaScript if CodeLanguage.TYPESCRIPT in candidates and CodeLanguage.JAVASCRIPT in candidates: if candidates[CodeLanguage.TYPESCRIPT] >= 2: candidates[CodeLanguage.JAVASCRIPT] = 0 # Disambiguation: C++ superset of C if CodeLanguage.CPP in candidates and CodeLanguage.C in candidates: if candidates[CodeLanguage.CPP] >= 2: candidates[CodeLanguage.C] = 0 # Phase 2: If tree-sitter available, parse with candidates and pick fewest errors if _check_tree_sitter_available(): best_lang = CodeLanguage.UNKNOWN min_errors = float("inf") best_node_count = 0 code_bytes = bytes(code[:10000], "utf-8") # Sort candidates by pre-filter score (try most likely first) sorted_candidates = sorted(candidates.items(), key=lambda x: x[1], reverse=True) for lang, _prefilter_score in sorted_candidates: if lang == CodeLanguage.UNKNOWN or candidates.get(lang, 0) == 0: continue try: parser = _get_parser(lang.value) tree = parser.parse(code_bytes) error_count = _count_error_nodes(tree.root_node) node_count = tree.root_node.child_count # Prefer: fewest errors, then most top-level nodes (richer parse) if error_count < min_errors or ( error_count == min_errors and node_count > best_node_count ): min_errors = error_count best_lang = lang best_node_count = node_count except (ValueError, ImportError): continue if best_lang != CodeLanguage.UNKNOWN: # Confidence based on error ratio total_lines = max(1, len(code.strip().split("\n"))) error_ratio = min_errors / total_lines confidence = max(0.3, min(1.0, 1.0 - error_ratio)) return best_lang, confidence # Phase 3: Fallback — regex-only scoring (no tree-sitter) best_lang = max(candidates, key=lambda k: candidates[k]) best_score = candidates[best_lang] if best_score == 0: return CodeLanguage.UNKNOWN, 0.0 confidence = min(1.0, 0.3 + (best_score * 0.1)) return best_lang, confidence # ========================================================================= # Symbol importance analysis # ========================================================================= @dataclass class _SymbolAnalysis: """Result of intra-file symbol importance analysis. All dicts are keyed by qualified name (e.g., 'ClassName.method') to avoid collisions between identically-named methods in different classes. """ scores: dict[str, float] = field(default_factory=dict) calls: dict[str, set[str]] = field(default_factory=dict) ref_counts: dict[str, int] = field(default_factory=dict) body_line_counts: dict[str, int] = field(default_factory=dict) bare_names: dict[str, str] = field(default_factory=dict) # qname -> short_name class CodeAwareCompressor(Transform): """AST-preserving compression for source code. This compressor uses tree-sitter to parse code into an AST, then selectively compresses function bodies while preserving structure. The output is guaranteed to be syntactically valid. Key advantages over token-level compression: - Syntax validity guaranteed - Preserves imports, signatures, types - Better compression ratios for code (5-8x vs 3-5x) - Lower latency (~20-50ms vs 50-200ms for token-level compressors) - Smaller memory footprint (~50MB vs ~1GB) - Thread-safe (thread-local tree-sitter parsers, no shared mutable state) Example: >>> compressor = CodeAwareCompressor() >>> result = compressor.compress(''' ... import os ... from typing import List ... ... def process_data(items: List[str]) -> List[str]: ... \"\"\"Process a list of items.\"\"\" ... results = [] ... for item in items: ... # Validate item ... if not item: ... continue ... # Process valid item ... processed = item.strip().lower() ... results.append(processed) ... return results ... ''') >>> print(result.compressed) import os from typing import List def process_data(items: List[str]) -> List[str]: \"\"\"Process a list of items.\"\"\" # ... (body compressed: 10 lines → 2 lines) pass """ name: str = "code_aware_compressor" def __init__(self, config: CodeCompressorConfig | None = None): """Initialize code-aware compressor. Args: config: Compression configuration. If None, uses defaults. Note: Tree-sitter parsers are loaded lazily on first use to avoid startup overhead when the compressor isn't used. """ self.config = config or CodeCompressorConfig() # ========================================================================= # Token estimation # ========================================================================= @staticmethod def _estimate_tokens(text: str, tokenizer: Tokenizer | None = None) -> int: """Count or estimate tokens for text. Uses real tokenizer when available; falls back to chars/4 which is a much closer approximation for code than word count. """ if tokenizer is not None: return tokenizer.count_text(text) # chars/4 is a reasonable approximation for code tokens # (code has lots of punctuation that tokenizes separately) return max(1, len(text) // 4) # ========================================================================= # Symbol importance analysis # ========================================================================= def _analyze_symbol_importance( self, root: Any, code: str, language: CodeLanguage, context: str = "", ) -> _SymbolAnalysis: """Analyze symbol importance using distribution-based scoring. Collects raw signals (reference count, fan-out, visibility, context match, convention importance) per symbol, then normalizes using min-max scaling so scores are relative within the file. This adapts to any file structure: utility libraries, test files, orchestrators, etc. Returns _SymbolAnalysis with normalized scores (0.0-1.0) per symbol. """ if not self.config.semantic_analysis: return _SymbolAnalysis() lang_config = _LANG_CONFIGS.get(language) if not lang_config: return _SymbolAnalysis() all_definition_types = lang_config.function_nodes | lang_config.class_nodes # Use qualified keys (ClassName.method) to avoid collisions definitions: dict[str, Any] = {} # qualified_name -> node bare_names: dict[str, str] = {} # qualified_name -> short_name all_identifiers: dict[str, int] = {} # short_name -> count function_calls: dict[str, set[str]] = {} def collect_definitions(node: Any, parent_name: str = "") -> None: if node.type in all_definition_types: short_name = _get_definition_name(node) if short_name: qualified = f"{parent_name}.{short_name}" if parent_name else short_name definitions[qualified] = node bare_names[qualified] = short_name for child in node.children: collect_definitions(child, parent_name=qualified) return # Also check for decorated definitions if lang_config.decorator_node and node.type == lang_config.decorator_node: for child in node.children: if child.type in all_definition_types: short_name = _get_definition_name(child) if short_name: qualified = f"{parent_name}.{short_name}" if parent_name else short_name definitions[qualified] = child bare_names[qualified] = short_name for grandchild in child.children: collect_definitions(grandchild, parent_name=qualified) return for child in node.children: collect_definitions(child, parent_name) def collect_identifiers(node: Any) -> None: if node.type in ("identifier", "property_identifier", "type_identifier"): text = node.text name = text.decode("utf-8") if isinstance(text, bytes) else str(text) all_identifiers[name] = all_identifiers.get(name, 0) + 1 for child in node.children: collect_identifiers(child) def collect_calls_in_function(func_node: Any, func_qname: str) -> None: func_short = bare_names[func_qname] defined_short_names = set(bare_names.values()) calls: set[str] = set() def walk(node: Any) -> None: if node.type in ("identifier", "property_identifier"): text = node.text name = text.decode("utf-8") if isinstance(text, bytes) else str(text) if name in defined_short_names and name != func_short: calls.add(name) for child in node.children: walk(child) walk(func_node) function_calls[func_qname] = calls # Pass 1: Collect definitions with qualified names collect_definitions(root) if not definitions: return _SymbolAnalysis() # Pass 2: Collect all identifiers collect_identifiers(root) # Pass 3: Collect call relationships and body sizes body_line_counts: dict[str, int] = {} for qname, node in definitions.items(): collect_calls_in_function(node, qname) node_text = _slice_code_bytes(code, node.start_byte, node.end_byte) body_line_counts[qname] = max(1, len(node_text.split("\n")) - 2) # Reference counts: subtract definition occurrences short_name_def_count: dict[str, int] = {} for short in bare_names.values(): short_name_def_count[short] = short_name_def_count.get(short, 0) + 1 ref_counts: dict[str, int] = {} for qname in definitions: short = bare_names[qname] count = all_identifiers.get(short, 0) ref_counts[qname] = max(0, count - short_name_def_count.get(short, 1)) # Raw importance signals per symbol context_words, context_lower, context_has_cjk = _query_context_tokens(context) raw_signals: dict[str, float] = {} for qname in definitions: short = bare_names[qname] refs = ref_counts.get(qname, 0) fan_out = len(function_calls.get(qname, set())) is_public = _is_public_symbol(short, language) raw = float(refs) raw += 1.0 if is_public else 0.0 raw += fan_out * 0.5 # Convention importance (language-specific) if language == CodeLanguage.PYTHON: if short.startswith("__") and short.endswith("__"): raw += 2.0 elif language == CodeLanguage.GO: if short and short[0].isupper(): raw += 1.0 # Context boost: the relevance query named this symbol. if _symbol_in_context(short.lower(), context_words, context_lower, context_has_cjk): raw += 3.0 raw_signals[qname] = raw # Normalize to 0-1 using min-max scaling values = list(raw_signals.values()) min_val = min(values) max_val = max(values) range_val = max_val - min_val if range_val > 0: scores = {name: round((v - min_val) / range_val, 3) for name, v in raw_signals.items()} else: scores = dict.fromkeys(raw_signals, 0.5) return _SymbolAnalysis( scores=scores, calls=function_calls, ref_counts=ref_counts, body_line_counts=body_line_counts, bare_names=bare_names, ) def _allocate_body_budget(self, analysis: _SymbolAnalysis, code: str) -> dict[str, int]: """Allocate body line budget across functions using target_compression_rate. Returns dict mapping symbol name to max body lines to keep. """ if not analysis.scores or not analysis.body_line_counts: return {} scores = analysis.scores body_sizes = analysis.body_line_counts target_rate = self.config.target_compression_rate total_lines = len(code.strip().split("\n")) total_body_lines = sum(body_sizes.values()) fixed_lines = max(0, total_lines - total_body_lines) target_total = total_lines * target_rate body_budget = max(0.0, target_total - fixed_lines) if total_body_lines == 0: return {} score_floor = 0.05 weights: dict[str, float] = {} for name in scores: score = max(scores.get(name, 0.5), score_floor) size = body_sizes.get(name, 0) weights[name] = score * size total_weight = sum(weights.values()) if total_weight == 0: per_func = max(0, int(body_budget / max(len(scores), 1))) return {name: min(per_func, body_sizes.get(name, 0)) for name in scores} limits: dict[str, int] = {} for qname in scores: allocation = body_budget * weights[qname] / total_weight max_lines = body_sizes.get(qname, 0) limit = min(int(round(allocation)), max_lines) limits[qname] = limit # Also store by short name so _get_body_limit can find it. short = analysis.bare_names.get(qname, qname) if short not in limits or limit > limits[short]: limits[short] = limit return limits # ========================================================================= # Core compression # ========================================================================= def compress( self, code: str, language: str | None = None, context: str = "", tokenizer: Tokenizer | None = None, ) -> CodeCompressionResult: """Compress code while preserving syntax validity. Args: code: Source code to compress. language: Language name (e.g., 'python'). Auto-detected if None. context: Optional context for relevance-aware compression. tokenizer: Optional tokenizer for accurate token counting. Returns: CodeCompressionResult with compressed code and metadata. """ if not code or not code.strip(): return CodeCompressionResult( compressed=code, original=code, original_tokens=0, compressed_tokens=0, compression_ratio=1.0, syntax_valid=True, ) original_tokens = self._estimate_tokens(code, tokenizer) # Skip small content if original_tokens < self.config.min_tokens_for_compression: return CodeCompressionResult( compressed=code, original=code, original_tokens=original_tokens, compressed_tokens=original_tokens, compression_ratio=1.0, syntax_valid=True, ) # Detect or use specified language. An explicit hint or fence tag may be # an alias (js/ts/py/...) or something we don't recognize — coerce it # instead of constructing CodeLanguage() directly (which raises), and # fall back to content detection when the hint is unknown. if language: detected_lang = coerce_language(language) if detected_lang == CodeLanguage.UNKNOWN: detected_lang, confidence = detect_language(code) else: confidence = 1.0 elif self.config.language_hint: detected_lang = coerce_language(self.config.language_hint) if detected_lang == CodeLanguage.UNKNOWN: detected_lang, confidence = detect_language(code) else: confidence = 1.0 else: detected_lang, confidence = detect_language(code) # If language unknown and fallback enabled, try Kompress if detected_lang == CodeLanguage.UNKNOWN: if self.config.fallback_to_kompress: return self._fallback_compress(code, original_tokens) else: return CodeCompressionResult( compressed=code, original=code, original_tokens=original_tokens, compressed_tokens=original_tokens, compression_ratio=1.0, language=CodeLanguage.UNKNOWN, language_confidence=0.0, syntax_valid=True, ) # Check if tree-sitter is available if not _check_tree_sitter_available(): logger.warning("tree-sitter not available. Install with: pip install headroom-ai[code]") if self.config.fallback_to_kompress: return self._fallback_compress(code, original_tokens) return CodeCompressionResult( compressed=code, original=code, original_tokens=original_tokens, compressed_tokens=original_tokens, compression_ratio=1.0, language=detected_lang, language_confidence=confidence, syntax_valid=True, ) # Parse and compress try: compressed, structure, symbol_scores = self._compress_with_ast( code, detected_lang, context, tokenizer ) compressed_tokens = self._estimate_tokens(compressed, tokenizer) # Verify syntax validity (checks both ERROR and MISSING nodes) syntax_valid = self._verify_syntax(compressed, detected_lang) # If syntax invalid, return original (never serve broken code) if not syntax_valid: logger.warning( "Code compression produced invalid syntax for %s (%d tokens), " "returning original", detected_lang.value, original_tokens, ) return CodeCompressionResult( compressed=code, original=code, original_tokens=original_tokens, compressed_tokens=original_tokens, compression_ratio=1.0, language=detected_lang, language_confidence=confidence, syntax_valid=True, ) ratio = compressed_tokens / max(original_tokens, 1) # Guard against over-aggressive compression (data loss) if ratio < 0.05: logger.warning( "Code compression too aggressive (ratio=%.3f), returning original", ratio, ) return CodeCompressionResult( compressed=code, original=code, original_tokens=original_tokens, compressed_tokens=original_tokens, compression_ratio=1.0, language=detected_lang, language_confidence=confidence, syntax_valid=True, ) # Store in CCR if significant compression cache_key = None if self.config.enable_ccr and ratio < 0.8: cache_key = self._store_in_ccr(code, compressed, original_tokens) if cache_key: from .compression_summary import summarize_compressed_code code_summary = summarize_compressed_code( structure.function_bodies, len(structure.function_bodies), ) summary_str = f" {code_summary}." if code_summary else "" # Use the actual config attribute (not the wrong name) ttl_min = max(1, self.config.ccr_ttl // 60) compressed += ( f"\n# [{original_tokens - compressed_tokens} tokens compressed." f"{summary_str}" f" Retrieve more: hash={cache_key}." f" Expires in {ttl_min}m.]" ) return CodeCompressionResult( compressed=compressed, original=code, original_tokens=original_tokens, compressed_tokens=compressed_tokens, compression_ratio=ratio, language=detected_lang, language_confidence=confidence, preserved_imports=len(structure.imports), preserved_signatures=len(structure.function_signatures), compressed_bodies=len(structure.function_bodies), syntax_valid=syntax_valid, cache_key=cache_key, symbol_scores=symbol_scores, ) except Exception as e: logger.warning("AST compression failed: %s, falling back", e) if self.config.fallback_to_kompress: return self._fallback_compress(code, original_tokens) return CodeCompressionResult( compressed=code, original=code, original_tokens=original_tokens, compressed_tokens=original_tokens, compression_ratio=1.0, language=detected_lang, language_confidence=confidence, syntax_valid=True, ) def _compress_with_ast( self, code: str, language: CodeLanguage, context: str, tokenizer: Tokenizer | None = None, ) -> tuple[str, CodeStructure, dict[str, float]]: """Compress code using AST parsing with symbol importance analysis. Thread-safe: all mutable state is passed through parameters, not stored on self. Args: code: Source code. language: Detected language. context: User context for relevance. tokenizer: Optional tokenizer for accurate token counting. Returns: Tuple of (compressed code, extracted structure, symbol scores). """ parser = _get_parser(language.value) tree = parser.parse(bytes(code, "utf-8")) root = tree.root_node # Analyze symbol importance and allocate compression budget analysis = self._analyze_symbol_importance(root, code, language, context) body_limits = self._allocate_body_budget(analysis, code) # Extract structure using data-driven language config lang_config = _LANG_CONFIGS.get(language) if lang_config: structure = self._extract_structure( root, code, language, lang_config, body_limits, analysis ) else: structure = self._extract_generic_structure(root, code) # Assemble compressed code compressed = self._assemble_compressed(structure, language) # Expose scores with short names for the public API symbol_scores: dict[str, float] = {} if analysis.scores: for qname, score in analysis.scores.items(): short = analysis.bare_names.get(qname, qname) if short not in symbol_scores or score > symbol_scores[short]: symbol_scores[short] = score return compressed, structure, symbol_scores # ========================================================================= # Unified structure extraction (data-driven, replaces per-language methods) # ========================================================================= def _extract_structure( self, root: Any, code: str, language: CodeLanguage, lang_config: LangConfig, body_limits: dict[str, int], analysis: _SymbolAnalysis, ) -> CodeStructure: """Extract structure from AST using data-driven language config. A single visitor handles all languages by checking node types against the LangConfig tables. No per-language extraction methods needed. """ structure = CodeStructure() captured_byte_ranges: list[tuple[int, int]] = [] def visit(node: Any) -> None: node_type = node.type # Package declarations (Go, Java) if lang_config.package_node and node_type == lang_config.package_node: leading = _get_leading_comment_text(node, code, captured_byte_ranges) structure.imports.insert(0, leading + _get_node_text(node, code)) captured_byte_ranges.append((node.start_byte, node.end_byte)) return # Import statements if node_type in lang_config.import_nodes: leading = _get_leading_comment_text(node, code, captured_byte_ranges) structure.imports.append(leading + _get_node_text(node, code)) captured_byte_ranges.append((node.start_byte, node.end_byte)) return # Export statements (JS/TS) — may contain functions or re-exports if node_type == "export_statement": leading = _get_leading_comment_text(node, code, captured_byte_ranges) text = _get_node_text(node, code) # Check if this export wraps a function or class has_func_or_class = False for child in node.children: if ( child.type in lang_config.function_nodes or child.type in lang_config.class_nodes ): has_func_or_class = True compressed = self._compress_function_ast( child, code, language, lang_config, body_limits, analysis ) # Reconstruct export with compressed inner definition export_prefix = _slice_code_bytes(code, node.start_byte, child.start_byte) export_suffix = _slice_code_bytes(code, child.end_byte, node.end_byte) structure.function_signatures.append( leading + export_prefix + compressed + export_suffix ) break if not has_func_or_class: structure.imports.append(leading + text) captured_byte_ranges.append((node.start_byte, node.end_byte)) return # Decorated definitions (Python) if lang_config.decorator_node and node_type == lang_config.decorator_node: leading = _get_leading_comment_text(node, code, captured_byte_ranges) decorator_text = [] definition_compressed = None for child in node.children: if child.type == "decorator": decorator_text.append(_get_node_text(child, code)) elif child.type in lang_config.function_nodes: definition_compressed = self._compress_function_ast( child, code, language, lang_config, body_limits, analysis ) elif child.type in lang_config.class_nodes: definition_compressed = self._compress_class_ast( child, code, language, lang_config, body_limits, analysis ) if decorator_text and definition_compressed: full_def = leading + "\n".join(decorator_text) + "\n" + definition_compressed # Route to correct list based on inner definition type for child in node.children: if child.type in lang_config.class_nodes: structure.class_definitions.append(full_def) break else: structure.function_signatures.append(full_def) elif definition_compressed: structure.function_signatures.append(leading + definition_compressed) captured_byte_ranges.append((node.start_byte, node.end_byte)) return # Function/method definitions if node_type in lang_config.function_nodes: leading = _get_leading_comment_text(node, code, captured_byte_ranges) compressed = self._compress_function_ast( node, code, language, lang_config, body_limits, analysis ) structure.function_signatures.append(leading + compressed) captured_byte_ranges.append((node.start_byte, node.end_byte)) return if lang_config.container_node_types and node_type in lang_config.container_node_types: compressed = self._compress_class_ast( node, code, language, lang_config, body_limits, analysis ) structure.class_definitions.append(compressed) captured_byte_ranges.append((node.start_byte, node.end_byte)) return # Class definitions — compress each method individually if node_type in lang_config.class_nodes: leading = _get_leading_comment_text(node, code, captured_byte_ranges) compressed = self._compress_class_ast( node, code, language, lang_config, body_limits, analysis ) structure.class_definitions.append(leading + compressed) captured_byte_ranges.append((node.start_byte, node.end_byte)) trailing_semicolon = _get_same_line_trailing_semicolon(node) if trailing_semicolon is not None: captured_byte_ranges.append( (trailing_semicolon.start_byte, trailing_semicolon.end_byte) ) return # Type definitions if node_type in lang_config.type_nodes: leading = _get_leading_comment_text(node, code, captured_byte_ranges) structure.type_definitions.append(leading + _get_node_text(node, code)) captured_byte_ranges.append((node.start_byte, node.end_byte)) return # Opaque regions (e.g. C# preprocessor conditionals) — preserved # verbatim, never recursed into. Recursing captures descendants # individually while the top-level pass re-emits the uncaptured # wrapper verbatim, duplicating its whole content in the output. # Blocks wrapping only import directives (`#if NET6_0\nusing X;`) # are emitted with the imports: usings must precede type # declarations, so appending them as trailing top-level code would # produce invalid output (and fall back to no compression). if lang_config.opaque_node_types and node_type in lang_config.opaque_node_types: child_types = [child.type for child in node.named_children] has_import = any(t in lang_config.import_nodes for t in child_types) has_declaration = any( t in lang_config.class_nodes or t in lang_config.type_nodes or t in lang_config.function_nodes or (lang_config.container_node_types and t in lang_config.container_node_types) for t in child_types ) if has_import and not has_declaration: structure.imports.append(_get_node_text(node, code)) else: structure.top_level_code.append(_get_node_text(node, code)) captured_byte_ranges.append((node.start_byte, node.end_byte)) return # Recurse into children for child in node.children: visit(child) visit(root) # Capture top-level code that wasn't handled by any of the above. # This preserves global variables, constants, if __name__ blocks, # module-level assignments, etc. Uncaptured nodes that precede the # first captured node are file headers (license banners, C# `#region # License` blocks, module comments): they are emitted first, in # original order, rather than relocated to the end of the output — # e.g. tree-sitter-c-sharp rejects top-level `#region` after a type # declaration, which would fail validation and forfeit compression. first_captured = min((r[0] for r in captured_byte_ranges), default=None) for child in root.children: child_range = (child.start_byte, child.end_byte) if child_range not in captured_byte_ranges: text = _get_node_text(child, code).strip() if text: if first_captured is not None and child.end_byte <= first_captured: structure.header_code.append(text) else: structure.top_level_code.append(text) return structure # ========================================================================= # Unified function/class compression (data-driven) # ========================================================================= def _compress_function_ast( self, node: Any, code: str, language: CodeLanguage, lang_config: LangConfig, body_limits: dict[str, int], analysis: _SymbolAnalysis, ) -> str: """Compress a function/class/impl block using AST body detection. Uses the AST to find the body node directly instead of string-scanning for '{' or ':'. Works for all languages via lang_config.body_node_types. Key insight: tree-sitter byte offsets may not include leading whitespace on the first line. We use LINE-based slicing from the original code to preserve indentation faithfully. """ # Use line-based slicing from original code (not byte offsets) to # preserve indentation. This is critical for nested definitions # (methods inside classes). code_lines = code.split("\n") start_row = node.start_point[0] node_lines = _get_node_lines(node, code_lines) node_text = "\n".join(node_lines) func_name = _get_definition_name(node) body_limit = _get_body_limit(func_name, body_limits, self.config.max_body_lines) # Small enough to keep as-is if len(node_lines) <= body_limit + 2: return node_text # Find the body node using AST (not string scanning) body_node = None for child in node.children: if child.type in lang_config.body_node_types: body_node = child break if body_node is None: return node_text # Use line numbers to slice: this preserves original indentation. # tree-sitter gives 0-based row numbers. node_start_line = node.start_point[0] body_start_line = body_node.start_point[0] body_end_line = body_node.end_point[0] # Lines within the node (0-indexed relative to node start) sig_end = body_start_line - node_start_line # exclusive body_end_rel = body_end_line - node_start_line + 1 # inclusive # Handle case where signature and body start on the SAME line # (common in brace languages: `function foo(arg) { ... }`) if sig_end == 0 and not lang_config.uses_colon_after_signature: # Signature and body on same line: `function foo(arg) { ... }` # Keep them together: sig includes up to and including `{` first_line = node_lines[0] # Include the opening brace in the signature line sig_with_brace = first_line.rstrip() signature_lines = [sig_with_brace] # Body lines are everything between { and } (inner content only) body_lines = node_lines[1:body_end_rel] after_lines = node_lines[body_end_rel:] # We've already included { in signature, so mark it _brace_in_signature = True else: signature_lines = node_lines[:sig_end] body_lines = node_lines[sig_end:body_end_rel] after_lines = node_lines[body_end_rel:] _brace_in_signature = False # For brace languages, detect opening/closing braces in the body lines. opening_brace_line = None closing_brace_line = None if not lang_config.uses_colon_after_signature: if _brace_in_signature: # Opening brace already in signature line — just find closing pass elif body_lines and body_lines[0].strip().endswith("{"): # Matches both a bare `{` line and a multi-line signature's # closing line (e.g. Go's `) error {`), where the brace # shares a line with the closing paren/return type rather # than starting one of its own. opening_brace_line = body_lines[0] body_lines = body_lines[1:] if body_lines and body_lines[-1].strip().endswith("}"): closing_brace_line = body_lines[-1] body_lines = body_lines[:-1] # Handle Python docstrings via AST docstring_text = "" ds_skip_lines = 0 if language == CodeLanguage.PYTHON and body_node.child_count > 0: first_child = body_node.children[0] # tree-sitter Python may represent docstrings as: # - bare `string` node directly in block, OR # - `expression_statement` containing a `string` node ds_node = None if first_child.type == "string": ds_node = first_child elif first_child.type == "expression_statement" and first_child.child_count > 0: if first_child.children[0].type == "string": ds_node = first_child if ds_node is not None: ds_lines_count = ds_node.end_point[0] - ds_node.start_point[0] + 1 ds_start_rel = ds_node.start_point[0] - body_node.start_point[0] if self.config.docstring_mode == DocstringMode.FULL: # Keep entire docstring as-is (preserve indentation from body_lines) docstring_text = "\n".join( body_lines[ds_start_rel : ds_start_rel + ds_lines_count] ) elif self.config.docstring_mode == DocstringMode.FIRST_LINE: # Use source lines directly (safe — preserves original quoting) if ds_lines_count == 1: # Single-line docstring: keep as-is docstring_text = body_lines[ds_start_rel] else: # Multi-line docstring: keep first line, close it properly first_ds_line = body_lines[ds_start_rel] ds_indent = first_ds_line[ : len(first_ds_line) - len(first_ds_line.lstrip()) ] stripped = first_ds_line.strip() # Detect quote style from source quote = '"""' for q in ('r"""', "r'''", '"""', "'''"): if stripped.startswith(q): quote = q[-3:] break # Find where content starts (after opening quotes + prefix) content_start = 0 for opener in ('r"""', "r'''", '"""', "'''"): if stripped.startswith(opener): content_start = len(opener) break first_content = stripped[content_start:].strip() # Remove trailing closing quotes if the first line has them for q in ('"""', "'''"): if first_content.endswith(q): first_content = first_content[: -len(q)].strip() if first_content: # """Some text here\n...\n""" → """Some text here""" prefix_part = stripped[:content_start] docstring_text = f"{ds_indent}{prefix_part}{first_content}{quote}" else: # Opening quote on its own line: """\n text\n""" if ds_start_rel + 1 < len(body_lines): second_line = body_lines[ds_start_rel + 1].strip() for q in ('"""', "'''"): if second_line.endswith(q): second_line = second_line[: -len(q)].strip() if second_line: docstring_text = f"{ds_indent}{quote}{second_line}{quote}" else: docstring_text = first_ds_line else: docstring_text = first_ds_line # elif REMOVE: docstring_text stays empty ds_skip_lines = ds_start_rel + ds_lines_count # --- Statement-based body truncation (never cuts mid-expression) --- # # Walk body_node.children (AST statements) instead of slicing lines. # Each child is a complete, syntactically valid statement. We keep # whole statements until the line budget is exhausted, so the output # always parses correctly. # Detect indentation from actual body code (preserves whatever the file uses) indent = _detect_indent(body_lines) if body_lines else " " # Collect non-docstring body statements from the AST body_stmts: list[tuple[int, int]] = [] # (start_row, end_row) absolute ds_end_row = -1 if ds_skip_lines > 0 and body_node.child_count > 0: # The docstring node occupies the first ds_skip_lines lines ds_end_row = body_node.start_point[0] + ds_skip_lines - 1 # Punctuation tokens to skip (brace-language body delimiters, semicolons) _SKIP_TYPES = frozenset({"{", "}", ";", ",", "comment", "line_comment", "block_comment"}) for child in body_node.children: # Skip docstring node (already handled separately) if child.start_point[0] <= ds_end_row: continue # Skip punctuation and comment nodes if child.type in _SKIP_TYPES: continue # Skip unnamed tokens (tree-sitter anonymous nodes like braces) if not child.is_named: continue # Some grammars (e.g. Go) wrap all body statements in one generic # list node instead of exposing them as direct siblings of the # block. Treating that wrapper as a single statement makes its # row range swallow the block's own closing brace line, causing # a duplicated `}` later. Unwrap it into its real statements. if child.type == "statement_list": for inner in child.children: if inner.type in _SKIP_TYPES or not inner.is_named: continue body_stmts.append((inner.start_point[0], inner.end_point[0])) continue body_stmts.append((child.start_point[0], child.end_point[0])) # Calculate lines per statement and keep whole statements until budget kept_lines: list[str] = [] kept_line_count = 0 stmts_kept = 0 total_body_lines_count = sum(end - start + 1 for start, end in body_stmts) for start_row, end_row in body_stmts: stmt_lines = code_lines[start_row : end_row + 1] stmt_line_count = len(stmt_lines) # If adding this statement would exceed budget and we already have # at least one statement, stop here if kept_line_count + stmt_line_count > body_limit and stmts_kept > 0: break kept_lines.extend(stmt_lines) kept_line_count += stmt_line_count stmts_kept += 1 omitted_lines = total_body_lines_count - kept_line_count # Build compressed output preserving original indentation result_parts: list[str] = [] # Signature lines (may be multi-line) if signature_lines: result_parts.extend(signature_lines) else: sig_text = _slice_code_bytes(code, node.start_byte, body_node.start_byte).rstrip() result_parts.append(sig_text) if opening_brace_line is not None: result_parts.append(opening_brace_line) if docstring_text and self.config.docstring_mode not in ( DocstringMode.NONE, DocstringMode.REMOVE, ): result_parts.append(docstring_text) if kept_lines: result_parts.extend(kept_lines) if omitted_lines > 0: result_parts.append( _make_omitted_comment( func_name, omitted_lines, indent, lang_config.comment_prefix, analysis ) ) if lang_config.uses_colon_after_signature: result_parts.append(f"{indent}pass") if closing_brace_line is not None: result_parts.append(closing_brace_line) elif after_lines: result_parts.extend(after_lines) return "\n".join(result_parts) def _compress_class_ast( self, node: Any, code: str, language: CodeLanguage, lang_config: LangConfig, body_limits: dict[str, int], analysis: _SymbolAnalysis, ) -> str: """Compress a class by individually compressing each method. Preserves class-level attributes, type annotations, and decorators while compressing method bodies individually. This ensures correct indentation for each method's omitted-body comment. """ # Use line-based extraction to preserve indentation code_lines = code.split("\n") start_row = node.start_point[0] node_lines = _get_node_lines(node, code_lines) node_text = "\n".join(node_lines) # Find the class/member container. For some languages this is not the # same node type as a function body's executable block. class_body_node_types = lang_config.class_body_node_types or lang_config.body_node_types body_node = None for child in node.children: if child.type in class_body_node_types: body_node = child break if body_node is None: return node_text # Class header (signature) — everything before the body node_start_line = node.start_point[0] body_start_line = body_node.start_point[0] sig_end = body_start_line - node_start_line if sig_end > 0: header_lines = node_lines[:sig_end] brace_line = node_lines[sig_end] if brace_line.strip().startswith("{"): header_lines = [*header_lines, brace_line] else: header_lines = [node_lines[0]] # Process each child of the class body individually body_parts: list[str] = [] processed_ranges: list[tuple[int, int]] = [] for child in body_node.children: if not child.is_named: continue # Use line-based extraction for children too child_start = child.start_point[0] child_end = child.end_point[0] if child.end_point[1] == 0 and child_end > child_start: child_end -= 1 child_text = "\n".join(code_lines[child_start : child_end + 1]) # Methods/functions inside the class — compress individually if child.type in lang_config.function_nodes: compressed = self._compress_function_ast( child, code, language, lang_config, body_limits, analysis ) body_parts.append(compressed) processed_ranges.append((child.start_byte, child.end_byte)) # Decorated methods elif lang_config.decorator_node and child.type == lang_config.decorator_node: decorator_lines = [] method_compressed = None for deco_child in child.children: if deco_child.type == "decorator": deco_start = deco_child.start_point[0] deco_end = deco_child.end_point[0] decorator_lines.append("\n".join(code_lines[deco_start : deco_end + 1])) elif deco_child.type in lang_config.function_nodes: method_compressed = self._compress_function_ast( deco_child, code, language, lang_config, body_limits, analysis ) if decorator_lines and method_compressed: body_parts.append("\n".join(decorator_lines) + "\n" + method_compressed) elif method_compressed: body_parts.append(method_compressed) else: body_parts.append(child_text) processed_ranges.append((child.start_byte, child.end_byte)) # Nested classes / containers (e.g. nested namespaces) — recurse elif child.type in lang_config.class_nodes or ( lang_config.container_node_types and child.type in lang_config.container_node_types ): compressed = self._compress_class_ast( child, code, language, lang_config, body_limits, analysis ) body_parts.append(compressed) processed_ranges.append((child.start_byte, child.end_byte)) else: # Class-level attributes, type annotations, docstrings, etc. # Keep them as-is with original indentation if child_text.strip(): body_parts.append(child_text) processed_ranges.append((child.start_byte, child.end_byte)) # Reconstruct class with proper indentation result_parts = list(header_lines) for part in body_parts: result_parts.append(part) # Handle closing brace for brace-delimited languages. The class body # node ends at the brace, while C++ class_specifier excludes the # trailing semicolon; keeping only the body node span prevents a second # semicolon from being rendered later as top-level code. body_end_line = body_node.end_point[0] body_end_rel = body_end_line - node_start_line + 1 after_lines = node_lines[body_end_rel:] if not lang_config.uses_colon_after_signature: if body_end_line != start_row: closing_line = code_lines[body_end_line] closing_text = closing_line[: body_node.end_point[1]] if _get_same_line_trailing_semicolon(node) is not None: closing_text += ";" if closing_text.strip(): result_parts.append(closing_text) elif after_lines: result_parts.extend(after_lines) return "\n".join(result_parts) def _extract_generic_structure(self, root: Any, code: str) -> CodeStructure: """Extract structure from generic/unknown code. For languages without a LangConfig, we can't reliably separate imports from other code. Just preserve everything in 'other'. """ structure = CodeStructure() structure.other = code.split("\n") return structure def _assemble_compressed( self, structure: CodeStructure, language: CodeLanguage, ) -> str: """Assemble compressed code from structure.""" parts: list[str] = [] # File header (license banners, top-of-file comments) stays on top if structure.header_code: parts.extend(structure.header_code) parts.append("") # Imports first if structure.imports: parts.extend(structure.imports) parts.append("") # Type definitions if structure.type_definitions: parts.extend(structure.type_definitions) parts.append("") # Class definitions if structure.class_definitions: parts.extend(structure.class_definitions) parts.append("") # Function signatures/definitions if structure.function_signatures: parts.extend(structure.function_signatures) parts.append("") # Top-level code (global variables, constants, if __name__, etc.) if structure.top_level_code: parts.extend(structure.top_level_code) parts.append("") # Other content (used by generic extraction) if structure.other: parts.extend(structure.other) # Remove trailing empty lines while parts and not parts[-1].strip(): parts.pop() return "\n".join(parts) def _verify_syntax(self, code: str, language: CodeLanguage) -> bool: """Verify that code is syntactically valid. Checks for both ERROR nodes (parse failures) and MISSING nodes (tokens the parser expected but didn't find). """ try: if language == CodeLanguage.PYTHON: import ast ast.parse(code) compile(code, "", "exec") parser = _get_parser(language.value) tree = parser.parse(bytes(code, "utf-8")) return not _has_syntax_issues(tree.root_node) except Exception: return False def _fallback_compress(self, code: str, original_tokens: int) -> CodeCompressionResult: """Fall back to Kompress compression.""" try: from .kompress_compressor import KompressCompressor, is_kompress_available if is_kompress_available(): compressor = KompressCompressor() result = compressor.compress(code) return CodeCompressionResult( compressed=result.compressed, original=code, original_tokens=result.original_tokens, compressed_tokens=result.compressed_tokens, compression_ratio=result.compression_ratio, language=CodeLanguage.UNKNOWN, language_confidence=0.0, # Kompress does NOT guarantee syntax validity syntax_valid=False, ) except ImportError: pass # No fallback available, return original return CodeCompressionResult( compressed=code, original=code, original_tokens=original_tokens, compressed_tokens=original_tokens, compression_ratio=1.0, language=CodeLanguage.UNKNOWN, language_confidence=0.0, syntax_valid=True, ) def _store_in_ccr( self, original: str, compressed: str, original_tokens: int, ) -> str | None: """Store original in CCR for later retrieval.""" try: from ..cache.compression_store import get_compression_store store = get_compression_store() return store.store( original, compressed, original_tokens=original_tokens, compressed_tokens=self._estimate_tokens(compressed), compression_strategy="code_aware", ) except ImportError: return None except Exception as e: logger.debug("CCR storage failed: %s", e) return None def apply( self, messages: list[dict[str, Any]], tokenizer: Tokenizer, **kwargs: Any, ) -> TransformResult: """Apply code-aware compression to messages. Handles both string content and Anthropic content block format (list of {"type": "text", "text": "..."} dicts). Args: messages: List of message dicts to transform. tokenizer: Tokenizer for accurate token counting. **kwargs: Additional arguments (e.g., 'context'). Returns: TransformResult with compressed messages and metadata. """ tokens_before = sum(tokenizer.count_text(str(m.get("content", ""))) for m in messages) context = kwargs.get("context", "") transformed_messages = [] transforms_applied: list[str] = [] warnings: list[str] = [] for message in messages: content = message.get("content", "") # Handle content blocks (Anthropic format) if isinstance(content, list): new_blocks = [] for block in content: if isinstance(block, dict) and block.get("type") == "text": text = block.get("text", "") compressed_text = self._try_compress_text( text, context, tokenizer, transforms_applied ) new_blocks.append({**block, "text": compressed_text}) else: new_blocks.append(block) transformed_messages.append({**message, "content": new_blocks}) continue # Handle string content if not content or not isinstance(content, str): transformed_messages.append(message) continue compressed_content = self._try_compress_text( content, context, tokenizer, transforms_applied ) if compressed_content != content: transformed_messages.append({**message, "content": compressed_content}) else: transformed_messages.append(message) tokens_after = sum( tokenizer.count_text(str(m.get("content", ""))) for m in transformed_messages ) if not _check_tree_sitter_available(): warnings.append( "tree-sitter not installed. Install with: pip install headroom-ai[code]" ) return TransformResult( messages=transformed_messages, tokens_before=tokens_before, tokens_after=tokens_after, transforms_applied=transforms_applied if transforms_applied else ["code_aware:noop"], warnings=warnings, ) def _try_compress_text( self, text: str, context: str, tokenizer: Tokenizer, transforms_applied: list[str], ) -> str: """Try to compress a text string if it contains code.""" from .content_detector import ContentType, detect_content_type if not text: return text detection = detect_content_type(text) if detection.content_type == ContentType.SOURCE_CODE: language = detection.metadata.get("language") result = self.compress(text, language=language, context=context, tokenizer=tokenizer) if result.compression_ratio < 0.9: transforms_applied.append( f"code_aware:{result.language.value}:{result.compression_ratio:.2f}" ) return result.compressed return text def should_apply( self, messages: list[dict[str, Any]], tokenizer: Tokenizer, **kwargs: Any, ) -> bool: """Check if code-aware compression should be applied. Returns True if: - tree-sitter is available, AND - Content contains detected source code Args: messages: Messages to check. tokenizer: Tokenizer for counting. **kwargs: Additional arguments. Returns: True if compression should be applied. """ if not _check_tree_sitter_available(): return False from .content_detector import ContentType, detect_content_type for message in messages: content = message.get("content", "") # Handle string content if content and isinstance(content, str): detection = detect_content_type(content) if detection.content_type == ContentType.SOURCE_CODE: return True # Handle content blocks elif isinstance(content, list): for block in content: if isinstance(block, dict) and block.get("type") == "text": text = block.get("text", "") if text: detection = detect_content_type(text) if detection.content_type == ContentType.SOURCE_CODE: return True return False # ========================================================================= # Module-level helper functions (stateless, used by the class) # ========================================================================= def _slice_code_bytes(code: str, start_byte: int, end_byte: int) -> str: """Extract source text using tree-sitter UTF-8 byte offsets.""" return code.encode("utf-8")[start_byte:end_byte].decode("utf-8") def _get_node_text(node: Any, code: str) -> str: """Extract text from AST node.""" return _slice_code_bytes(code, node.start_byte, node.end_byte) _COMMENT_NODE_TYPES = frozenset({"comment", "line_comment", "block_comment"}) def _get_leading_comment_text( node: Any, code: str, captured_byte_ranges: list[tuple[int, int]] ) -> str: """Collect contiguous doc-comment siblings immediately preceding a node. Doc comments are top-level siblings of the declaration they document, not children of it. Left uncaptured, they fall through to the leftover top-level sweep and get grouped separately from the declarations they document instead of staying attached. Only comments with no blank line before the node (or the next comment) are treated as attached. """ comments: list[Any] = [] sibling = getattr(node, "prev_sibling", None) anchor_row = node.start_point[0] while ( sibling is not None and sibling.type in _COMMENT_NODE_TYPES and (anchor_row - sibling.end_point[0] <= 1) ): comments.append(sibling) anchor_row = sibling.start_point[0] sibling = getattr(sibling, "prev_sibling", None) if not comments: return "" comments.reverse() captured_byte_ranges.extend((c.start_byte, c.end_byte) for c in comments) return "\n".join(_get_node_text(c, code) for c in comments) + "\n" def _get_node_lines(node: Any, code_lines: list[str]) -> list[str]: """Line-based slice of a node's source, preserving original indentation. Line-based (not byte-offset) slicing is used deliberately so leading whitespace survives for indented nested definitions (e.g. methods inside a class). But when a node shares its first line with a preceding sibling (e.g. the `export` keyword in `export function foo() {`), a naive full-line slice pulls in that sibling's text too — and callers that reconstruct the wrapper (re-adding the sibling text themselves) end up duplicating it. Trim the sibling prefix from the first line when it's not pure whitespace; keep the whole line (indentation intact) otherwise. """ start_row, start_col = node.start_point end_row = node.end_point[0] node_lines = list(code_lines[start_row : end_row + 1]) if node_lines and node_lines[0][:start_col].strip(): node_lines[0] = node_lines[0][start_col:] return node_lines def _get_same_line_trailing_semicolon(node: Any) -> Any | None: """Return a trailing semicolon sibling that belongs to this declaration.""" next_sibling = getattr(node, "next_sibling", None) if ( next_sibling is not None and next_sibling.type == ";" and next_sibling.start_point[0] == node.end_point[0] ): return next_sibling return None def _get_definition_name(node: Any) -> str | None: """Extract the name identifier from a definition AST node.""" for child in node.children: if child.type in ("identifier", "name", "type_identifier", "property_identifier"): text = child.text return text.decode("utf-8") if isinstance(text, bytes) else str(text) return None # Symbol names are ASCII identifiers; CJK relevance queries have no spaces and use # CJK/full-width punctuation, so the ASCII-only delimiter class would collapse the # whole query into one blob and never isolate an ASCII name the user asked to keep. _CONTEXT_DELIMS = re.compile(r"[\s,;:.()\[\]{}\"'，、；：。．！？（）【】「」『』《》〈〉·…—　]+") _CJK_CHARS = re.compile(r"[　-鿿가-힯＀-￯]") def _query_context_tokens(context: str) -> tuple[set[str], str, bool]: """Tokenize a relevance query for symbol-name matching (CJK-aware). Returns (word set, lowercased query, has_cjk). CJK/full-width punctuation and the ideographic space are delimiters so an ASCII symbol name wrapped in CJK is still isolated as its own token. """ if not context: return set(), "", False lowered = context.lower() words = set(_CONTEXT_DELIMS.split(lowered)) words.discard("") return words, lowered, bool(_CJK_CHARS.search(lowered)) def _symbol_in_context(name_lower: str, words: set[str], context_lower: str, has_cjk: bool) -> bool: """Whether the relevance query names this symbol. Exact token match, or a substring fallback gated by len>3 for ASCII queries (avoids spurious short-name matches) but relaxed for CJK queries -- a short ASCII name glued to CJK has no delimiter to isolate it, so exact-match can't fire and the guard would wrongly drop it. """ if not words or not name_lower: return False if name_lower in words: return True return name_lower in context_lower and (len(name_lower) > 3 or has_cjk) def _is_public_symbol(name: str, language: CodeLanguage) -> bool: """Heuristic for whether a symbol is public/exported.""" if not name: return False if language == CodeLanguage.GO: return name[0].isupper() return not name.startswith("_") def _get_body_limit( func_name: str | None, body_limits: dict[str, int], max_body_lines: int, ) -> int: """Look up the allocated body line limit for a function. Falls back to max_body_lines if no budget allocation was computed. max_body_lines always acts as a hard cap. """ if body_limits and func_name and func_name in body_limits: return min(body_limits[func_name], max_body_lines) return max_body_lines def _make_omitted_comment( func_name: str | None, omitted_count: int, indent: str, comment_prefix: str, analysis: _SymbolAnalysis | None, ) -> str: """Build omitted comment with call information from analysis.""" calls_info = "" if analysis and func_name: for key in ( func_name, *(k for k in analysis.calls if k.endswith(f".{func_name}")), ): if key in analysis.calls: called = analysis.calls[key] if called: sorted_calls = sorted(called)[:5] calls_info = "; calls: " + ", ".join(sorted_calls) if len(called) > 5: calls_info += f" +{len(called) - 5} more" break return f"{indent}{comment_prefix} [{omitted_count} lines omitted{calls_info}]" def _detect_indent(lines: list[str]) -> str: """Detect the indentation used in a list of code lines.""" for line in lines: if line.strip(): return line[: len(line) - len(line.lstrip())] return " " def _has_syntax_issues(node: Any) -> bool: """Check if AST contains ERROR or MISSING nodes.""" if node.type == "ERROR" or node.is_missing: return True for child in node.children: if _has_syntax_issues(child): return True return False def compress_code( code: str, language: str | None = None, target_rate: float = 0.2, context: str = "", ) -> str: """Convenience function for one-off code compression. Args: code: Source code to compress. language: Language hint (auto-detected if None). target_rate: Target compression rate (0.2 = keep 20%). context: Optional context for relevance. Returns: Compressed code string. Example: >>> compressed = compress_code(large_python_file) >>> print(compressed) # Valid Python code """ config = CodeCompressorConfig( target_compression_rate=target_rate, language_hint=language, ) compressor = CodeAwareCompressor(config) result = compressor.compress(code, language=language, context=context) return result.compressed