diff --git a/headroom/compression/handlers/code_handler.py b/headroom/compression/handlers/code_handler.py index 309876316..3de5a20cb 100644 --- a/headroom/compression/handlers/code_handler.py +++ b/headroom/compression/handlers/code_handler.py @@ -20,7 +20,6 @@ import logging import re import threading from dataclasses import dataclass -from enum import Enum from typing import Any from headroom.compression.handlers.base import BaseStructureHandler, HandlerResult @@ -107,19 +106,6 @@ def _ts_children(node: Any) -> list[Any]: return [node.child(i) for i in range(node.child_count())] -class CodeLanguage(Enum): - """Supported programming languages.""" - - PYTHON = "python" - JAVASCRIPT = "javascript" - TYPESCRIPT = "typescript" - GO = "go" - RUST = "rust" - JAVA = "java" - C = "c" - CPP = "cpp" - - @dataclass class CodeSpan: """A span of code with its structural role.""" @@ -236,6 +222,16 @@ _CONTAINER_BODY_TYPES: frozenset[str] = frozenset( } ) +# Language-detection markers for _detect_language +_LANGUAGE_MARKERS: dict[str, list[str]] = { + "python": ["def ", "import ", "from ", "class ", "async def"], + "javascript": ["function ", "const ", "let ", "var ", "=>"], + "typescript": ["interface ", "type ", ": string", ": number"], + "go": ["func ", "package ", "import (", "type "], + "rust": ["fn ", "let mut", "impl ", "pub fn", "use "], + "java": ["public class", "private ", "protected ", "void "], +} + # 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), @@ -608,8 +604,10 @@ class CodeStructureHandler(BaseStructureHandler): for span in spans: if span.is_structural: - for i in range(span.start, min(span.end, length)): - mask[i] = True + start = min(span.start, length) + end = min(span.end, length) + if start < end: + mask[start:end] = [True] * (end - start) return mask @@ -622,18 +620,8 @@ class CodeStructureHandler(BaseStructureHandler): Returns: Language name (lowercase). """ - # Check for language-specific markers - markers = { - "python": ["def ", "import ", "from ", "class ", "async def"], - "javascript": ["function ", "const ", "let ", "var ", "=>"], - "typescript": ["interface ", "type ", ": string", ": number"], - "go": ["func ", "package ", "import (", "type "], - "rust": ["fn ", "let mut", "impl ", "pub fn", "use "], - "java": ["public class", "private ", "protected ", "void "], - } - scores: dict[str, int] = {} - for lang, patterns in markers.items(): + for lang, patterns in _LANGUAGE_MARKERS.items(): scores[lang] = sum(1 for p in patterns if p in content) if not scores or max(scores.values()) == 0: diff --git a/headroom/compression/handlers/json_handler.py b/headroom/compression/handlers/json_handler.py index 80d4f7ec5..9c54908ba 100644 --- a/headroom/compression/handlers/json_handler.py +++ b/headroom/compression/handlers/json_handler.py @@ -183,11 +183,8 @@ class JSONStructureHandler(BaseStructureHandler): for i in range(token.start, min(token.end, len(mask))): mask[i] = True - # Convert to character tokens if needed - char_tokens = list(content) if tokens == list(content) else tokens - return HandlerResult( - mask=StructureMask(tokens=char_tokens, mask=mask), + mask=StructureMask(tokens=tokens, mask=mask), handler_name=self.name, confidence=1.0, metadata={ @@ -326,7 +323,9 @@ class JSONStructureHandler(BaseStructureHandler): i += 1 while i < n and content[i] != '"': if content[i] == "\\": - i += 2 # Skip escaped character + # Clamp: a trailing backslash at EOF must not + # step past the buffer. + i = min(i + 2, n) else: i += 1 i += 1 # Include closing quote