diff --git a/headroom/transforms/code_compressor.py b/headroom/transforms/code_compressor.py index fe05e86b7..03acb1648 100644 --- a/headroom/transforms/code_compressor.py +++ b/headroom/transforms/code_compressor.py @@ -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: diff --git a/tests/test_transforms/test_code_compressor.py b/tests/test_transforms/test_code_compressor.py index 76de4e3ae..774d9b6eb 100644 --- a/tests/test_transforms/test_code_compressor.py +++ b/tests/test_transforms/test_code_compressor.py @@ -1591,6 +1591,50 @@ class TestRealASTRuns: # Output is still valid Python. 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 + + 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, "", "exec") + def test_ast_runs_for_rust_no_fallback(self): """A second supported language (Rust) also runs through real AST.""" code = (