From 2d2d478a3f5720b4f6f4caabddafd3d16e981efe Mon Sep 17 00:00:00 2001
From: Davide Grilli <davide.grilli@outlook.com>
Date: Mon, 8 Jun 2026 16:25:49 +0200
Subject: [PATCH] =?UTF-8?q?feat(chunks):=20aggiunge=20segmenter.py=20?=
 =?UTF-8?q?=E2=80=94=20token=20stream=20=E2=86=92=20Block[]?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
---
 chunks/segmenter.py            | 215 +++++++++++++++++++++++++++++++++
 tests/chunks/conftest.py       |  15 +++
 tests/chunks/test_segmenter.py | 128 ++++++++++++++++++++
 3 files changed, 358 insertions(+)
 create mode 100644 chunks/segmenter.py
 create mode 100644 tests/chunks/conftest.py
 create mode 100644 tests/chunks/test_segmenter.py

diff --git a/chunks/segmenter.py b/chunks/segmenter.py
new file mode 100644
index 0000000..e95ab8f
--- /dev/null
+++ b/chunks/segmenter.py
@@ -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
diff --git a/tests/chunks/conftest.py b/tests/chunks/conftest.py
new file mode 100644
index 0000000..12d7cf2
--- /dev/null
+++ b/tests/chunks/conftest.py
@@ -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
diff --git a/tests/chunks/test_segmenter.py b/tests/chunks/test_segmenter.py
new file mode 100644
index 0000000..81b87fa
--- /dev/null
+++ b/tests/chunks/test_segmenter.py
@@ -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