105 lines
4.0 KiB
Python
105 lines
4.0 KiB
Python
"""ContextualChunker — structure preservation for technical corpora."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from thicket.chunker import ContextualChunker
|
|
|
|
|
|
def _filler(n: int, tag: str = "w") -> str:
|
|
return " ".join(f"{tag}{i}" for i in range(n))
|
|
|
|
|
|
def test_empty_text_yields_no_chunks():
|
|
assert ContextualChunker().chunk("T", "") == []
|
|
assert ContextualChunker().chunk("T", "\n\n\n") == []
|
|
|
|
|
|
def test_fenced_block_is_atomic_and_verbatim():
|
|
code = "\n".join([
|
|
"```python",
|
|
"def harden(host):",
|
|
" if host == 'prod':",
|
|
" sys.exit('no')",
|
|
"",
|
|
" return True",
|
|
"```",
|
|
])
|
|
chunks = ContextualChunker(chunk_size=10, overlap=0).chunk("T", code)
|
|
assert len(chunks) == 1
|
|
assert chunks[0].kind == "code"
|
|
assert chunks[0].lang == "python"
|
|
assert " if host == 'prod':" in chunks[0].text # indentation kept
|
|
assert "\n\n" in chunks[0].text # blank line kept
|
|
assert "| code:python]" in chunks[0].contextual_text
|
|
|
|
|
|
def test_code_never_merges_with_prose():
|
|
doc = "Intro prose here.\n\n```bash\nls -la\n```\n\nOutro prose here."
|
|
chunks = ContextualChunker(chunk_size=100, overlap=0).chunk("T", doc)
|
|
kinds = [c.kind for c in chunks]
|
|
assert kinds.count("code") == 1
|
|
code = next(c for c in chunks if c.kind == "code")
|
|
assert "Intro" not in code.text and "Outro" not in code.text
|
|
|
|
|
|
def test_shebang_and_comments_are_not_headings():
|
|
doc = "#!/usr/bin/env python3\nimport os\n\n#no-space-comment\n\n## Real Heading\n\nBody text."
|
|
chunks = ContextualChunker(chunk_size=100).chunk("T", doc)
|
|
assert chunks[-1].header == "Real Heading"
|
|
|
|
|
|
def test_prose_paragraph_lines_are_preserved():
|
|
doc = "# H\n\nFirst paragraph.\nSecond line of same paragraph."
|
|
chunks = ContextualChunker(chunk_size=100).chunk("T", doc)
|
|
assert "First paragraph.\nSecond line of same paragraph." in chunks[0].text
|
|
|
|
|
|
def test_paragraph_aligned_budget():
|
|
paras = [_filler(40, f"p{i}_") for i in range(6)] # 6 x 40 words
|
|
doc = "# H\n\n" + "\n\n".join(paras)
|
|
chunks = ContextualChunker(chunk_size=100, overlap=0).chunk("T", doc)
|
|
# No paragraph is ever split: each chunk is 2 paragraphs (80 words).
|
|
assert all(c.text.count("_") >= 40 for c in chunks)
|
|
assert all(len(c.text.split()) <= 100 for c in chunks)
|
|
assert len(chunks) == 3
|
|
|
|
|
|
def test_oversized_code_is_windowed_with_lang_on_every_window():
|
|
lines = [f"x_{i} = {i} # {'filler ' * 20}" for i in range(60)]
|
|
doc = "```python\n" + "\n".join(lines) + "\n```"
|
|
chunks = ContextualChunker(chunk_size=40, overlap=0).chunk("T", doc)
|
|
assert len(chunks) > 1
|
|
assert all(c.kind == "code" and c.lang == "python" for c in chunks)
|
|
# Window overlap: first line of window N+1 appeared in window N.
|
|
assert any(chunks[i + 1].text.split("\n")[0] in chunks[i].text
|
|
for i in range(len(chunks) - 1))
|
|
|
|
|
|
def test_section_boundaries_do_not_straddle():
|
|
doc = ("# A\n\n" + _filler(30, "a") + "\n\n"
|
|
"# B\n\n" + _filler(30, "b"))
|
|
chunks = ContextualChunker(chunk_size=100, overlap=50).chunk("T", doc)
|
|
assert len(chunks) == 2
|
|
assert chunks[0].header == "A" and chunks[1].header == "B"
|
|
assert "b0" not in chunks[0].text
|
|
|
|
|
|
def test_indented_block_treated_as_code():
|
|
doc = "Prose.\n\n permit root no\n retries 3\n\nMore prose."
|
|
chunks = ContextualChunker(chunk_size=100).chunk("T", doc)
|
|
code = [c for c in chunks if c.kind == "code"]
|
|
assert code and "permit root no" in code[0].text
|
|
|
|
|
|
def test_default_header_is_introduction():
|
|
chunks = ContextualChunker(chunk_size=50).chunk("T", _filler(10))
|
|
assert chunks[0].header == "Introduction"
|
|
|
|
|
|
def test_source_provenance_attaches_to_every_chunk():
|
|
doc = "# H\n\nprose\n\n```python\nx = 1\n```"
|
|
chunks = ContextualChunker(chunk_size=100).chunk(
|
|
"T", doc, source_path="sub/deploy.py")
|
|
assert {c.source for c in chunks} == {"sub/deploy.py"}
|
|
assert ContextualChunker().chunk("T", doc)[0].source is None
|