diff --git a/headroom/transforms/code_compressor.py b/headroom/transforms/code_compressor.py index 1023d1d12..e955bd74f 100644 --- a/headroom/transforms/code_compressor.py +++ b/headroom/transforms/code_compressor.py @@ -45,6 +45,7 @@ from __future__ import annotations import logging import re import threading +from collections.abc import Callable from dataclasses import dataclass, field from enum import Enum from typing import Any @@ -1196,6 +1197,20 @@ class CodeAwareCompressor(Transform): # Verify syntax validity (checks both ERROR and MISSING nodes) syntax_valid = self._verify_syntax(compressed, detected_lang) + original_syntax_valid = True + + if not syntax_valid and detected_lang == CodeLanguage.PYTHON: + original_syntax_valid = self._verify_syntax(code, detected_lang) + if original_syntax_valid: + compressed, structure, symbol_scores = self._compress_with_ast( + code, + detected_lang, + context, + tokenizer, + recover_invalid_python_nodes=True, + ) + compressed_tokens = self._estimate_tokens(compressed, tokenizer) + syntax_valid = self._verify_syntax(compressed, detected_lang) # If syntax invalid, return original (never serve broken code) if not syntax_valid: @@ -1294,6 +1309,7 @@ class CodeAwareCompressor(Transform): language: CodeLanguage, context: str, tokenizer: Tokenizer | None = None, + recover_invalid_python_nodes: bool = False, ) -> tuple[str, CodeStructure, dict[str, float]]: """Compress code using AST parsing with symbol importance analysis. @@ -1312,6 +1328,23 @@ class CodeAwareCompressor(Transform): parser = _get_parser(language.value) tree = parser.parse(bytes(code, "utf-8")) root = tree.root_node + candidate_validator: Callable[[Any, str], str] | None = None + + if recover_invalid_python_nodes and language == CodeLanguage.PYTHON: + code_bytes = code.encode("utf-8") + + def candidate_validator(node: Any, candidate_text: str) -> str: + original_text = _slice_code_bytes(code, node.start_byte, node.end_byte) + if candidate_text == original_text: + return candidate_text + candidate_module = ( + code_bytes[: node.start_byte] + + candidate_text.encode("utf-8") + + code_bytes[node.end_byte :] + ).decode("utf-8") + if self._verify_syntax(candidate_module, CodeLanguage.PYTHON): + return candidate_text + return original_text # Analyze symbol importance and allocate compression budget analysis = self._analyze_symbol_importance(root, code, language, context) @@ -1321,7 +1354,13 @@ class CodeAwareCompressor(Transform): lang_config = _LANG_CONFIGS.get(language) if lang_config: structure = self._extract_structure( - root, code, language, lang_config, body_limits, analysis + root, + code, + language, + lang_config, + body_limits, + analysis, + candidate_validator=candidate_validator, ) else: structure = self._extract_generic_structure(root, code) @@ -1351,6 +1390,7 @@ class CodeAwareCompressor(Transform): lang_config: LangConfig, body_limits: dict[str, int], analysis: _SymbolAnalysis, + candidate_validator: Callable[[Any, str], str] | None = None, ) -> CodeStructure: """Extract structure from AST using data-driven language config. @@ -1360,6 +1400,11 @@ class CodeAwareCompressor(Transform): structure = CodeStructure() captured_byte_ranges: list[tuple[int, int]] = [] + def _validated_candidate(node: Any, compressed: str) -> str: + if candidate_validator is None: + return compressed + return candidate_validator(node, compressed) + def visit(node: Any) -> None: node_type = node.type @@ -1396,7 +1441,11 @@ class CodeAwareCompressor(Transform): 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 + leading + + _validated_candidate( + node, + export_prefix + compressed + export_suffix, + ) ) break if not has_func_or_class: @@ -1421,16 +1470,21 @@ class CodeAwareCompressor(Transform): child, code, language, lang_config, body_limits, analysis ) if decorator_text and definition_compressed: - full_def = leading + "\n".join(decorator_text) + "\n" + definition_compressed + full_def = _validated_candidate( + node, + "\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) + structure.class_definitions.append(leading + full_def) break else: - structure.function_signatures.append(full_def) + structure.function_signatures.append(leading + full_def) elif definition_compressed: - structure.function_signatures.append(leading + definition_compressed) + structure.function_signatures.append( + leading + _validated_candidate(node, definition_compressed) + ) captured_byte_ranges.append((node.start_byte, node.end_byte)) return @@ -1440,7 +1494,9 @@ class CodeAwareCompressor(Transform): compressed = self._compress_function_ast( node, code, language, lang_config, body_limits, analysis ) - structure.function_signatures.append(leading + compressed) + structure.function_signatures.append( + leading + _validated_candidate(node, compressed) + ) captured_byte_ranges.append((node.start_byte, node.end_byte)) return @@ -1448,7 +1504,7 @@ class CodeAwareCompressor(Transform): compressed = self._compress_class_ast( node, code, language, lang_config, body_limits, analysis ) - structure.class_definitions.append(compressed) + structure.class_definitions.append(_validated_candidate(node, compressed)) captured_byte_ranges.append((node.start_byte, node.end_byte)) return @@ -1458,7 +1514,7 @@ class CodeAwareCompressor(Transform): compressed = self._compress_class_ast( node, code, language, lang_config, body_limits, analysis ) - structure.class_definitions.append(leading + compressed) + structure.class_definitions.append(leading + _validated_candidate(node, 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: diff --git a/tests/test_transforms/test_code_compressor.py b/tests/test_transforms/test_code_compressor.py index 574b438a1..33c71fac3 100644 --- a/tests/test_transforms/test_code_compressor.py +++ b/tests/test_transforms/test_code_compressor.py @@ -14,6 +14,7 @@ from unittest.mock import patch import pytest +import headroom.transforms.code_compressor as cc from headroom.transforms.code_compressor import ( CodeAwareCompressor, CodeCompressionResult, @@ -1535,6 +1536,38 @@ class TestRealASTRuns: ) ) + def _recovery_compressor(self, **overrides): + defaults = { + "min_tokens_for_compression": 1, + "max_body_lines": 2, + "enable_ccr": False, + "semantic_analysis": False, + } + defaults.update(overrides) + return CodeAwareCompressor(CodeCompressorConfig(**defaults)) + + def _python_recovery_fixture(self) -> str: + return textwrap.dedent( + """\ + from pathlib import Path + + def expand_search_roots(user_root: str) -> list[Path]: + root = Path(user_root) + candidates = [root] + for child in root.iterdir(): + candidates.append(child.resolve()) + return candidates + + def load_user_overrides(config_path: str) -> dict[str, str]: + config: dict[str, str] = {} + for line in Path(config_path).read_text().splitlines(): + if "=" in line: + key, value = line.split("=", 1) + config[key.strip()] = value.strip() + return config + """ + ) + def test_get_parser_returns_stock_node_api(self): """The parser must yield nodes with the stock tree_sitter property API that the tree-walking code relies on (``.type``/``.children``/...).""" @@ -1612,6 +1645,101 @@ class TestRealASTRuns: # Output is still valid Python. compile(result.compressed, "", "exec") + def test_python_invalid_node_falls_back_locally(self): + """One invalid Python rewrite must preserve only that definition.""" + compressor = self._recovery_compressor() + code = self._python_recovery_fixture() + original_compress_function_ast = compressor._compress_function_ast + + def _patched(node, code_text, language, lang_config, body_limits, analysis): + name = cc._get_definition_name(node) + if name == "expand_search_roots": + return "def expand_search_roots(user_root: str) -> list[Path]:\n if True\n" + if name == "load_user_overrides": + return ( + "def load_user_overrides(config_path: str) -> dict[str, str]:\n" + " config: dict[str, str] = {}\n" + " # [4 lines omitted]\n" + " return config" + ) + return original_compress_function_ast( + node, + code_text, + language, + lang_config, + body_limits, + analysis, + ) + + with patch.object(compressor, "_compress_function_ast", side_effect=_patched): + result = compressor.compress(code, language="python") + + assert result.syntax_valid is True + assert result.compression_ratio < 1.0 + assert result.compressed != code + assert "if True" not in result.compressed + assert "for child in root.iterdir():" in result.compressed + assert "# [4 lines omitted]" in result.compressed + compile(result.compressed, "", "exec") + + def test_python_recovery_does_not_block_valid_modern_syntax(self): + """Valid decorators, nested defs, and match syntax still compress.""" + compressor = self._recovery_compressor(max_body_lines=4) + code = textwrap.dedent( + """\ + from dataclasses import dataclass + + def traced(fn): + return fn + + @dataclass + class Command: + name: str + payload: dict[str, int] + + @traced + def route_command(command: Command) -> str: + def normalize(value: str) -> str: + return value.strip().lower() + + match normalize(command.name): + case "ping": + return "pong" + case "echo": + return str(command.payload) + case _: + return "unknown" + """ + ) + + result = compressor.compress(code, language="python") + + assert result.syntax_valid is True + assert result.compression_ratio < 1.0 + assert "@traced" in result.compressed + assert "class Command:" in result.compressed + assert "def route_command(command: Command) -> str:" in result.compressed + compile(result.compressed, "", "exec") + + def test_python_recovery_still_returns_original_when_all_candidates_invalid(self): + """Recovery keeps the terminal whole-file fail-safe.""" + compressor = self._recovery_compressor() + code = self._python_recovery_fixture() + + def _patched(node, _code_text, _language, _lang_config, _body_limits, _analysis): + name = cc._get_definition_name(node) or "broken" + return f"def {name}(:\n pass" + + with patch.object(compressor, "_compress_function_ast", side_effect=_patched): + result = compressor.compress(code, language="python") + + assert result.syntax_valid is True + assert result.compression_ratio == 1.0 + assert "(:\n" not in result.compressed + assert "for child in root.iterdir():" in result.compressed + assert 'key, value = line.split("=", 1)' in result.compressed + compile(result.compressed, "", "exec") + def test_get_node_text_uses_utf8_byte_offsets(self): """tree-sitter byte offsets must not be sliced as Python str indexes.""" from headroom.transforms.code_compressor import _get_node_text, _get_parser