diff --git a/src/sqlbuild/runtime/output_capture/_helpers/text.py b/src/sqlbuild/runtime/output_capture/_helpers/text.py index 6bd41a1887..69e096b9b6 100644 --- a/src/sqlbuild/runtime/output_capture/_helpers/text.py +++ b/src/sqlbuild/runtime/output_capture/_helpers/text.py @@ -4,6 +4,11 @@ import re +from sqlbuild.runtime.output_capture.constants import ( + UTF8_CONTINUATION_BYTE_MASK, + UTF8_CONTINUATION_BYTE_TAG, +) + _ANSI_ESCAPE: re.Pattern[str] = re.compile( r"(?:\x1b\][^\x07\x1b]*(?:\x07|\x1b\\)|\x1b\[[0-?]*[ -/]*[@-~])" ) @@ -22,6 +27,48 @@ def chunk_text(*, text: str, max_bytes: int) -> tuple[str, ...]: if len(text) <= max_bytes: return (text,) return tuple(text[index : index + max_bytes] for index in range(0, len(text), max_bytes)) + try: + encoded: bytes = text.encode("utf-8") + except UnicodeEncodeError: + return _chunk_text_by_character(text=text, max_bytes=max_bytes) + if len(encoded) <= max_bytes: + return (text,) + return _chunk_utf8(encoded=encoded, max_bytes=max_bytes) + + +def _chunk_utf8(*, encoded: bytes, max_bytes: int) -> tuple[str, ...]: + """Split valid UTF-8 greedily at the last code point boundary within each limit.""" + + chunks: list[str] = [] + total: int = len(encoded) + start: int = 0 + while start < total: + end: int = start + max_bytes + if end >= total: + end = total + else: + while ( + end > start + and encoded[end] & UTF8_CONTINUATION_BYTE_MASK == UTF8_CONTINUATION_BYTE_TAG + ): + end -= 1 + if end <= start: + end = start + 1 + while ( + end < total + and encoded[end] & UTF8_CONTINUATION_BYTE_MASK == UTF8_CONTINUATION_BYTE_TAG + ): + end += 1 + chunks.append(encoded[start:end].decode("utf-8")) + start = end + if not chunks: + chunks.append("") + return tuple(chunks) + + +def _chunk_text_by_character(*, text: str, max_bytes: int) -> tuple[str, ...]: + """Split text containing lone surrogates, which surrogateescape maps to single bytes.""" + if len(text) <= max_bytes and len(text.encode("utf-8", "surrogateescape")) <= max_bytes: return (text,) diff --git a/src/sqlbuild/runtime/output_capture/constants.py b/src/sqlbuild/runtime/output_capture/constants.py index aa7622cb4f..440d894096 100644 --- a/src/sqlbuild/runtime/output_capture/constants.py +++ b/src/sqlbuild/runtime/output_capture/constants.py @@ -8,6 +8,8 @@ DEFAULT_OUTPUT_QUEUE_CAPACITY: int = 1024 DEFAULT_OUTPUT_SHUTDOWN_TIMEOUT_SECONDS: float = 2.0 MIN_OUTPUT_RECORD_BYTES: int = 4 +UTF8_CONTINUATION_BYTE_MASK: int = 0xC0 +UTF8_CONTINUATION_BYTE_TAG: int = 0x80 INVOCATION_CONTEXT_ENV: str = "SQLBUILD_INVOCATION_CONTEXT_JSON" MAX_INVOCATION_CONTEXT_BYTES: int = 16 * 1024 MAX_INVOCATION_CONTEXT_DEPTH: int = 8 diff --git a/tests/unit/src/sqlbuild/runtime/output_capture/_helpers/__init__.py b/tests/unit/src/sqlbuild/runtime/output_capture/_helpers/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/unit/src/sqlbuild/runtime/output_capture/_helpers/_test_types.py b/tests/unit/src/sqlbuild/runtime/output_capture/_helpers/_test_types.py new file mode 100644 index 0000000000..240cb2de33 --- /dev/null +++ b/tests/unit/src/sqlbuild/runtime/output_capture/_helpers/_test_types.py @@ -0,0 +1,48 @@ +"""Parameterized text chunking test cases.""" + +from collections.abc import Callable +from dataclasses import dataclass + +type ChunkTextFunction = Callable[..., tuple[str, ...]] + + +@dataclass(frozen=True) +class ChunkTextTestCase: + """One exact chunk_text expectation.""" + + description: str + text: str + max_bytes: int + expected_chunks: tuple[str, ...] + + +@dataclass(frozen=True) +class ChunkTextErrorTestCase: + """One input the stored chunk contract rejects.""" + + description: str + text: str + max_bytes: int + expected_error: type[Exception] + + +@dataclass(frozen=True) +class ChunkTextOracleTestCase: + """One seeded random comparison against the reference implementation.""" + + description: str + seed: int + max_bytes: int + count: int + expected_mismatch_count: int + + +@dataclass(frozen=True) +class ChunkTextPerformanceTestCase: + """One large payload chunking time budget.""" + + description: str + unit: str + payload_bytes: int + max_bytes: int + expected_max_seconds: float diff --git a/tests/unit/src/sqlbuild/runtime/output_capture/_helpers/conftest.py b/tests/unit/src/sqlbuild/runtime/output_capture/_helpers/conftest.py new file mode 100644 index 0000000000..192ad68d57 --- /dev/null +++ b/tests/unit/src/sqlbuild/runtime/output_capture/_helpers/conftest.py @@ -0,0 +1,38 @@ +"""Shared fixtures for output capture text chunking tests.""" + +import pytest + +from tests.unit.src.sqlbuild.runtime.output_capture._helpers._test_types import ChunkTextFunction + + +def _reference_chunk_text(*, text: str, max_bytes: int) -> tuple[str, ...]: + """Copy of the original character-at-a-time chunker that defines the stored contract.""" + + if text.isascii() and max_bytes > 0: + if len(text) <= max_bytes: + return (text,) + return tuple(text[index : index + max_bytes] for index in range(0, len(text), max_bytes)) + if len(text) <= max_bytes and len(text.encode("utf-8", "surrogateescape")) <= max_bytes: + return (text,) + + chunks: list[str] = [] + current: list[str] = [] + current_bytes: int = 0 + for character in text: + character_bytes: int = len(character.encode("utf-8", "surrogateescape")) + if current and current_bytes + character_bytes > max_bytes: + chunks.append("".join(current)) + current = [] + current_bytes = 0 + current.append(character) + current_bytes += character_bytes + if current or not chunks: + chunks.append("".join(current)) + return tuple(chunks) + + +@pytest.fixture +def reference_chunk_text() -> ChunkTextFunction: + """Return the reference oracle for chunk boundaries and text.""" + + return _reference_chunk_text diff --git a/tests/unit/src/sqlbuild/runtime/output_capture/_helpers/helpers.py b/tests/unit/src/sqlbuild/runtime/output_capture/_helpers/helpers.py new file mode 100644 index 0000000000..92a805994c --- /dev/null +++ b/tests/unit/src/sqlbuild/runtime/output_capture/_helpers/helpers.py @@ -0,0 +1,43 @@ +"""Random inputs for output capture text chunking tests.""" + +import random + +_ALPHABET: tuple[str, ...] = ( + "a", + "Z", + " ", + "\n", + "é", + "ß", + "Ω", + "订", + "€", + "✓", + "📦", + "🙂", + "𝔸", + "e\u0301", + "\u0308", +) + + +def random_mixed_width_texts(*, seed: int, max_bytes: int, count: int) -> tuple[str, ...]: + """Build texts whose encoded lengths land around multiples of the chunk limit.""" + + generator: random.Random = random.Random(seed) + return tuple( + _random_mixed_width_text(generator=generator, max_bytes=max_bytes) for _ in range(count) + ) + + +def _random_mixed_width_text(*, generator: random.Random, max_bytes: int) -> str: + target_bytes: int = max( + 0, max_bytes * generator.choice((0, 1, 1, 2, 3)) + generator.randint(-5, 5) + ) + parts: list[str] = [] + size: int = 0 + while size < target_bytes: + piece: str = generator.choice(_ALPHABET) + parts.append(piece) + size += len(piece.encode("utf-8")) + return "".join(parts) diff --git a/tests/unit/src/sqlbuild/runtime/output_capture/_helpers/test_text.py b/tests/unit/src/sqlbuild/runtime/output_capture/_helpers/test_text.py new file mode 100644 index 0000000000..b32863bf60 --- /dev/null +++ b/tests/unit/src/sqlbuild/runtime/output_capture/_helpers/test_text.py @@ -0,0 +1,128 @@ +"""Tests for deterministic output capture text chunking.""" + +import time + +import pytest + +from sqlbuild.runtime.output_capture._helpers.text import chunk_text +from tests.unit.src.sqlbuild.runtime.output_capture._helpers._test_types import ( + ChunkTextErrorTestCase, + ChunkTextFunction, + ChunkTextOracleTestCase, + ChunkTextPerformanceTestCase, + ChunkTextTestCase, +) +from tests.unit.src.sqlbuild.runtime.output_capture._helpers.helpers import ( + random_mixed_width_texts, +) + + +@pytest.mark.parametrize( + "test_case", + ( + ChunkTextTestCase("two_byte_straddles_limit", "aé", 2, ("a", "é")), + ChunkTextTestCase("three_byte_straddles_limit", "ab订c", 4, ("ab", "订c")), + ChunkTextTestCase("four_byte_straddles_limit", "abc📦", 4, ("abc", "📦")), + ChunkTextTestCase("four_byte_exact_limit", "📦📦", 4, ("📦", "📦")), + ChunkTextTestCase("combining_mark_split_from_base", "ée\u0301", 4, ("ée", "\u0301")), + ChunkTextTestCase("character_wider_than_limit", "a📦b", 2, ("a", "📦", "b")), + ChunkTextTestCase("each_character_when_limit_zero", "a订", 0, ("a", "订")), + ChunkTextTestCase("empty_text_negative_limit", "", -1, ("",)), + ChunkTextTestCase("escaped_surrogate_counts_one_byte", "é\udc80a", 3, ("é\udc80", "a")), + ), + ids=lambda case: case.description, +) +def test_given_multibyte_boundaries_when_chunking_then_code_points_stay_whole( + test_case: ChunkTextTestCase, + reference_chunk_text: ChunkTextFunction, +) -> None: + result: tuple[str, ...] = chunk_text(text=test_case.text, max_bytes=test_case.max_bytes) + + assert result == test_case.expected_chunks + assert result == reference_chunk_text(text=test_case.text, max_bytes=test_case.max_bytes) + + +@pytest.mark.parametrize( + "test_case", + (ChunkTextErrorTestCase("unescapable_surrogate", "é\ud800", 2, UnicodeEncodeError),), + ids=lambda case: case.description, +) +def test_given_unescapable_surrogate_when_chunking_then_raises_like_reference( + test_case: ChunkTextErrorTestCase, + reference_chunk_text: ChunkTextFunction, +) -> None: + with pytest.raises(test_case.expected_error): + _ = reference_chunk_text(text=test_case.text, max_bytes=test_case.max_bytes) + with pytest.raises(test_case.expected_error): + _ = chunk_text(text=test_case.text, max_bytes=test_case.max_bytes) + + +@pytest.mark.parametrize( + "test_case", + [ + ChunkTextOracleTestCase( + f"limit_{max_bytes}", + seed=max_bytes, + max_bytes=max_bytes, + count=400, + expected_mismatch_count=0, + ) + for max_bytes in (1, 2, 3, 4, 5, 6, 7, 8, 13, 64, 257) + ], + ids=lambda case: case.description, +) +def test_given_random_mixed_width_text_when_chunking_then_matches_reference( + test_case: ChunkTextOracleTestCase, + reference_chunk_text: ChunkTextFunction, +) -> None: + texts: tuple[str, ...] = random_mixed_width_texts( + seed=test_case.seed, max_bytes=test_case.max_bytes, count=test_case.count + ) + + actual: tuple[tuple[str, ...], ...] = tuple( + chunk_text(text=text, max_bytes=test_case.max_bytes) for text in texts + ) + + expected: tuple[tuple[str, ...], ...] = tuple( + reference_chunk_text(text=text, max_bytes=test_case.max_bytes) for text in texts + ) + mismatch_count: int = sum( + result != reference for result, reference in zip(actual, expected, strict=True) + ) + assert mismatch_count == test_case.expected_mismatch_count + assert tuple("".join(chunks) for chunks in actual) == texts + + +@pytest.mark.parametrize( + "test_case", + ( + ChunkTextPerformanceTestCase( + "non_ascii_ten_megabytes", + unit="order ünïcode 订单 📦 e\u0301 ", + payload_bytes=10_000_000, + max_bytes=64 * 1024, + expected_max_seconds=0.1, + ), + ), + ids=lambda case: case.description, +) +def test_given_large_non_ascii_payload_when_chunking_then_completes_within_budget( + test_case: ChunkTextPerformanceTestCase, +) -> None: + unit_bytes: int = len(test_case.unit.encode("utf-8")) + text: str = test_case.unit * (test_case.payload_bytes // unit_bytes + 1) + timings: list[float] = [] + result: tuple[str, ...] = () + + for _ in range(3): + started: float = time.perf_counter() + result = chunk_text(text=text, max_bytes=test_case.max_bytes) + timings.append(time.perf_counter() - started) + + assert "".join(result) == text + assert all(len(chunk.encode("utf-8")) <= test_case.max_bytes for chunk in result) + assert min(timings) < test_case.expected_max_seconds + + +if __name__ == "__main__": + pytest.main([__file__, "-vv"])