From 4e8228fd838900620b31a8273a6267a01665fe47 Mon Sep 17 00:00:00 2001 From: Awe Srush Date: Fri, 17 Jul 2026 15:49:20 -0700 Subject: [PATCH 1/2] feat(language_parsers): add TreeSitterParserBase for tree-sitter language parsers --- vanir/language_parsers/tree_sitter_base.py | 256 +++++++++++++++++++++ 1 file changed, 256 insertions(+) create mode 100644 vanir/language_parsers/tree_sitter_base.py diff --git a/vanir/language_parsers/tree_sitter_base.py b/vanir/language_parsers/tree_sitter_base.py new file mode 100644 index 0000000..3b09920 --- /dev/null +++ b/vanir/language_parsers/tree_sitter_base.py @@ -0,0 +1,256 @@ +# Copyright 2025 Google LLC +# +# Use of this source code is governed by a BSD-style +# license that can be found in the LICENSE file or at +# https://developers.google.com/open-source/licenses/bsd + +"""Shared base class for tree-sitter-backed Vanir language parsers. + +This module does NOT import tree-sitter at module level. Subclasses guard +their own tree-sitter imports with try/except ImportError so that the module +can be imported on Python 3.9 (where tree-sitter is not installed) without +raising an error. + +Inheritance chain: + AbstractLanguageParser (existing ABC — unchanged) + └── TreeSitterParserBase (this module) + └── PythonParser (thin subclass) + └── GoParser (future) +""" + +import abc +from typing import Optional, Sequence, Tuple + +from vanir.language_parsers import abstract_language_parser +from vanir.language_parsers import common + + +# --------------------------------------------------------------------------- +# Iterative cursor utilities (no tree-sitter import needed — callers pass nodes) +# --------------------------------------------------------------------------- + +def _collect_tokens_cursor(node, skip_types, string_type, out_dict): + """Iterative TreeCursor DFS — populates out_dict[line_number, list[str]]. + + - Nodes in skip_types are skipped (their whole subtree is pruned). + - Nodes of string_type are emitted as a single opaque token. + - Leaf nodes (no children) are emitted as individual tokens. + - All other nodes are descended into. + + out_dict: dict[int, list[str]] key = 1-indexed line number. + """ + cursor = node.walk() + while True: + n = cursor.node + ntype = n.type + if ntype in skip_types: + pass # prune subtree — fall through to sibling/parent nav + elif ntype == string_type or not n.children: + line = n.start_point[0] + 1 + out_dict.setdefault(line, []).append( + n.text.decode('utf-8', errors='replace')) + elif cursor.goto_first_child(): + continue # descended — restart loop from new position + # Sibling / parent navigation + while not cursor.goto_next_sibling(): + if not cursor.goto_parent(): + return + + +def _flat_tokens_cursor(node, skip_types, string_type): + """Return a flat ordered list of token strings from node's subtree.""" + buf = {} + _collect_tokens_cursor(node, skip_types, string_type, buf) + return [tok for line in sorted(buf) for tok in buf[line]] + + +def _collect_errors_cursor(root): + """Iterative TreeCursor walk that collects ERROR nodes. + + Returns a list of common.ParseError, one per ERROR node encountered. + The walk visits every node (ERROR nodes are not pruned) so that + errors inside partially-parsed subtrees are also reported. + """ + errors = [] + cursor = root.walk() + while True: + n = cursor.node + if n.type == 'ERROR': + line = n.start_point[0] + 1 + col = n.start_point[1] + bad_token = n.text.decode('utf-8', errors='replace') + display = bad_token[:80] + ('...' if len(bad_token) > 80 else '') + errors.append(common.ParseError( + line=line, + column=col, + bad_token=display, + message='syntax error', + )) + # Always try to descend — errors inside ERROR subtrees are also reported. + if cursor.goto_first_child(): + continue + while not cursor.goto_next_sibling(): + if not cursor.goto_parent(): + return errors + + +def _overlaps(func_start, func_end, ranges): + """Return True if [func_start, func_end] overlaps any (start, end) range. + + An empty ranges list means "no filter — include all functions". + """ + if not ranges: + return True + return any(func_start <= r_end and func_end >= r_start + for r_start, r_end in ranges) + + +# --------------------------------------------------------------------------- +# TreeSitterParserBase +# --------------------------------------------------------------------------- + +class TreeSitterParserBase(abstract_language_parser.AbstractLanguageParser): + """Shared base for tree-sitter-backed Vanir language parsers. + + Subclasses must set the following class attributes and implement the + three abstract methods below. + """ + + # Subclasses set these at class level. + LANGUAGE = None # tree_sitter.Language instance, or None if unavailable + _FUNC_QUERY = None # compiled Query for function discovery + SKIP_TOKEN_TYPES = frozenset() + STRING_NODE_TYPE = 'string' + + def __init__(self, filename: str): + # Lazy import so the module itself has no tree-sitter dependency at + # import time. + from tree_sitter import Parser # pylint: disable=import-outside-toplevel + with open(filename, 'rb') as f: + self._source = f.read() + self._tree = Parser(self.LANGUAGE).parse(self._source) + if self._tree is None: + raise RuntimeError( + f'tree-sitter failed to parse {filename!r}') + + def get_chunks( + self, + affected_line_ranges_for_functions: Optional[ + Sequence[Tuple[int, int]] + ] = None, + ) -> common.ParseResults: + from tree_sitter import QueryCursor # pylint: disable=import-outside-toplevel + + if affected_line_ranges_for_functions is None: + affected_line_ranges_for_functions = [] + + root = self._tree.root_node + + # --- Step 1: function discovery via C-level Query --- + all_matches = list(QueryCursor(self._FUNC_QUERY).matches(root)) + all_func_nodes = [m['func.def'][0] for _, m in all_matches] + all_name_nodes = [m['func.name'][0] for _, m in all_matches] + + # --- Step 2: build one FunctionChunkBase per matched function --- + function_chunks = [] + for i, func_node in enumerate(all_func_nodes): + start_line = func_node.start_point[0] + 1 + end_line = func_node.end_point[0] + 1 + + if not _overlaps(start_line, end_line, + affected_line_ranges_for_functions): + continue + + name_node = all_name_nodes[i] if i < len(all_name_nodes) else None + name = (name_node.text.decode('utf-8', errors='replace') + if name_node else '') + + params_node = func_node.child_by_field_name('parameters') + parameters = (self._extract_param_names(params_node) + if params_node else []) + + return_types, used_data_types = self._extract_annotations(func_node) + + # Byte ranges of functions nested directly inside this function + # (used by _collect_locals_calls to exclude their captures). + nested_ranges = [ + (n.start_byte, n.end_byte) + for n in all_func_nodes + if (n.start_byte > func_node.start_byte + and n.end_byte <= func_node.end_byte) + ] + + local_vars_set: set = set() + called_fns_set: set = set() + body_node = func_node.child_by_field_name('body') + if body_node: + local_vars_set, called_fns_set = self._collect_locals_calls( + body_node, nested_ranges) + local_vars_set -= set(parameters) + + tokens = _flat_tokens_cursor( + func_node, self.SKIP_TOKEN_TYPES, self.STRING_NODE_TYPE) + + function_chunks.append(common.FunctionChunkBase( + name=name, + return_types=return_types, + parameters=parameters, + used_data_types=used_data_types, + local_variables=sorted(local_vars_set), + called_functions=sorted(called_fns_set), + tokens=tokens, + )) + + # --- Step 3: line chunk (whole-file token collection) --- + tokens_by_line: dict = {} + _collect_tokens_cursor( + root, self.SKIP_TOKEN_TYPES, self.STRING_NODE_TYPE, tokens_by_line) + line_chunk = common.LineChunkBase(tokens=tokens_by_line) + + # --- Step 4: parse errors --- + parse_errors = _collect_errors_cursor(root) + + return common.ParseResults( + function_chunks=function_chunks, + line_chunk=line_chunk, + parse_errors=parse_errors, + ) + + # --- Abstract methods (language-specific) --- + + @abc.abstractmethod + def _extract_param_names(self, params_node) -> list: + """Extract parameter names from a parameters node. + + Args: + params_node: The tree-sitter node for the parameter list. + + Returns: + List of parameter name strings (in order). + """ + + @abc.abstractmethod + def _extract_annotations(self, func_node) -> tuple: + """Extract type annotations from a function definition node. + + Args: + func_node: The tree-sitter function_definition node. + + Returns: + (return_types, used_data_types) where: + return_types — [] or [[token, ...]] (at most one return annotation). + used_data_types — [[token, ...], ...] (one list per param annotation). + """ + + @abc.abstractmethod + def _collect_locals_calls(self, body_node, nested_ranges) -> tuple: + """Collect local variable names and called function names from a body. + + Args: + body_node: The tree-sitter node for the function body. + nested_ranges: List of (start_byte, end_byte) pairs for nested + function definitions that should be excluded from the search. + + Returns: + (local_vars, called_fns) — both are sets of strings. + """ From 44204c26179d12fa29796b5220d428b906c1b5c1 Mon Sep 17 00:00:00 2001 From: Awe Srush Date: Thu, 23 Jul 2026 16:16:59 -0700 Subject: [PATCH 2/2] trigger CLA recheck