from __future__ import annotations import os import re from collections.abc import Iterable from pathlib import Path from typing import TYPE_CHECKING, Any, Literal, TypedDict, cast import sublime import sublime_plugin from .core import ( BUFFER_ID_TO_TREE, SCOPE_TO_LANGUAGE, byte_offset, check_scope, get_scope, get_view_text, make_tree_dict, mutable_settings, parse, publish_tree_update, trim_cached_trees, ) from .utils import ( PROJECT_ROOT, get_queries_path, get_scope_to_language_name, get_scope_to_queries_name, log, maybe_none, not_none, ) if TYPE_CHECKING: from tree_sitter import Language, Node, Tree from .core import Injection SYMBOLS_FILE = "symbols.scm" # # Public-facing API functions, and some helper functions # def get_tracked_buffer_ids(): """ Get buffer ids for all tracked buffers. """ return list(BUFFER_ID_TO_TREE.keys()) def get_tree_dict(buffer_id: int): """ Get tree dict being maintained for this buffer, or instantiate new tree dict on the fly. """ if not isinstance(cast(Any, buffer_id), int) or not (view := get_view_from_buffer_id(buffer_id)): return if not (scope := get_scope(view)) or not (scope := check_scope(scope)): BUFFER_ID_TO_TREE.pop(buffer_id, None) return tree_dict = BUFFER_ID_TO_TREE.get(buffer_id) if not tree_dict or tree_dict["scope"] != scope: from tree_sitter import Parser view_text = get_view_text(view) # `get_tree_dict` can be called synchronously from a command on the main thread; we pass `only_downloaded=True` # to ensure computing injections doesn't block on network I/O for a not-yet-cached injected language tree = parse(Parser(), scope, view_text) BUFFER_ID_TO_TREE[buffer_id] = make_tree_dict(tree, view_text, scope, only_downloaded=True) trim_cached_trees() publish_tree_update(view.window(), buffer_id=buffer_id, scope=scope) return BUFFER_ID_TO_TREE.get(buffer_id) def get_view_from_buffer_id(buffer_id: int) -> sublime.View | None: """ Utilify function. Ensures `None` returned if a "dead" buffer id passed. """ buffer = sublime.Buffer(buffer_id) view = buffer.primary_view() return view if maybe_none(view.id()) is not None else None def get_tree_from_code(scope: str, s: str | bytes): """ Get a syntax tree back for source code `s`. """ from tree_sitter import Parser if not (validated_scope := check_scope(scope)): return None parser = Parser(SCOPE_TO_LANGUAGE[validated_scope]) return parser.parse(s.encode() if isinstance(s, str) else s) def query_node_with_language(language: Language, node: Node | InjectedNode, query_s: str): """ Query a node with `query_s`, against `language` directly (see `query_node_with_s` for the scope-based version, used for a buffer's own top-level language; this one also backs symbol search in injected trees, each of which has its own `Language` - see `get_captures_from_injections`). Returns `(node, capture_name)` tuples, in the same document order in which their matches were found, i.e. an ancestor's captures always come before its descendants'. This ordering is relied on by `get_captures_from_nodes`, e.g. to build symbol breadcrumbs, since a node's breadcrumb ancestors must already have been seen by the time the node itself is processed. Note `QueryCursor.captures` doesn't preserve this ordering (it groups all captures by capture name), so we build this list from `QueryCursor.matches` instead. `QueryCursor.matches` needs a real `tree_sitter.Node`, so `node` is unwrapped if it's an `InjectedNode` (see `unwrap`). See https://github.com/tree-sitter/py-tree-sitter#pattern-matching """ from tree_sitter import Query, QueryCursor cursor = QueryCursor(Query(language, query_s)) return [ (captured_node, name) for _, captures in cursor.matches(unwrap(node)) for name, nodes in captures.items() for captured_node in nodes ] def query_node_with_s(scope: str | None, node: Node | InjectedNode, query_s: str): """ `query_node_with_language`, resolving `Language` from a Sublime scope (see `SCOPE_TO_LANGUAGE`) rather than taking one directly. """ if not (scope := check_scope(scope)): return return query_node_with_language(SCOPE_TO_LANGUAGE[scope], node, query_s) def get_query_s_from_file( queries_name: str, queries_path: str | Path = "", symbols_file: str = SYMBOLS_FILE, ignore_file_not_found: bool = False, ) -> str: """ Handle `inherits` "pragmas" of the following structure: `; inherits: lang(,other_lang)` Passing `ignore_file_not_found=True` to recursive calls of this function essentially makes inherits pragma not "strict". See https://github.com/sublime-treesitter/TreeSitter/pull/6 for more context. """ INHERITS_PREFIX = "; inherits:" queries_path = os.path.expanduser(queries_path or get_queries_path(mutable_settings.d)) path = Path(queries_path) / queries_name / symbols_file names: list[str] = [] try: with open(path, "r") as f: query_s = f.read() except FileNotFoundError: if not ignore_file_not_found: raise log(f"query file not found, so it was ignored:\n{path}") query_s = "" else: for line in query_s.splitlines(): if line.startswith(INHERITS_PREFIX): names = [name.strip() for name in line.split(INHERITS_PREFIX)[1].split(",") if name] queries = [ get_query_s_from_file( queries_name=name, queries_path=queries_path, symbols_file=symbols_file, ignore_file_not_found=True, ) for name in names ] return "\n".join([query_s, *queries]) def get_tags_query_s(language_name: str) -> str: """ Fall back to the "tags" query `tree_sitter_language_pack` bundles for `language_name`, when this plugin doesn't ship (and the user hasn't supplied) a `symbols.scm` for it. This is the community-standard convention for a language-agnostic symbol outline (also used by e.g. GitHub's semantic, and ctags-like tooling generally): its `@definition.class`/`@definition.function`/... captures already match `CAPTURE_NAME_PREFIX`/`CAPTURE_NAME_TO_KIND` here, so goto/select symbol works with no changes, just without breadcrumbs (`tags.scm` has no equivalent of this plugin's `@breadcrumb.N` pragma - see `TreeSitterGotoSymbolCommand`). """ try: from tree_sitter_language_pack import get_tags_query except ImportError: return "" return get_tags_query(language_name) or "" def walk_tree(tree_or_node: Tree | Node, max_depth: int | None = None): """ Walk all the nodes under `tree_or_node`. See https://github.com/tree-sitter/py-tree-sitter/issues/33#issuecomment-864557166 """ cursor = tree_or_node.walk() reached_root = False while not reached_root: yield cursor.node, cursor if (max_depth is None or cursor.depth < max_depth) and cursor.goto_first_child(): # Don't walk children if we've already reached `max_depth` continue if cursor.goto_next_sibling(): continue retracing = True while retracing: if not cursor.goto_parent(): retracing = False reached_root = True if cursor.goto_next_sibling(): retracing = False def get_injection_for_node(node: Node, injections: list[Injection]) -> Injection | None: """ Does `node`'s byte range exactly match one of `injections`' content ranges, i.e. is `node` an `@injection.content` node a tree was parsed for (see `core.compute_injections`)? Matched against `Injection.content_ranges`, not the parsed tree's own root node span, since those can differ (see `Injection`'s docstring). Doesn't match a node merely contained within one. """ for injection in injections: if (node.start_byte, node.end_byte) in injection["content_ranges"]: return injection return None def format_injected_tree( root_node: Node, injections: list[Injection], format_node, indent: str, depth_offset: int = 0, ) -> list[str]: """ Render `root_node`'s tree as indented lines with `format_node(node, field_name)`, one per node, splicing in each injected tree in `injections` (see `core.compute_injections`) right after the node it was injected into. Used by `TreeSitterPrintTreeCommand`. """ lines: list[str] = [] for n, cursor in walk_tree(root_node): node = not_none(n) depth = cursor.depth + depth_offset lines.append(f"{indent * depth}{format_node(node, cursor.field_name)}") if injection := get_injection_for_node(node, injections): lines.append(f"{indent * (depth + 1)}[injected: {injection['language_name']}]") lines.extend( format_injected_tree( injection["tree"].root_node, injection["children"], format_node, indent, depth_offset=depth + 2 ) ) return lines def descendant_for_byte_range(node: Node, start_byte: int, end_byte: int) -> Node | None: """ Get the smallest node within the given byte range. This API added in September 2023: https://github.com/tree-sitter/py-tree-sitter/pull/150/files See also: https://tree-sitter.github.io/tree-sitter/using-parsers#named-vs-anonymous-nodes """ return node.descendant_for_byte_range(start_byte, end_byte) class InjectedNode: """ Wraps a `tree_sitter.Node`, so that navigating via `.parent`/`.children` can cross injection boundaries (see `core.compute_injections`) as if the outer tree and every tree injected into it (transitively) were one tree. This exists because `tree_sitter.Tree`/`Node` can't literally be merged across languages: they're immutable native structures, one per `Language`, and there's no API for a `Node` in one `Tree` to have a `.parent` living in a different `Tree`. So instead, this is the "glue": `.children` descends into an injected tree if a child is exactly an `@injection.content` node (see `get_injection_for_node`), and `.parent` climbs back out to the outer node an injected tree's root replaced, once the wrapped node's own native `.parent` is exhausted. Every other attribute (`.type`, `.start_byte`, `.text`, `.id`, `.child_by_field_name`, ...) is delegated straight to the wrapped node, unchanged - use `.raw` to get it directly, e.g. to run a `Query` against it (`QueryCursor` needs a real `tree_sitter.Node`, not this wrapper). Deliberately not used by anything that needs to walk a whole subtree (`walk_tree`, `get_cousins`, symbol queries): for those, wrapping every visited node is wasted overhead, and (for `get_cousins`) crossing isn't even wanted - see `core.compute_injections`'s module docstring for which commands cross injection boundaries and which don't. """ __slots__ = ("_injections", "_outer", "_own_injection", "raw") def __init__( self, node: Node, injections: list[Injection], outer: InjectedNode | None = None, own_injection: Injection | None = None, ): self.raw = node self._injections = injections self._outer = outer self._own_injection = own_injection @property def parent(self) -> InjectedNode | None: if (parent := self.raw.parent) is not None: return InjectedNode(parent, self._injections, self._outer, self._own_injection) return self._outer @property def children(self) -> list[InjectedNode]: return [self._wrap_child(c) for c in self.raw.children] @property def injections(self) -> list[Injection]: """ The injections (see `core.compute_injections`) rooted within this node's own tree - i.e. the ones relevant to finding a language injected *under* this node, wherever in the injection hierarchy this node itself is. Used by `get_captures_from_nodes` to search for symbols in injected trees under an arbitrary starting node, not just a whole buffer's root. """ return self._injections @property def own_injection(self) -> Injection | None: """ The `Injection` (see `core.compute_injections`) this node's own tree was parsed for, or `None` if this node is still within the buffer's own top-level tree. Used by `get_captures_from_nodes` to resolve the right query and language for a starting node that's already inside an injected tree - e.g. one selected inside a Python fenced code block in a Markdown buffer - rather than always assuming the buffer's own. """ return self._own_injection def _wrap_child(self, child: Node) -> InjectedNode: if injection := get_injection_for_node(child, self._injections): outer = InjectedNode(child, self._injections, self._outer, self._own_injection) return InjectedNode(injection["tree"].root_node, injection["children"], outer, injection) return InjectedNode(child, self._injections, self._outer, self._own_injection) def __getattr__(self, name: str) -> Any: return getattr(self.raw, name) def __eq__(self, other: object) -> bool: return isinstance(other, InjectedNode) and self.raw == other.raw def __hash__(self) -> int: return hash(self.raw.id) def __repr__(self) -> str: return f"InjectedNode({self.raw!r})" def unwrap(node: Node | InjectedNode) -> Node: """ Get the plain `tree_sitter.Node` underneath `node`, whether or not it's an `InjectedNode`. """ return node.raw if isinstance(node, InjectedNode) else node def resolve_node_for_range( root: Node, injections: list[Injection], start_byte: int, end_byte: int, outer: InjectedNode | None = None, own_injection: Injection | None = None, ) -> InjectedNode | None: """ Find the smallest node spanning `[start_byte, end_byte)` starting from `root`, descending into an injected tree (see `core.compute_injections`) whenever one contains the range, rather than returning a node from the outer, pre-injection tree (e.g. an unparsed token inside Markdown's `inline` node). Wraps the result in an `InjectedNode` so further navigation (`.parent`, `.children`) can cross back out. Matched against `Injection.content_ranges`, not the parsed tree's own root node span (see `Injection`'s docstring) - e.g. a point in the indentation before the first statement in an HTML `