mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
fix(code): slice tree-sitter byte offsets as UTF-8 (#1332)
## Description CodeAwareCompressor was slicing Python strings with tree-sitter `start_byte` / `end_byte` offsets directly. That works for ASCII-only files, but it corrupts slices after non-ASCII source text such as CJK characters or emoji because tree-sitter offsets are UTF-8 byte offsets while Python string indexes are character offsets. This caused code-aware compression to produce invalid intermediate Python and then safely fall back to the original file, resulting in 0% compression on affected files. Closes #1319 ## 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 `_slice_code_bytes()` in `headroom/transforms/code_compressor.py` to slice source text using UTF-8 byte offsets. - Updated `_get_node_text()` to use byte-safe slicing. - Routed the other direct tree-sitter byte-offset slices through the same helper. - Added regression tests in `tests/test_transforms/test_code_compressor.py`: - `test_get_node_text_uses_utf8_byte_offsets` - `test_ast_compresses_python_after_non_ascii_source` ## 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 $ .venv/bin/python -m pytest tests/test_transforms/test_code_compressor.py 68 passed, 1 warning $ .venv/bin/python -m ruff check headroom/transforms/code_compressor.py tests/test_transforms/test_code_compressor.py All checks passed! $ .venv/bin/python -m ruff format --check headroom/transforms/code_compressor.py tests/test_transforms/test_code_compressor.py 2 files already formatted $ git diff --check # no output $ /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, tree-sitter 0.25.2, tree-sitter-language-pack 0.13.0 - Exact command / steps: On `main`, ran a local reproducer with a Python source string containing a CJK docstring before a second function; called `_get_node_text()` on the second tree-sitter function node; ran a full `CodeAwareCompressor.compress(...)` repro with non-ASCII module text before an import and a compressible function; re-ran both repros on this branch. - Observed result: Before fix, `_get_node_text()` returned the wrong slice (`'nd():\n return 2\n'` instead of `'def second():\n return 2'`) and full compression fell back to the original file with `compression_ratio: 1.0`; after fix, `_get_node_text()` returns the full expected function slice and full compression succeeds with `compression_ratio < 1.0`, `syntax_valid: True`, and does not return the original. - Not tested: Full repository test suite; live proxy/provider integrations; Windows/Linux platform-specific behavior. ## 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 - [ ] 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 was not updated because this is an internal bug fix with no user-facing API or behavior change beyond restoring intended compression. - `CHANGELOG.md` was not updated because the fix is narrow and issue-scoped; maintainers can advise if they want a changelog entry. - The fix is intentionally small and targeted: it only changes how tree-sitter byte offsets are converted back into Python source text, without changing compression heuristics or language behavior.
This commit is contained in:
Vinay Gupta
GitHub
parent
c35af858ea
commit
82384022bd
2 changed files with 54 additions and 5 deletions
|
|
@ -860,7 +860,7 @@ class CodeAwareCompressor(Transform):
|
|||
body_line_counts: dict[str, int] = {}
|
||||
for qname, node in definitions.items():
|
||||
collect_calls_in_function(node, qname)
|
||||
node_text = code[node.start_byte : node.end_byte]
|
||||
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
|
||||
|
|
@ -1266,8 +1266,8 @@ class CodeAwareCompressor(Transform):
|
|||
child, code, language, lang_config, body_limits, analysis
|
||||
)
|
||||
# Reconstruct export with compressed inner definition
|
||||
export_prefix = code[node.start_byte : child.start_byte]
|
||||
export_suffix = code[child.end_byte : node.end_byte]
|
||||
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(
|
||||
export_prefix + compressed + export_suffix
|
||||
)
|
||||
|
|
@ -1582,7 +1582,7 @@ class CodeAwareCompressor(Transform):
|
|||
if signature_lines:
|
||||
result_parts.extend(signature_lines)
|
||||
else:
|
||||
sig_text = code[node.start_byte : body_node.start_byte].rstrip()
|
||||
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:
|
||||
|
|
@ -2010,9 +2010,14 @@ class CodeAwareCompressor(Transform):
|
|||
# =========================================================================
|
||||
|
||||
|
||||
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 code[node.start_byte : node.end_byte]
|
||||
return _slice_code_bytes(code, node.start_byte, node.end_byte)
|
||||
|
||||
|
||||
def _get_same_line_trailing_semicolon(node: Any) -> Any | None:
|
||||
|
|
|
|||
|
|
@ -1591,6 +1591,50 @@ class TestRealASTRuns:
|
|||
# Output is still valid Python.
|
||||
compile(result.compressed, "<test>", "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
|
||||
|
||||
code = 'def first():\n """中文占位"""\n return 1\n\ndef second():\n return 2\n'
|
||||
root = _get_parser("python").parse(code.encode("utf-8")).root_node
|
||||
functions = [node for node in root.children if node.type == "function_definition"]
|
||||
|
||||
assert _get_node_text(functions[1], code) == "def second():\n return 2"
|
||||
|
||||
def test_ast_compresses_python_after_non_ascii_source(self):
|
||||
"""CJK/emoji before a later function must not corrupt downstream slices."""
|
||||
compressor = CodeAwareCompressor(
|
||||
CodeCompressorConfig(
|
||||
min_tokens_for_compression=1,
|
||||
max_body_lines=2,
|
||||
enable_ccr=False,
|
||||
semantic_analysis=False,
|
||||
)
|
||||
)
|
||||
code = (
|
||||
"def first():\n"
|
||||
' """中文占位 with emoji 🔥."""\n'
|
||||
" return 1\n"
|
||||
"\n"
|
||||
"def second():\n"
|
||||
" values = []\n"
|
||||
" for i in range(10):\n"
|
||||
" values.append(i)\n"
|
||||
" values.append(i * 2)\n"
|
||||
" values.append(i * 3)\n"
|
||||
" values.append(i * 4)\n"
|
||||
" return sum(values)\n"
|
||||
)
|
||||
|
||||
result = compressor.compress(code, language="python")
|
||||
|
||||
assert result.language == CodeLanguage.PYTHON
|
||||
assert result.syntax_valid is True
|
||||
assert result.compression_ratio < 1.0
|
||||
assert "def second():" in result.compressed
|
||||
assert "中文占位" in result.compressed
|
||||
compile(result.compressed, "<test>", "exec")
|
||||
|
||||
def test_ast_runs_for_rust_no_fallback(self):
|
||||
"""A second supported language (Rust) also runs through real AST."""
|
||||
code = (
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue