From c35af858eaa50b40cfcd74a80f12f79be69847c5 Mon Sep 17 00:00:00 2001 From: Vinay Gupta <58447456+aivinay@users.noreply.github.com> Date: Tue, 23 Jun 2026 15:41:36 -0400 Subject: [PATCH] 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. --- headroom/transforms/code_compressor.py | 53 +++++-- tests/test_transforms/test_code_compressor.py | 131 ++++++++++++++++++ 2 files changed, 175 insertions(+), 9 deletions(-) diff --git a/headroom/transforms/code_compressor.py b/headroom/transforms/code_compressor.py index d762cd669..fe05e86b7 100644 --- a/headroom/transforms/code_compressor.py +++ b/headroom/transforms/code_compressor.py @@ -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: diff --git a/tests/test_transforms/test_code_compressor.py b/tests/test_transforms/test_code_compressor.py index a94a4a235..76de4e3ae 100644 --- a/tests/test_transforms/test_code_compressor.py +++ b/tests/test_transforms/test_code_compressor.py @@ -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(