wip: call_only filter for JS/TS + checkpoint before GBNF rewrite
This commit is contained in:
parent
44415c5b42
commit
617b7ad578
2 changed files with 17 additions and 5 deletions
|
|
@ -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)
|
||||||
|
|
|
||||||
|
|
@ -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)
|
||||||
|
|
|
||||||
Loading…
Add table
Reference in a new issue