feat(chunks): aggiunge segmenter.py — token stream → Block[]
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
@@ -0,0 +1,215 @@
|
||||
from __future__ import annotations
|
||||
from markdown_it.token import Token
|
||||
from chunks.models import Block
|
||||
from chunks.config import ChunkerConfig
|
||||
|
||||
_COUNTER = 0
|
||||
|
||||
|
||||
def _reset_counter() -> None:
|
||||
global _COUNTER
|
||||
_COUNTER = 0
|
||||
|
||||
|
||||
def _make_id() -> str:
|
||||
global _COUNTER
|
||||
_COUNTER += 1
|
||||
return f"blk_{_COUNTER:04d}"
|
||||
|
||||
|
||||
def _inline_to_plain(token: Token) -> str:
|
||||
if token.children:
|
||||
parts = []
|
||||
for c in token.children:
|
||||
if c.type in ("text", "code_inline"):
|
||||
parts.append(c.content)
|
||||
elif c.type == "softbreak":
|
||||
parts.append(" ")
|
||||
return "".join(parts).strip()
|
||||
return (token.content or "").strip()
|
||||
|
||||
|
||||
def _lines_content(lines: list[str], start: int, end: int) -> str:
|
||||
return "\n".join(lines[start:end]).strip()
|
||||
|
||||
|
||||
def _current_path(stack: dict[int, str], depth: int) -> list[dict]:
|
||||
path = []
|
||||
for level in sorted(stack.keys()):
|
||||
path.append({"level": level, "text": stack[level]})
|
||||
return path[:depth]
|
||||
|
||||
|
||||
def _is_skip_heading(text: str, skip_set: set[str]) -> bool:
|
||||
t = text.lower().strip()
|
||||
return any(t == s or t.startswith(s) for s in {h.lower() for h in skip_set})
|
||||
|
||||
|
||||
def _find_close(tokens: list[Token], start: int, open_type: str, close_type: str) -> int:
|
||||
"""Ritorna indice del token close_type corrispondente a tokens[start]."""
|
||||
depth = 1
|
||||
i = start + 1
|
||||
while i < len(tokens) and depth > 0:
|
||||
if tokens[i].type == open_type:
|
||||
depth += 1
|
||||
elif tokens[i].type == close_type:
|
||||
depth -= 1
|
||||
i += 1
|
||||
return i - 1
|
||||
|
||||
|
||||
def segment(tokens: list[Token], lines: list[str], config: ChunkerConfig) -> list[Block]:
|
||||
_reset_counter()
|
||||
blocks: list[Block] = []
|
||||
heading_stack: dict[int, str] = {}
|
||||
pre_heading_done = False
|
||||
skip_mode = False
|
||||
skip_level: int | None = None
|
||||
|
||||
i = 0
|
||||
while i < len(tokens):
|
||||
tok = tokens[i]
|
||||
|
||||
# ── Heading ──────────────────────────────────────────────────────────
|
||||
if tok.type == "heading_open":
|
||||
level = int(tok.tag[1])
|
||||
inline = tokens[i + 1] if i + 1 < len(tokens) else None
|
||||
text = (inline.content if inline and inline.type == "inline" else "").strip()
|
||||
|
||||
if _is_skip_heading(text, config.skip_headings):
|
||||
skip_mode = True
|
||||
skip_level = level
|
||||
elif skip_mode and skip_level is not None and level <= skip_level:
|
||||
skip_mode = False
|
||||
skip_level = None
|
||||
|
||||
for lvl in [l for l in list(heading_stack.keys()) if l >= level]:
|
||||
del heading_stack[lvl]
|
||||
heading_stack[level] = text
|
||||
pre_heading_done = True
|
||||
i += 3 # heading_open, inline, heading_close
|
||||
continue
|
||||
|
||||
# ── Skip pre-heading ─────────────────────────────────────────────────
|
||||
if config.skip_pre_heading and not pre_heading_done:
|
||||
i += 1
|
||||
continue
|
||||
|
||||
# ── Skip section ─────────────────────────────────────────────────────
|
||||
if skip_mode:
|
||||
i += 1
|
||||
continue
|
||||
|
||||
header_path = _current_path(heading_stack, config.context_depth)
|
||||
|
||||
# ── Paragraph ────────────────────────────────────────────────────────
|
||||
if tok.type == "paragraph_open":
|
||||
inline = tokens[i + 1] if i + 1 < len(tokens) else None
|
||||
if tok.map:
|
||||
start, end = tok.map
|
||||
content = _lines_content(lines, start, end)
|
||||
plain = _inline_to_plain(inline) if inline and inline.type == "inline" else content
|
||||
blocks.append(Block(
|
||||
id=_make_id(),
|
||||
kind="paragraph",
|
||||
content=content,
|
||||
plain_text=plain,
|
||||
atomic=False,
|
||||
start_line=start,
|
||||
end_line=end,
|
||||
header_path=header_path,
|
||||
chars=len(content),
|
||||
))
|
||||
i += 3 # paragraph_open, inline, paragraph_close
|
||||
continue
|
||||
|
||||
# ── Code fence ───────────────────────────────────────────────────────
|
||||
if tok.type == "fence":
|
||||
if tok.map:
|
||||
start, end = tok.map
|
||||
content = _lines_content(lines, start, end)
|
||||
plain = f"[codice {tok.info.strip()}]" if tok.info.strip() else "[codice]"
|
||||
blocks.append(Block(
|
||||
id=_make_id(),
|
||||
kind="code",
|
||||
content=content,
|
||||
plain_text=plain,
|
||||
atomic=True,
|
||||
start_line=start,
|
||||
end_line=end,
|
||||
header_path=header_path,
|
||||
chars=len(content),
|
||||
))
|
||||
i += 1
|
||||
continue
|
||||
|
||||
# ── Container blocks ─────────────────────────────────────────────────
|
||||
_CONTAINERS = {
|
||||
"table_open": ("table_close", "table", True),
|
||||
"bullet_list_open": ("bullet_list_close", "list", True),
|
||||
"ordered_list_open": ("ordered_list_close", "list", True),
|
||||
"blockquote_open": ("blockquote_close", "blockquote", False),
|
||||
}
|
||||
if tok.type in _CONTAINERS:
|
||||
close_type, kind, atomic = _CONTAINERS[tok.type]
|
||||
atomic = atomic or kind in config.atomic_types
|
||||
close_idx = _find_close(tokens, i, tok.type, close_type)
|
||||
close_tok = tokens[close_idx]
|
||||
|
||||
start = tok.map[0] if tok.map else 0
|
||||
end = (close_tok.map[1] if close_tok.map else None) or (tok.map[1] if tok.map else start + 1)
|
||||
content = _lines_content(lines, start, end)
|
||||
blocks.append(Block(
|
||||
id=_make_id(),
|
||||
kind=kind,
|
||||
content=content,
|
||||
plain_text=content,
|
||||
atomic=atomic,
|
||||
start_line=start,
|
||||
end_line=end,
|
||||
header_path=header_path,
|
||||
chars=len(content),
|
||||
))
|
||||
i = close_idx + 1
|
||||
continue
|
||||
|
||||
# ── Thematic break ────────────────────────────────────────────────────
|
||||
if tok.type == "hr":
|
||||
start = tok.map[0] if tok.map else 0
|
||||
end = tok.map[1] if tok.map else start + 1
|
||||
blocks.append(Block(
|
||||
id=_make_id(),
|
||||
kind="thematic_break",
|
||||
content="---",
|
||||
plain_text="",
|
||||
atomic=False,
|
||||
start_line=start,
|
||||
end_line=end,
|
||||
header_path=header_path,
|
||||
chars=3,
|
||||
))
|
||||
i += 1
|
||||
continue
|
||||
|
||||
# ── HTML block ────────────────────────────────────────────────────────
|
||||
if tok.type == "html_block":
|
||||
start = tok.map[0] if tok.map else 0
|
||||
end = tok.map[1] if tok.map else start + 1
|
||||
content = tok.content.strip()
|
||||
blocks.append(Block(
|
||||
id=_make_id(),
|
||||
kind="html",
|
||||
content=content,
|
||||
plain_text="",
|
||||
atomic="html" in config.atomic_types,
|
||||
start_line=start,
|
||||
end_line=end,
|
||||
header_path=header_path,
|
||||
chars=len(content),
|
||||
))
|
||||
i += 1
|
||||
continue
|
||||
|
||||
i += 1
|
||||
|
||||
return blocks
|
||||
@@ -0,0 +1,15 @@
|
||||
import pytest
|
||||
from chunks.parser import parse
|
||||
from chunks.config import ChunkerConfig
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def cfg():
|
||||
return ChunkerConfig()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def parse_md():
|
||||
def _parse(md: str):
|
||||
return parse(md)
|
||||
return _parse
|
||||
@@ -0,0 +1,128 @@
|
||||
import pytest
|
||||
from chunks.parser import parse
|
||||
from chunks.config import ChunkerConfig
|
||||
from chunks.segmenter import segment
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def cfg():
|
||||
return ChunkerConfig()
|
||||
|
||||
|
||||
def test_paragraph_block(cfg):
|
||||
tokens, lines = parse("# Titolo\n\nParagrafo di testo.\n")
|
||||
blocks = segment(tokens, lines, cfg)
|
||||
para = next(b for b in blocks if b.kind == "paragraph")
|
||||
assert "Paragrafo di testo." in para.content
|
||||
assert para.header_path == [{"level": 1, "text": "Titolo"}]
|
||||
|
||||
|
||||
def test_code_block_is_atomic(cfg):
|
||||
tokens, lines = parse("# T\n\n```python\nprint('hello')\n```\n")
|
||||
blocks = segment(tokens, lines, cfg)
|
||||
code = next(b for b in blocks if b.kind == "code")
|
||||
assert code.atomic is True
|
||||
|
||||
|
||||
def test_table_block_is_atomic(cfg):
|
||||
tokens, lines = parse("# T\n\n| A | B |\n|---|---|\n| 1 | 2 |\n")
|
||||
blocks = segment(tokens, lines, cfg)
|
||||
table = next(b for b in blocks if b.kind == "table")
|
||||
assert table.atomic is True
|
||||
|
||||
|
||||
def test_list_block_is_atomic(cfg):
|
||||
tokens, lines = parse("# T\n\n- item 1\n- item 2\n")
|
||||
blocks = segment(tokens, lines, cfg)
|
||||
lst = next(b for b in blocks if b.kind == "list")
|
||||
assert lst.atomic is True
|
||||
|
||||
|
||||
def test_heading_stack_depth(cfg):
|
||||
md = "# H1\n\n## H2\n\n### H3\n\nTesto.\n"
|
||||
tokens, lines = parse(md)
|
||||
blocks = segment(tokens, lines, cfg)
|
||||
para = next(b for b in blocks if b.kind == "paragraph")
|
||||
assert para.header_path == [
|
||||
{"level": 1, "text": "H1"},
|
||||
{"level": 2, "text": "H2"},
|
||||
{"level": 3, "text": "H3"},
|
||||
]
|
||||
|
||||
|
||||
def test_context_depth_limits_header_path():
|
||||
cfg = ChunkerConfig(context_depth=2)
|
||||
md = "# H1\n\n## H2\n\n### H3\n\nTesto.\n"
|
||||
tokens, lines = parse(md)
|
||||
blocks = segment(tokens, lines, cfg)
|
||||
para = next(b for b in blocks if b.kind == "paragraph")
|
||||
assert len(para.header_path) == 2
|
||||
|
||||
|
||||
def test_skip_headings(cfg):
|
||||
md = "# Titolo\n\n## Indice\n\nContenuto saltato.\n\n## Sezione Reale\n\nContenuto reale.\n"
|
||||
tokens, lines = parse(md)
|
||||
blocks = segment(tokens, lines, cfg)
|
||||
contents = [b.content for b in blocks]
|
||||
assert not any("saltato" in c for c in contents)
|
||||
assert any("reale" in c for c in contents)
|
||||
|
||||
|
||||
def test_skip_pre_heading(cfg):
|
||||
md = "Testo prima del primo heading.\n\n# Titolo\n\nTesto dopo.\n"
|
||||
tokens, lines = parse(md)
|
||||
blocks = segment(tokens, lines, cfg)
|
||||
contents = [b.content for b in blocks]
|
||||
assert not any("prima del primo" in c for c in contents)
|
||||
assert any("dopo" in c for c in contents)
|
||||
|
||||
|
||||
def test_skip_pre_heading_disabled():
|
||||
cfg = ChunkerConfig(skip_pre_heading=False)
|
||||
md = "Testo prima.\n\n# Titolo\n\nTesto dopo.\n"
|
||||
tokens, lines = parse(md)
|
||||
blocks = segment(tokens, lines, cfg)
|
||||
contents = [b.content for b in blocks]
|
||||
assert any("prima" in c for c in contents)
|
||||
|
||||
|
||||
def test_block_ids_unique(cfg):
|
||||
md = "# T\n\nPara 1.\n\nPara 2.\n\nPara 3.\n"
|
||||
tokens, lines = parse(md)
|
||||
blocks = segment(tokens, lines, cfg)
|
||||
ids = [b.id for b in blocks]
|
||||
assert len(ids) == len(set(ids))
|
||||
|
||||
|
||||
def test_heading_reset_on_same_level(cfg):
|
||||
md = "# H1a\n\n## H2a\n\n# H1b\n\nTesto.\n"
|
||||
tokens, lines = parse(md)
|
||||
blocks = segment(tokens, lines, cfg)
|
||||
para = next(b for b in blocks if b.kind == "paragraph")
|
||||
assert para.header_path == [{"level": 1, "text": "H1b"}]
|
||||
|
||||
|
||||
def test_thematic_break(cfg):
|
||||
md = "# T\n\nPara.\n\n---\n\nPara 2.\n"
|
||||
tokens, lines = parse(md)
|
||||
blocks = segment(tokens, lines, cfg)
|
||||
kinds = [b.kind for b in blocks]
|
||||
assert "thematic_break" in kinds
|
||||
|
||||
|
||||
def test_blockquote(cfg):
|
||||
md = "# T\n\n> Una citazione.\n"
|
||||
tokens, lines = parse(md)
|
||||
blocks = segment(tokens, lines, cfg)
|
||||
bq = next((b for b in blocks if b.kind == "blockquote"), None)
|
||||
assert bq is not None
|
||||
assert "citazione" in bq.content
|
||||
|
||||
|
||||
def test_source_line_numbers(cfg):
|
||||
md = "# Titolo\n\nParagrafo.\n"
|
||||
tokens, lines = parse(md)
|
||||
blocks = segment(tokens, lines, cfg)
|
||||
para = next(b for b in blocks if b.kind == "paragraph")
|
||||
assert para.start_line == 2
|
||||
assert para.end_line == 3
|
||||
Reference in New Issue
Block a user