Coverage for src/lilbee/data/extract/code_chunker.py: 100%
84 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-08-14 11:46 +0000
1"""Code chunking via tree-sitter.
3Consumes the parser's size-bounded chunks and prepends a header naming the
4relative source and the symbols each chunk defines. Falls back to plain text
5chunking when the language is unsupported or parsing fails.
6"""
8import logging
9from dataclasses import dataclass
10from pathlib import Path
11from typing import Any
13from tree_sitter_language_pack import (
14 PackConfig,
15 ProcessConfig,
16 detect_language,
17 has_language,
18 init,
19 process,
20)
22from lilbee.core.config import cfg
24from .chunk import CHARS_PER_TOKEN, chunk_text
26log = logging.getLogger(__name__)
29@dataclass
30class CodeChunk:
31 """A chunk of source code with line location metadata."""
33 chunk: str
34 line_start: int
35 line_end: int
36 chunk_index: int
39def _detect_language(file_path: Path) -> str | None:
40 """Detect language from file path using tree-sitter-language-pack."""
41 result: str | None = detect_language(str(file_path))
42 return result
45def _ensure_language(lang: str) -> bool:
46 """Download language parser if not already available."""
47 try:
48 if has_language(lang):
49 return True
50 # tslp 1.8.0 mistypes init() against _native.PackConfig, but the
51 # public re-export is options.PackConfig (a dataclass). Runtime is
52 # fine. Both share the same fields.
53 init(PackConfig(languages=[lang])) # type: ignore[arg-type]
54 return has_language(lang)
55 except Exception:
56 log.debug("Failed to download tree-sitter language: %s", lang)
57 return False
60def find_line(needle: str, lines: list[str], start: int) -> int:
61 """Find the first line index (1-based) containing needle, from start."""
62 for i in range(start, len(lines)):
63 if needle and needle in lines[i]:
64 return i + 1
65 return start + 1
68def _fallback_chunks(text: str) -> list[CodeChunk]:
69 """Fallback text chunking with approximate line tracking."""
70 raw = chunk_text(text)
71 lines = text.split("\n")
72 results: list[CodeChunk] = []
73 search_from = 0
75 for idx, chunk in enumerate(raw):
76 first_line = chunk.split("\n")[0][:80]
77 line_start = find_line(first_line, lines, search_from)
78 line_end = min(line_start + chunk.count("\n"), len(lines))
79 results.append(
80 CodeChunk(
81 chunk=chunk,
82 line_start=line_start,
83 line_end=line_end,
84 chunk_index=idx,
85 )
86 )
87 # line_start is 1-based; find_line's `start` is a 0-based index. Convert so
88 # the next search begins at this chunk's start line (not one past it), which
89 # matters when overlapping chunks share a first line.
90 search_from = line_start - 1
92 return results
95def _chunk_header(source_name: str, symbols: list[str], line_start: int, line_end: int) -> str:
96 """Build the metadata header prepended to a code chunk.
98 ``source_name`` is the relative source path, never the host's absolute path,
99 so an exported/shared corpus does not leak the operator's disk layout.
100 ``symbols`` are the names defined in the chunk (empty for an anonymous or
101 symbol-free span, in which case the name segment is omitted entirely).
102 """
103 header = f"# File: {source_name}"
104 if symbols:
105 header += f" | {', '.join(symbols)}"
106 header += f" (lines {line_start}-{line_end})"
107 return header
110def _line_span(start_line_zero_based: int, content: str) -> tuple[int, int]:
111 """1-based inclusive ``(line_start, line_end)`` the chunk's content covers.
113 ``line_end`` is derived from the content's own line count rather than the
114 parser's ``end_line``: tree-sitter reports ``end_line`` as the count of
115 newline-terminated lines, so a chunk whose final line has no trailing newline
116 (the last chunk of a file that does not end in one, or a sub-line split) would
117 otherwise under-report its last line by one.
118 """
119 line_start = start_line_zero_based + 1
120 # count("\n") plus one for a final line without a trailing newline; empty
121 # content yields a single (degenerate) line so line_end never precedes start.
122 spanned = content.count("\n") + (0 if content.endswith("\n") else 1)
123 line_end = line_start + max(spanned, 1) - 1
124 return line_start, line_end
127def _chunks_from_result(result: Any, source_name: str) -> list[CodeChunk]:
128 """Build :class:`CodeChunk` records from tree-sitter's size-bounded ``result.chunks``.
130 ``result.chunks`` is the ``chunk_max_size``-aware output: each entry's
131 ``content`` is the already-extracted UTF-8 text (so there is no manual
132 byte-offset slicing, which corrupts non-ASCII source) and its size honors
133 the configured budget. ``metadata.symbols_defined`` lists the symbols in the
134 chunk, including methods of a class, so nested symbols are not folded away.
135 """
136 chunks: list[CodeChunk] = []
137 for i, tc in enumerate(result.chunks):
138 symbols = list(tc.metadata.symbols_defined) if tc.metadata is not None else []
139 line_start, line_end = _line_span(tc.start_line, tc.content)
140 header = _chunk_header(source_name, symbols, line_start, line_end)
141 chunks.append(
142 CodeChunk(
143 chunk=f"{header}\n\n{tc.content}",
144 line_start=line_start,
145 line_end=line_end,
146 chunk_index=i,
147 )
148 )
149 return chunks
152def chunk_code(file_path: Path, source_name: str | None = None) -> list[CodeChunk]:
153 """Chunk a source file using tree-sitter-language-pack's process() API.
155 Emits the parser's size-bounded chunks (``chunk_max_size``-aware) with a
156 metadata header naming the relative ``source_name`` and the symbols defined
157 in each chunk. Falls back to token-based chunking when the language is
158 unsupported, the parser is unavailable, parsing fails, or it produces no
159 chunks. ``source_name`` defaults to the file's basename so the header never
160 carries the absolute path.
161 """
162 source_text = file_path.read_text(encoding="utf-8", errors="replace")
163 if not source_text.strip():
164 return []
166 label = source_name if source_name is not None else file_path.name
168 lang = _detect_language(file_path)
169 if not lang:
170 return _fallback_chunks(source_text)
172 try:
173 if not _ensure_language(lang):
174 return _fallback_chunks(source_text)
175 config = ProcessConfig(
176 lang,
177 structure=True,
178 symbols=True,
179 docstrings=True,
180 # tree-sitter's chunk_max_size is bytes; cfg.chunk_size is tokens,
181 # so convert with the same char budget the text path uses.
182 chunk_max_size=cfg.chunk_size * CHARS_PER_TOKEN,
183 )
184 result = process(source_text, config) # type: ignore[arg-type] # tslp 1.8.0 typing bug, see init() above
185 except Exception:
186 log.debug("tree-sitter process() failed for %s", file_path, exc_info=True)
187 return _fallback_chunks(source_text)
189 chunks = _chunks_from_result(result, label)
190 if not chunks:
191 return _fallback_chunks(source_text)
192 return chunks
195def is_code_file(file_path: Path) -> bool:
196 """Check if a file is supported by tree-sitter chunking."""
197 return detect_language(str(file_path)) is not None