mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
fix(code): compress class member containers (#1334)
## Description CodeAwareCompressor used the same `body_node_types` config to find both executable function bodies and class/impl member containers. That works when those AST nodes happen to match, but it misses member containers such as Java `class_body`, C++ `field_declaration_list`, and Rust `declaration_list`, so class methods were returned essentially uncompressed. This adds an optional `class_body_node_types` override for class/impl member containers and uses it only in class compression. It also skips anonymous punctuation tokens while reconstructing class bodies and keeps same-line C++ class semicolons attached to the compressed class declaration. Closes #1318 ## 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 - Added `LangConfig.class_body_node_types` for languages whose class/impl member container differs from executable method-body nodes. - Configured class member containers for JavaScript, TypeScript, Java, C++, and Rust. - Updated `_compress_class_ast` to use class-member containers, skip anonymous punctuation children, and preserve C++ `};` output without creating stray top-level semicolons. - Added regression coverage proving class/impl methods compress for JavaScript, TypeScript, Java, C++, and Rust. ## 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 $ PYTHONPATH=. /tmp/headroom-1319-venv/bin/python -m pytest tests/test_transforms/test_code_compressor.py -q collected 71 items tests/test_transforms/test_code_compressor.py .......................... [ 36%] ............................................. [100%] 71 passed, 1 warning in 0.36s $ /tmp/headroom-1319-venv/bin/python -m ruff check . All checks passed! $ /tmp/headroom-1319-venv/bin/python -m ruff format --check . 965 files already formatted $ PYTHONPATH=. /tmp/headroom-1319-venv/bin/python -m mypy headroom headroom/proxy/server.py:1186: note: By default the bodies of untyped functions are not checked, consider using --check-untyped-defs [annotation-unchecked] headroom/proxy/server.py:1257: note: By default the bodies of untyped functions are not checked, consider using --check-untyped-defs [annotation-unchecked] headroom/proxy/server.py:1261: note: By default the bodies of untyped functions are not checked, consider using --check-untyped-defs [annotation-unchecked] pyproject.toml: note: unused section(s): module = ['mlx.*'] Success: no issues found in 394 source files ``` ## Real Behavior Proof - Environment: macOS arm64, Python 3.14.5, branch `fix-code-compressor-class-members`, tree-sitter grammar pack installed in `/tmp/headroom-1319-venv`, repo imported with `PYTHONPATH=.`. - Exact command / steps: Reproduced class-method compression with `CodeAwareCompressor(CodeCompressorConfig(enable_ccr=False, min_tokens_for_compression=1, max_body_lines=1))` for Java/C++/Rust before the fix, then reran the pytest/ruff/mypy commands listed above after the patch. - Observed result: Java/C++/Rust class methods now compress below 1.0 while `syntax_valid` remains true; C++ output preserves `};`; regression coverage also verifies JavaScript/TypeScript class member containers. - Not tested: Full repository pytest suite; local `uv run` editable builds are blocked on this machine by native C++ header failures in optional/native dependencies (`hnswlib` / Rust `esaxx-rs`), so validation used a lightweight venv with `PYTHONPATH=.`. ## 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 ## Additional Notes Documentation and CHANGELOG updates are not applicable for this narrow bug fix. The pytest warning shown above is from running without `pytest-asyncio` in the lightweight verification venv (`asyncio_mode` config is unknown there); it is unrelated to this change.
This commit is contained in:
Vinay Gupta
GitHub
parent
cbd361de2a
commit
c35af858ea
2 changed files with 175 additions and 9 deletions
|
|
@ -257,6 +257,8 @@ class LangConfig:
|
|||
|
||||
# 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
|
||||
|
||||
|
||||
_LANG_CONFIGS: dict[CodeLanguage, LangConfig] = {
|
||||
|
|
@ -283,6 +285,7 @@ _LANG_CONFIGS: dict[CodeLanguage, LangConfig] = {
|
|||
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"}),
|
||||
|
|
@ -294,6 +297,7 @@ _LANG_CONFIGS: dict[CodeLanguage, LangConfig] = {
|
|||
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"}),
|
||||
|
|
@ -317,6 +321,7 @@ _LANG_CONFIGS: dict[CodeLanguage, LangConfig] = {
|
|||
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"}),
|
||||
|
|
@ -329,6 +334,7 @@ _LANG_CONFIGS: dict[CodeLanguage, LangConfig] = {
|
|||
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"}),
|
||||
|
|
@ -351,6 +357,7 @@ _LANG_CONFIGS: dict[CodeLanguage, LangConfig] = {
|
|||
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"}),
|
||||
|
|
@ -1315,6 +1322,11 @@ class CodeAwareCompressor(Transform):
|
|||
)
|
||||
structure.class_definitions.append(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
|
||||
|
|
@ -1623,10 +1635,12 @@ class CodeAwareCompressor(Transform):
|
|||
node_lines = code_lines[start_row : end_row + 1]
|
||||
node_text = "\n".join(node_lines)
|
||||
|
||||
# Find the body node
|
||||
# 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 lang_config.body_node_types:
|
||||
if child.type in class_body_node_types:
|
||||
body_node = child
|
||||
break
|
||||
|
||||
|
|
@ -1644,6 +1658,9 @@ class CodeAwareCompressor(Transform):
|
|||
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]
|
||||
|
|
@ -1695,17 +1712,23 @@ class CodeAwareCompressor(Transform):
|
|||
for part in body_parts:
|
||||
result_parts.append(part)
|
||||
|
||||
# Handle closing brace for brace-delimited languages
|
||||
# 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 after_lines:
|
||||
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)
|
||||
elif not lang_config.uses_colon_after_signature:
|
||||
# Ensure closing brace
|
||||
last_body_line = node_lines[-1] if node_lines else ""
|
||||
if last_body_line.strip() == "}":
|
||||
result_parts.append(last_body_line)
|
||||
|
||||
return "\n".join(result_parts)
|
||||
|
||||
|
|
@ -1992,6 +2015,18 @@ def _get_node_text(node: Any, code: str) -> str:
|
|||
return code[node.start_byte : node.end_byte]
|
||||
|
||||
|
||||
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:
|
||||
|
|
|
|||
|
|
@ -770,6 +770,137 @@ class TestTreeSitterIntegration:
|
|||
assert result.language == CodeLanguage.GO
|
||||
assert result.compressed # Some output is produced
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
(
|
||||
"language",
|
||||
"code",
|
||||
"expected_signature",
|
||||
"expected_omitted_lines",
|
||||
"expected_removed_line",
|
||||
"expected_closing",
|
||||
),
|
||||
[
|
||||
(
|
||||
"javascript",
|
||||
(
|
||||
"class Calc {\n"
|
||||
" compute(x) {\n"
|
||||
" let a = x + 1;\n"
|
||||
" let b = a * 2;\n"
|
||||
" let c = b - 3;\n"
|
||||
" return c;\n"
|
||||
" }\n"
|
||||
"}\n"
|
||||
),
|
||||
"compute(x) {",
|
||||
3,
|
||||
"return c;",
|
||||
"}\n}",
|
||||
),
|
||||
(
|
||||
"typescript",
|
||||
(
|
||||
"class Calc {\n"
|
||||
" compute(x: number): number {\n"
|
||||
" let a = x + 1;\n"
|
||||
" let b = a * 2;\n"
|
||||
" let c = b - 3;\n"
|
||||
" return c;\n"
|
||||
" }\n"
|
||||
"}\n"
|
||||
),
|
||||
"compute(x: number): number {",
|
||||
3,
|
||||
"return c;",
|
||||
"}\n}",
|
||||
),
|
||||
(
|
||||
"java",
|
||||
(
|
||||
"public class Calc {\n"
|
||||
" public int compute(int x) {\n"
|
||||
" int a = x + 1;\n"
|
||||
" int b = a * 2;\n"
|
||||
" int c = b - 3;\n"
|
||||
" int d = c / 4;\n"
|
||||
" int e = d + 5;\n"
|
||||
" return e;\n"
|
||||
" }\n"
|
||||
"}\n"
|
||||
),
|
||||
"public int compute(int x) {",
|
||||
5,
|
||||
"return e;",
|
||||
"}\n}",
|
||||
),
|
||||
(
|
||||
"cpp",
|
||||
(
|
||||
"class Calc {\n"
|
||||
"public:\n"
|
||||
" int compute(int x) {\n"
|
||||
" int a = x + 1;\n"
|
||||
" int b = a * 2;\n"
|
||||
" int c = b - 3;\n"
|
||||
" int d = c / 4;\n"
|
||||
" int e = d + 5;\n"
|
||||
" return e;\n"
|
||||
" }\n"
|
||||
"};\n"
|
||||
),
|
||||
"int compute(int x) {",
|
||||
5,
|
||||
"return e;",
|
||||
"};",
|
||||
),
|
||||
(
|
||||
"rust",
|
||||
(
|
||||
"impl Calc {\n"
|
||||
" pub fn compute(&self, x: i32) -> i32 {\n"
|
||||
" let a = x + 1;\n"
|
||||
" let b = a * 2;\n"
|
||||
" let c = b - 3;\n"
|
||||
" let d = c / 4;\n"
|
||||
" let e = d + 5;\n"
|
||||
" e\n"
|
||||
" }\n"
|
||||
"}\n"
|
||||
),
|
||||
"pub fn compute(&self, x: i32) -> i32 {",
|
||||
5,
|
||||
" e\n",
|
||||
"}\n}",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_compresses_methods_inside_class_member_containers(
|
||||
self,
|
||||
language,
|
||||
code,
|
||||
expected_signature,
|
||||
expected_omitted_lines,
|
||||
expected_removed_line,
|
||||
expected_closing,
|
||||
):
|
||||
"""Class/impl member containers are distinct from executable method bodies."""
|
||||
config = CodeCompressorConfig(
|
||||
min_tokens_for_compression=1,
|
||||
max_body_lines=1,
|
||||
enable_ccr=False,
|
||||
)
|
||||
compressor = CodeAwareCompressor(config)
|
||||
|
||||
result = compressor.compress(code, language=language)
|
||||
|
||||
assert result.language == CodeLanguage(language)
|
||||
assert result.syntax_valid is True
|
||||
assert result.compression_ratio < 1.0
|
||||
assert expected_signature in result.compressed
|
||||
assert f"// [{expected_omitted_lines} lines omitted]" in result.compressed
|
||||
assert expected_removed_line not in result.compressed
|
||||
assert result.compressed.endswith(expected_closing)
|
||||
|
||||
def test_imports_preserved(self):
|
||||
"""Imports are preserved in compressed output."""
|
||||
config = CodeCompressorConfig(
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue