mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-10 14:27:00 -04:00
fix(compression): keep container bodies compressible in code handler (#890)
## Description Two bugs in `CodeStructureHandler`'s tree-sitter path. (1) Container nodes (class/impl/trait/decorated definitions) were marked structural over their full span and `_spans_to_mask` never un-marks, so every method body inside a class was preserved and compression silently no-opped at confidence 0.95. (2) Discovered while testing: `tree-sitter-language-pack >= 1.0` switched to a Rust binding (methods, not attributes; `parse(str)`), so the handler raised `TypeError` on every call and silently fell back to regex. Closes # <!-- compression-handler review --> ## Type of Change - [x] Bug fix (non-breaking change that fixes an issue) - [ ] New feature (non-breaking change that adds functionality) - [ ] Breaking change (fix or feature that would cause existing functionality to change) - [ ] Documentation update - [ ] Performance improvement - [ ] Code refactoring (no functional changes) ## Changes Made - `headroom/compression/handlers/code_handler.py`: containers emit a signature-only span (start to body start); recursion gives nested functions their own signature/body split; decorated definitions emit no whole-node span. - `headroom/compression/handlers/code_handler.py`: small compat shim supporting both the classic attribute API and the new Rust-binding method API. - `tests/test_compression/test_code_handler.py`: new file (the handler had zero dedicated tests) covering class/decorated/impl body compressibility, regex fallback, and a preservation-ratio bound. ## Testing - [x] Unit tests pass (`pytest`) - [x] Linting passes (`ruff check .`) - [x] Type checking passes (`mypy headroom`) - [x] New tests added for new functionality - [x] Manual testing performed ### Test Output ```text $ pytest tests/test_compression/ -q 92 passed, 8 skipped ``` ## Real Behavior Proof - Environment: local macOS, Python 3.11, tree-sitter-language-pack 1.8.1 installed, branch `fix/code-container-bodies`. - Exact command / steps: `pytest tests/test_compression/ -q`. - Observed result: Class method bodies are now compressible (preservation ratio drops from ~1.0 to roughly the signature fraction); the tree-sitter path runs instead of falling back to regex. - Not tested: End-to-end through the live proxy pipeline. ## Review Readiness - [x] I have performed a self-review - [x] This PR is ready for human review ## Checklist - [x] My code follows the project's style guidelines - [x] I have performed a self-review of my code - [x] I have commented my code, particularly in hard-to-understand areas - [ ] I have made corresponding changes to the documentation - [x] My changes generate no new warnings - [x] I have added tests that prove my fix is effective or that my feature works - [x] New and existing unit tests pass locally with my changes - [ ] I have updated the CHANGELOG.md if applicable ## Screenshots (if applicable) N/A — library change. See Test Output. ## Additional Notes PR 3 of 7; branched fresh from main (independent of #887/#889). 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
Ashish
GitHub
parent
d6f0f0f642
commit
16ed73bca6
2 changed files with 272 additions and 21 deletions
|
|
@ -63,6 +63,50 @@ def _get_parser(language: str) -> Any:
|
|||
return _tree_sitter_parsers[language]
|
||||
|
||||
|
||||
# tree-sitter API compatibility. tree-sitter-language-pack switched to a
|
||||
# Rust binding (>=1.0) where node accessors are METHODS (kind(),
|
||||
# start_byte(), child(i)) and parse() takes str; the classic pybind API
|
||||
# uses attributes (.type, .start_byte, .children) and parse(bytes).
|
||||
# Without this shim the tree-sitter path raises TypeError on modern
|
||||
# installs and silently falls back to regex.
|
||||
|
||||
|
||||
def _ts_parse(parser: Any, content: str) -> Any:
|
||||
try:
|
||||
return parser.parse(content.encode("utf-8"))
|
||||
except TypeError:
|
||||
return parser.parse(content)
|
||||
|
||||
|
||||
def _ts_root(tree: Any) -> Any:
|
||||
root = tree.root_node
|
||||
return root() if callable(root) else root
|
||||
|
||||
|
||||
def _ts_kind(node: Any) -> str:
|
||||
kind = getattr(node, "type", None)
|
||||
if isinstance(kind, str):
|
||||
return kind
|
||||
return str(node.kind())
|
||||
|
||||
|
||||
def _ts_start_byte(node: Any) -> int:
|
||||
start = node.start_byte
|
||||
return int(start()) if callable(start) else int(start)
|
||||
|
||||
|
||||
def _ts_end_byte(node: Any) -> int:
|
||||
end = node.end_byte
|
||||
return int(end()) if callable(end) else int(end)
|
||||
|
||||
|
||||
def _ts_children(node: Any) -> list[Any]:
|
||||
children = getattr(node, "children", None)
|
||||
if children is not None and not callable(children):
|
||||
return list(children)
|
||||
return [node.child(i) for i in range(node.child_count())]
|
||||
|
||||
|
||||
class CodeLanguage(Enum):
|
||||
"""Supported programming languages."""
|
||||
|
||||
|
|
@ -175,6 +219,23 @@ _SIGNATURE_PATTERNS: dict[str, list[re.Pattern[str]]] = {
|
|||
],
|
||||
}
|
||||
|
||||
# Body child node types for container definitions (classes, impls,
|
||||
# traits). A container's span up to its body is structural (the
|
||||
# signature); the body itself is NOT marked — recursion into the body
|
||||
# emits signature spans for nested functions/methods, leaving their
|
||||
# bodies compressible.
|
||||
_CONTAINER_BODY_TYPES: frozenset[str] = frozenset(
|
||||
{
|
||||
"block", # python class body
|
||||
"statement_block", # js/ts
|
||||
"compound_statement", # c/cpp
|
||||
"class_body", # js/ts/java class body
|
||||
"interface_body", # java/ts interface body
|
||||
"declaration_list", # rust impl/trait body
|
||||
"enum_body", # java enum body
|
||||
}
|
||||
)
|
||||
|
||||
# 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),
|
||||
|
|
@ -301,15 +362,16 @@ class CodeStructureHandler(BaseStructureHandler):
|
|||
HandlerResult with mask.
|
||||
"""
|
||||
parser = _get_parser(language)
|
||||
tree = parser.parse(content.encode("utf-8"))
|
||||
tree = _ts_parse(parser, content)
|
||||
|
||||
# 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
|
||||
node_type = _ts_kind(node)
|
||||
structural_types = _STRUCTURAL_NODE_TYPES.get(language, set())
|
||||
children = _ts_children(node)
|
||||
|
||||
# Check if this is a structural node type
|
||||
if node_type in structural_types:
|
||||
|
|
@ -317,8 +379,8 @@ class CodeStructureHandler(BaseStructureHandler):
|
|||
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"):
|
||||
for child in children:
|
||||
if _ts_kind(child) in ("block", "statement_block", "compound_statement"):
|
||||
body_node = child
|
||||
break
|
||||
|
||||
|
|
@ -326,8 +388,8 @@ class CodeStructureHandler(BaseStructureHandler):
|
|||
# Signature is from start to body start
|
||||
spans.append(
|
||||
CodeSpan(
|
||||
start=node.start_byte,
|
||||
end=body_node.start_byte,
|
||||
start=_ts_start_byte(node),
|
||||
end=_ts_start_byte(body_node),
|
||||
role="signature",
|
||||
is_structural=True,
|
||||
)
|
||||
|
|
@ -335,8 +397,8 @@ class CodeStructureHandler(BaseStructureHandler):
|
|||
# Body is compressible
|
||||
spans.append(
|
||||
CodeSpan(
|
||||
start=body_node.start_byte,
|
||||
end=body_node.end_byte,
|
||||
start=_ts_start_byte(body_node),
|
||||
end=_ts_end_byte(body_node),
|
||||
role="body",
|
||||
is_structural=False,
|
||||
)
|
||||
|
|
@ -345,37 +407,75 @@ class CodeStructureHandler(BaseStructureHandler):
|
|||
# No body found, preserve whole thing
|
||||
spans.append(
|
||||
CodeSpan(
|
||||
start=node.start_byte,
|
||||
end=node.end_byte,
|
||||
start=_ts_start_byte(node),
|
||||
end=_ts_end_byte(node),
|
||||
role=node_type,
|
||||
is_structural=True,
|
||||
)
|
||||
)
|
||||
elif node_type == "decorated_definition":
|
||||
# Wrapper around decorator(s) + definition. Emit no
|
||||
# span: recursion marks the decorators and gives the
|
||||
# inner function its signature/body split. A whole-
|
||||
# node span here would preserve the function body.
|
||||
pass
|
||||
else:
|
||||
# Non-function structural nodes
|
||||
spans.append(
|
||||
CodeSpan(
|
||||
start=node.start_byte,
|
||||
end=node.end_byte,
|
||||
role=node_type,
|
||||
is_structural=True,
|
||||
# Container definitions (class, impl, trait): the
|
||||
# signature runs to the body start; the body is NOT
|
||||
# marked, so nested function bodies stay compressible
|
||||
# (recursion emits their signature spans). Leaf
|
||||
# declarations (imports, type aliases, structs) have
|
||||
# no such body child and are preserved whole.
|
||||
body_node = None
|
||||
for child in children:
|
||||
if _ts_kind(child) in _CONTAINER_BODY_TYPES:
|
||||
body_node = child
|
||||
break
|
||||
|
||||
if body_node is not None:
|
||||
spans.append(
|
||||
CodeSpan(
|
||||
start=_ts_start_byte(node),
|
||||
end=_ts_start_byte(body_node),
|
||||
role="signature",
|
||||
is_structural=True,
|
||||
)
|
||||
)
|
||||
else:
|
||||
spans.append(
|
||||
CodeSpan(
|
||||
start=_ts_start_byte(node),
|
||||
end=_ts_end_byte(node),
|
||||
role=node_type,
|
||||
is_structural=True,
|
||||
)
|
||||
)
|
||||
elif node_type == "decorator":
|
||||
# Decorators are structural (preserved) on their own so
|
||||
# the decorated_definition wrapper doesn't need a span.
|
||||
spans.append(
|
||||
CodeSpan(
|
||||
start=_ts_start_byte(node),
|
||||
end=_ts_end_byte(node),
|
||||
role="decorator",
|
||||
is_structural=True,
|
||||
)
|
||||
)
|
||||
elif node_type == "comment" and self.preserve_comments:
|
||||
spans.append(
|
||||
CodeSpan(
|
||||
start=node.start_byte,
|
||||
end=node.end_byte,
|
||||
start=_ts_start_byte(node),
|
||||
end=_ts_end_byte(node),
|
||||
role="comment",
|
||||
is_structural=True,
|
||||
)
|
||||
)
|
||||
|
||||
# Recurse into children
|
||||
for child in node.children:
|
||||
for child in children:
|
||||
visit_node(child, depth + 1)
|
||||
|
||||
visit_node(tree.root_node)
|
||||
visit_node(_ts_root(tree))
|
||||
|
||||
# Build mask from spans
|
||||
mask = self._spans_to_mask(spans, len(content))
|
||||
|
|
|
|||
151
tests/test_compression/test_code_handler.py
Normal file
151
tests/test_compression/test_code_handler.py
Normal file
|
|
@ -0,0 +1,151 @@
|
|||
"""Tests for code structure handler."""
|
||||
|
||||
import pytest
|
||||
|
||||
from headroom.compression.handlers.code_handler import (
|
||||
CodeStructureHandler,
|
||||
is_tree_sitter_available,
|
||||
)
|
||||
|
||||
requires_tree_sitter = pytest.mark.skipif(
|
||||
not is_tree_sitter_available(),
|
||||
reason="tree-sitter-language-pack not installed",
|
||||
)
|
||||
|
||||
|
||||
class TestCanHandle:
|
||||
@pytest.fixture
|
||||
def handler(self):
|
||||
return CodeStructureHandler()
|
||||
|
||||
def test_detects_python(self, handler):
|
||||
assert handler.can_handle("def foo():\n pass\n") is True
|
||||
|
||||
def test_detects_javascript(self, handler):
|
||||
assert handler.can_handle("function foo() { return 1; }") is True
|
||||
|
||||
def test_rejects_prose(self, handler):
|
||||
assert handler.can_handle("This is a plain sentence.") is False
|
||||
|
||||
|
||||
class TestRegexFallback:
|
||||
"""Regex path runs regardless of tree-sitter availability."""
|
||||
|
||||
@pytest.fixture
|
||||
def handler(self):
|
||||
return CodeStructureHandler(use_tree_sitter=False)
|
||||
|
||||
def test_python_signature_preserved_body_compressible(self, handler):
|
||||
code = "def hello(name: str) -> str:\n message = name\n return message\n"
|
||||
result = handler.get_mask(code, language="python")
|
||||
|
||||
assert result.metadata["parser"] == "regex"
|
||||
sig = "def hello(name: str) -> str:"
|
||||
start = code.index(sig)
|
||||
assert all(result.mask.mask[i] for i in range(start, start + len(sig)))
|
||||
|
||||
body_char = code.index("message = name")
|
||||
assert result.mask.mask[body_char] is False
|
||||
|
||||
def test_python_import_preserved(self, handler):
|
||||
code = "import os\n\nx = 1\n"
|
||||
result = handler.get_mask(code, language="python")
|
||||
assert all(result.mask.mask[i] for i in range(len("import os")))
|
||||
|
||||
|
||||
@requires_tree_sitter
|
||||
class TestTreeSitterContainers:
|
||||
"""Container bodies must stay compressible (signature-only spans).
|
||||
|
||||
Regression: class_definition / decorated_definition / impl_item were
|
||||
marked structural over their FULL span, so every method body inside a
|
||||
class (i.e. most real code) was preserved and compression no-opped at
|
||||
confidence 0.95.
|
||||
"""
|
||||
|
||||
@pytest.fixture
|
||||
def handler(self):
|
||||
return CodeStructureHandler()
|
||||
|
||||
def test_class_method_bodies_compressible(self, handler):
|
||||
code = (
|
||||
"class Foo:\n"
|
||||
" def method_a(self):\n"
|
||||
" body_line_a = 1\n"
|
||||
" return body_line_a\n"
|
||||
"\n"
|
||||
" def method_b(self):\n"
|
||||
" body_line_b = 2\n"
|
||||
" return body_line_b\n"
|
||||
)
|
||||
result = handler.get_mask(code, language="python")
|
||||
assert result.metadata["parser"] == "tree-sitter"
|
||||
|
||||
# Class signature and method signatures preserved
|
||||
assert all(result.mask.mask[i] for i in range(len("class Foo:")))
|
||||
sig = "def method_a(self):"
|
||||
start = code.index(sig)
|
||||
assert all(result.mask.mask[i] for i in range(start, start + len(sig)))
|
||||
|
||||
# Method bodies compressible
|
||||
for body in ("body_line_a = 1", "body_line_b = 2"):
|
||||
start = code.index(body)
|
||||
assert not any(result.mask.mask[i] for i in range(start, start + len(body))), (
|
||||
f"method body {body!r} must be compressible"
|
||||
)
|
||||
|
||||
def test_decorated_function_body_compressible(self, handler):
|
||||
code = "@decorator\ndef decorated():\n body_line = 4\n return body_line\n"
|
||||
result = handler.get_mask(code, language="python")
|
||||
|
||||
# Decorator and signature preserved
|
||||
assert all(result.mask.mask[i] for i in range(len("@decorator")))
|
||||
sig = "def decorated():"
|
||||
start = code.index(sig)
|
||||
assert all(result.mask.mask[i] for i in range(start, start + len(sig)))
|
||||
|
||||
# Body compressible
|
||||
start = code.index("body_line = 4")
|
||||
assert not any(result.mask.mask[i] for i in range(start, start + len("body_line = 4"))), (
|
||||
"decorated function body must be compressible"
|
||||
)
|
||||
|
||||
def test_module_function_body_compressible(self, handler):
|
||||
code = "def standalone():\n body_line = 3\n return body_line\n"
|
||||
result = handler.get_mask(code, language="python")
|
||||
|
||||
start = code.index("body_line = 3")
|
||||
assert not any(result.mask.mask[i] for i in range(start, start + len("body_line = 3")))
|
||||
|
||||
def test_rust_impl_method_bodies_compressible(self, handler):
|
||||
code = (
|
||||
"struct Foo { x: i32 }\n"
|
||||
"impl Foo {\n"
|
||||
" fn method(&self) -> i32 {\n"
|
||||
" let body_line = 5;\n"
|
||||
" body_line\n"
|
||||
" }\n"
|
||||
"}\n"
|
||||
)
|
||||
result = handler.get_mask(code, language="rust")
|
||||
|
||||
# impl signature preserved
|
||||
start = code.index("impl Foo")
|
||||
assert all(result.mask.mask[i] for i in range(start, start + len("impl Foo")))
|
||||
|
||||
# method body compressible
|
||||
start = code.index("let body_line = 5;")
|
||||
assert not any(
|
||||
result.mask.mask[i] for i in range(start, start + len("let body_line = 5;"))
|
||||
), "impl method body must be compressible"
|
||||
|
||||
def test_preservation_ratio_sane_for_class_code(self, handler):
|
||||
"""A class with substantial method bodies should NOT preserve
|
||||
everything — the whole point of the handler."""
|
||||
body = "\n".join(f" line_{i} = {i}" for i in range(20))
|
||||
code = f"class Big:\n def method(self):\n{body}\n return 0\n"
|
||||
result = handler.get_mask(code, language="python")
|
||||
assert result.preservation_ratio < 0.5, (
|
||||
f"class code preserved {result.preservation_ratio:.0%} — "
|
||||
"container bodies are leaking into the structural mask"
|
||||
)
|
||||
Loading…
Add table
Add a link
Reference in a new issue