diff --git a/requirements/requirements.txt b/requirements/requirements.txt index cba4317..ccac70f 100644 --- a/requirements/requirements.txt +++ b/requirements/requirements.txt @@ -15,3 +15,8 @@ google-cloud-iam protobuf jsonpath_ng +# tree-sitter and the Python grammar require Python >= 3.10. +# On Python 3.9, the PythonParser is unavailable and .py files are skipped. +tree-sitter>=0.21.0; python_version>='3.10' +tree-sitter-python>=0.21.0; python_version>='3.10' + diff --git a/vanir/language_parsers/BUILD.bazel b/vanir/language_parsers/BUILD.bazel index aa2acf7..5a9565c 100644 --- a/vanir/language_parsers/BUILD.bazel +++ b/vanir/language_parsers/BUILD.bazel @@ -31,6 +31,15 @@ py_library( ], ) +py_library( + name = "tree_sitter_base", + srcs = ["tree_sitter_base.py"], + deps = [ + ":abstract_language_parser", + ":common", + ], +) + py_library( name = "language_parsers", srcs = ["language_parsers.py"], @@ -39,5 +48,6 @@ py_library( ":common", "//vanir/language_parsers/cpp:cpp_parser", "//vanir/language_parsers/java:java_parser", + "//vanir/language_parsers/python:python_parser", ], ) diff --git a/vanir/language_parsers/language_parsers.py b/vanir/language_parsers/language_parsers.py index 0bc5a46..169f5e8 100644 --- a/vanir/language_parsers/language_parsers.py +++ b/vanir/language_parsers/language_parsers.py @@ -23,14 +23,25 @@ # pylint: disable=unused-import from vanir.language_parsers.cpp import cpp_parser from vanir.language_parsers.java import java_parser +from vanir.language_parsers.python import python_parser # pylint: enable=unused-import _P = TypeVar('_P', bound=abstract_language_parser.AbstractLanguageParser) +def _all_subclasses(cls): + """Recursively collect all concrete (non-abstract) subclasses of cls.""" + result = [] + for sub in cls.__subclasses__(): + if not sub.__abstractmethods__: + result.append(sub) + result.extend(_all_subclasses(sub)) + return result + + def get_parser_class(filename: str) -> Optional[Type[_P]]: """Returns the language parser class that handles the given file, or None.""" - parsers = abstract_language_parser.AbstractLanguageParser.__subclasses__() + parsers = _all_subclasses(abstract_language_parser.AbstractLanguageParser) ext = os.path.splitext(filename)[1] for parser_class in parsers: if ext in parser_class.get_supported_extensions(): diff --git a/vanir/language_parsers/python/BUILD.bazel b/vanir/language_parsers/python/BUILD.bazel new file mode 100644 index 0000000..f8092ea --- /dev/null +++ b/vanir/language_parsers/python/BUILD.bazel @@ -0,0 +1,25 @@ +# 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 + +# Bazel build rules for Vanir Python language parser. +load("@rules_python//python:defs.bzl", "py_library") +load("@vanir_deps//:requirements.bzl", "requirement") + +package(default_visibility = [ + "//visibility:public", +]) + +py_library( + name = "python_parser", + srcs = ["python_parser.py"], + deps = [ + "//vanir/language_parsers:abstract_language_parser", + "//vanir/language_parsers:common", + "//vanir/language_parsers:tree_sitter_base", + requirement("tree-sitter"), + requirement("tree-sitter-python"), + ], +) diff --git a/vanir/language_parsers/python/__init__.py b/vanir/language_parsers/python/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/vanir/language_parsers/python/python_parser.py b/vanir/language_parsers/python/python_parser.py new file mode 100644 index 0000000..a440287 --- /dev/null +++ b/vanir/language_parsers/python/python_parser.py @@ -0,0 +1,298 @@ +# 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 + +"""Vanir Python parser using tree-sitter. + +tree-sitter and tree-sitter-python require Python >= 3.10. On Python 3.9 +the packages are not installed (requirements.txt version guard), so +_TREE_SITTER_AVAILABLE will be False and get_supported_extensions() returns +[] — the dispatcher finds no parser for .py files and skips them silently. +""" + +from typing import Sequence + +from vanir.language_parsers.tree_sitter_base import _flat_tokens_cursor +from vanir.language_parsers.tree_sitter_base import TreeSitterParserBase + +try: + import tree_sitter_python as _tspython + from tree_sitter import Language as _Language + from tree_sitter import QueryCursor as _QueryCursor + _PY_LANGUAGE = _Language(_tspython.language()) + # Compiled once at module import — never per file. + _FUNC_QUERY = _PY_LANGUAGE.query(""" + (function_definition + name: (identifier) @func.name + body: (block) @func.body) @func.def + """) + _LOCALS_QUERY = _PY_LANGUAGE.query(""" + (assignment left: (_) @assign.lhs) + (augmented_assignment left: (identifier) @aug.lhs) + (for_statement left: (_) @for.lhs) + (named_expression name: (identifier) @walrus.name) + (call function: (identifier) @call.name) + """) + _TREE_SITTER_AVAILABLE = True +except ImportError: + _PY_LANGUAGE = None + _FUNC_QUERY = None + _LOCALS_QUERY = None + _TREE_SITTER_AVAILABLE = False + + +# --------------------------------------------------------------------------- +# Helper functions (module-level, reused by _extract_param_names and +# _collect_locals_calls) +# --------------------------------------------------------------------------- + +def _get_param_name(param_node): + """Return the identifier name from a parameter node, or None. + + Handles: identifier, typed_parameter, default_parameter, + typed_default_parameter, list_splat_pattern (*args), + dictionary_splat_pattern (**kwargs), bare * separator. + """ + t = param_node.type + if t == 'identifier': + return param_node.text.decode('utf-8', errors='replace') + if t in ('typed_parameter', 'default_parameter', 'typed_default_parameter', + 'list_splat_pattern', 'dictionary_splat_pattern'): + for child in param_node.named_children: + if child.type == 'identifier': + return child.text.decode('utf-8', errors='replace') + return None + + +def _collect_ids_from_lhs(lhs_node, out): + """Add identifier name(s) from an assignment LHS node into out (set).""" + if lhs_node is None: + return + t = lhs_node.type + if t == 'identifier': + out.add(lhs_node.text.decode('utf-8', errors='replace')) + elif t in ('pattern_list', 'tuple_pattern'): + for child in lhs_node.named_children: + _collect_ids_from_lhs(child, out) + + +# --------------------------------------------------------------------------- +# PythonParser — thin subclass of TreeSitterParserBase +# --------------------------------------------------------------------------- + +class PythonParser(TreeSitterParserBase): + """Vanir Python parser backed by tree-sitter. + + Implements AbstractLanguageParser via TreeSitterParserBase. On Python 3.9 + (where tree-sitter is not installed), get_supported_extensions() returns [] + so the dispatcher skips .py files without raising an error. + """ + + LANGUAGE = _PY_LANGUAGE + _FUNC_QUERY = _FUNC_QUERY + SKIP_TOKEN_TYPES = frozenset({ + 'comment', + 'newline', + 'indent', + 'dedent', + 'line_continuation', + }) + STRING_NODE_TYPE = 'string' + + @classmethod + def get_supported_extensions(cls) -> Sequence[str]: + return ['.py'] if _TREE_SITTER_AVAILABLE else [] + + def _extract_param_names(self, params_node) -> list: + """Extract parameter names from a Python parameters node.""" + parameters = [] + for child in params_node.named_children: + param_name = _get_param_name(child) + if param_name: + parameters.append(param_name) + return parameters + + def _extract_annotations(self, func_node) -> tuple: + """Extract return type and parameter type annotations.""" + return_types = [] + used_data_types = [] + + return_type_node = func_node.child_by_field_name('return_type') + if return_type_node: + toks = _flat_tokens_cursor( + return_type_node, self.SKIP_TOKEN_TYPES, self.STRING_NODE_TYPE) + if toks: + return_types = [toks] + + params_node = func_node.child_by_field_name('parameters') + if params_node: + for child in params_node.named_children: + if child.type in ('typed_parameter', 'typed_default_parameter'): + type_node = child.child_by_field_name('type') + if type_node: + toks = _flat_tokens_cursor( + type_node, + self.SKIP_TOKEN_TYPES, + self.STRING_NODE_TYPE) + if toks: + used_data_types.append(toks) + + return return_types, used_data_types + + def _collect_locals_calls(self, body_node, nested_ranges) -> tuple: + """Collect local vars and called functions using a scoped C-level query. + + Uses QueryCursor(_LOCALS_QUERY).captures(body_node) and filters out + any captures that fall inside a nested function's byte range. + """ + local_vars: set = set() + called_fns: set = set() + + captures = _QueryCursor(_LOCALS_QUERY).captures(body_node) + + def _is_nested(cap_node): + sb = cap_node.start_byte + return any(s < sb < e for s, e in nested_ranges) + + for node in captures.get('assign.lhs', []): + if not _is_nested(node): + _collect_ids_from_lhs(node, local_vars) + + for node in captures.get('aug.lhs', []): + if not _is_nested(node) and node.type == 'identifier': + local_vars.add(node.text.decode('utf-8', errors='replace')) + + for node in captures.get('for.lhs', []): + if not _is_nested(node): + _collect_ids_from_lhs(node, local_vars) + + for node in captures.get('walrus.name', []): + if not _is_nested(node) and node.type == 'identifier': + local_vars.add(node.text.decode('utf-8', errors='replace')) + + for node in captures.get('call.name', []): + if not _is_nested(node) and node.type == 'identifier': + called_fns.add(node.text.decode('utf-8', errors='replace')) + + return local_vars, called_fns + + +# --------------------------------------------------------------------------- +# Debug harness (run directly: python python_parser.py [file.py]) +# --------------------------------------------------------------------------- + +def _debug_print_tree(node, indent=0): + """Recursively print the tree-sitter AST for inspection.""" + prefix = ' ' * indent + is_leaf = not node.children + if is_leaf: + text = node.text.decode('utf-8', errors='replace') + display = repr(text) if len(text) <= 60 else repr(text[:57] + '...') + print(f"{prefix}[{'named' if node.is_named else 'anon '}] " + f"{node.type!r:35s} {display} " + f"L{node.start_point[0]+1}:{node.start_point[1]}") + else: + print(f"{prefix}[{'named' if node.is_named else 'anon '}] " + f"{node.type!r:35s} " + f"L{node.start_point[0]+1}:{node.start_point[1]}" + f"–L{node.end_point[0]+1}:{node.end_point[1]}") + for child in node.children: + _debug_print_tree(child, indent + 1) + + +if __name__ == '__main__': + import sys + + if not _TREE_SITTER_AVAILABLE: + print('tree-sitter is not available (requires Python >= 3.10)') + sys.exit(1) + + if len(sys.argv) > 1: + with open(sys.argv[1], 'rb') as f: + source = f.read() + filename = sys.argv[1] + else: + source = b'''\ +import os +from typing import Optional, List + +GLOBAL_CONST = 42 + +class MyClass: + class_var: int = 0 + + @staticmethod + def simple_method(x, y): + z = x + y + return z + + def annotated_method( + self, + data: List[str], + count: int = 0, + ) -> Optional[bool]: + """Triple-quoted docstring for testing.""" + result = None + for item in data: + value = transform(item) + result = validate(value, count) + return result + +def top_level(arg1: int, arg2: str = "hi") -> bool: + local_a, local_b = arg1, arg2 + x = f"formatted {arg1}" + if (n := len(arg2)) > 10: + helper(n) + return True + +async def async_func(items): + for i, item in enumerate(items): + process(item) +''' + import tempfile, os + tmp = tempfile.NamedTemporaryFile(suffix='.py', delete=False) + tmp.write(source) + tmp.close() + filename = tmp.name + + parser_obj = PythonParser(filename) + results = parser_obj.get_chunks() + + print('=' * 70) + print(f'FILE: {filename}') + print(f'parse_errors : {results.parse_errors}') + print() + + print('LINE CHUNK (first 10 lines with tokens):') + for ln in sorted(results.line_chunk.tokens)[:10]: + print(f' {ln:3d}: {results.line_chunk.tokens[ln]}') + print() + + print('FUNCTION CHUNKS:') + from vanir import normalizer as _normalizer + for chunk in results.function_chunks: + normalized = _normalizer.normalize_function_chunk(chunk) + print(f'\n func : {chunk.name}') + print(f' parameters : {list(chunk.parameters)}') + print(f' return_types : {list(chunk.return_types)}') + print(f' used_data_types : {list(chunk.used_data_types)}') + print(f' local_variables : {list(chunk.local_variables)}') + print(f' called_functions: {list(chunk.called_functions)}') + print(f' tokens : {list(chunk.tokens)}') + print(f' normalized : {normalized}') + + print() + print('=' * 70) + print('ERROR HANDLING TEST (broken Python):') + broken = b'$$$ %%% @@@\ndef foo():\n pass\n' + tmp2 = tempfile.NamedTemporaryFile(suffix='.py', delete=False) + tmp2.write(broken) + tmp2.close() + broken_results = PythonParser(tmp2.name).get_chunks() + print(f' parse_errors: {broken_results.parse_errors}') + os.unlink(tmp2.name) + + if len(sys.argv) <= 1: + os.unlink(filename) 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. + """