wip: call_only filter for JS/TS + checkpoint before GBNF rewrite

This commit is contained in:
tobjend 2026-07-12 15:20:55 +02:00
parent 44415c5b42
commit 617b7ad578
2 changed files with 17 additions and 5 deletions

View file

@ -220,12 +220,17 @@ def frequency_filter(sequences, min_coverage=0.2):
return filtered return filtered
_CALL_ONLY_EXTS = {".ts", ".tsx", ".js", ".jsx", ".mts", ".mjs"}
def _preprocess_file(fp): def _preprocess_file(fp):
"""Preprocess one file. Module-level for ProcessPoolExecutor.""" """Preprocess one file. Module-level for ProcessPoolExecutor."""
ext = os.path.splitext(fp)[1].lower()
call_only = ext in _CALL_ONLY_EXTS
with open(fp) as f: with open(fp) as f:
code = f.read() code = f.read()
sequences = [] sequences = []
for method_seq in preprocess_by_method(fp, code): for method_seq in preprocess_by_method(fp, code, call_only=call_only):
if method_seq: if method_seq:
sequences.append(method_seq) sequences.append(method_seq)
return (fp, sequences) return (fp, sequences)

View file

@ -359,11 +359,16 @@ def _find_method_bodies(tree):
return bodies return bodies
def preprocess_by_method(file_path: str, code: str): def preprocess_by_method(file_path: str, code: str, call_only=False):
"""Preprocess and group behavioral tokens by enclosing method body. """Preprocess and group behavioral tokens by enclosing method body.
Returns list of sequences, one per function/method found. Returns list of sequences, one per function/method found.
Each sequence is [(capture_name, text, line_number), ...]. Each sequence is [(capture_name, text, line_number), ...].
Args:
call_only: When True, only keep CALL_PREFIXES captures (function
calls, method calls). Use for JS/TS where keyword captures
produce truncated text due to tree-sitter node boundary issues.
""" """
ext = os.path.splitext(file_path)[1].lower() ext = os.path.splitext(file_path)[1].lower()
lang, query_name = _load_grammar(ext) lang, query_name = _load_grammar(ext)
@ -377,9 +382,10 @@ def preprocess_by_method(file_path: str, code: str):
cursor = QueryCursor(query) cursor = QueryCursor(query)
captures = cursor.captures(tree.root_node) captures = cursor.captures(tree.root_node)
prefix_filter = CALL_PREFIXES if call_only else BEHAVIORAL_PREFIXES
items = [] items = []
for capname, nodes in captures.items(): for capname, nodes in captures.items():
if not capname.startswith(BEHAVIORAL_PREFIXES): if not capname.startswith(prefix_filter):
continue continue
for node in nodes: for node in nodes:
text = sanitize_symbol(code[node.start_byte:node.end_byte], capname) text = sanitize_symbol(code[node.start_byte:node.end_byte], capname)
@ -403,7 +409,7 @@ def preprocess_by_method(file_path: str, code: str):
return sequences return sequences
def preprocess(file_path: str, code: str): def preprocess(file_path: str, code: str, call_only=False):
ext = os.path.splitext(file_path)[1].lower() ext = os.path.splitext(file_path)[1].lower()
lang, query_name = _load_grammar(ext) lang, query_name = _load_grammar(ext)
query_src = _load_query(query_name) query_src = _load_query(query_name)
@ -420,9 +426,10 @@ def preprocess(file_path: str, code: str):
cursor = QueryCursor(query) cursor = QueryCursor(query)
captures = cursor.captures(tree.root_node) captures = cursor.captures(tree.root_node)
prefix_filter = CALL_PREFIXES if call_only else BEHAVIORAL_PREFIXES
items = [] items = []
for capname, nodes in captures.items(): for capname, nodes in captures.items():
if not capname.startswith(BEHAVIORAL_PREFIXES): if not capname.startswith(prefix_filter):
continue continue
for node in nodes: for node in nodes:
text = sanitize_symbol(code[node.start_byte:node.end_byte], capname) text = sanitize_symbol(code[node.start_byte:node.end_byte], capname)