tree-sitter-seminar/src/ts-example/ts_minimal.py

import tempfile
from typing import Any

import tree_sitter_cpp as tscpp
import tree_sitter_python as tspython
from arg_conf import parser as arg_parser
from command_runner import ICommandRunner
from log_conf import logger
from tree_sitter import Language, Parser, Query, QueryCursor, Tree

PY_LANGUAGE = Language(tspython.language())
CPP_LANGUAGE = Language(tscpp.language())


def arguments():
    global arg_parser
    subparsers = arg_parser.add_subparsers(
        dest="command",
        required=True,
    )

    module_parser = subparsers.add_parser("minimal")

    module_parser.add_argument(
        "--language", help="Setup parsing language", default="Python"
    )
    module_parser.add_argument("--file", "-f", help="File to parse", required=True)
    logger.debug("Added Tree Sitter Minimal Example Arguments to the list")


class SingleFileParser(ICommandRunner):
    def __init__(self, args):
        self.normalized_language: Language
        self.tree: Tree
        self.parser: Parser
        language = args.language.casefold()
        if language in {"c++", "cpp"}:
            self.parser = Parser(CPP_LANGUAGE)
            self.normalized_language = CPP_LANGUAGE
            self.language_name = "cpp"
        else:
            self.parser = Parser(PY_LANGUAGE)
            self.normalized_language = PY_LANGUAGE
            self.language_name = "python"
        logger.info(f"Parser is installed to {self.language_name}")

    def node_text(self, node, source: bytes) -> str:
        return source[node.start_byte : node.end_byte].decode(
            "utf-8",
            errors="replace",
        )

    def final_callee_name(self, callee) -> str:
        """Return the last named component of a call expression.

        Selecting this from the syntax tree is more reliable than splitting
        source text: template arguments (``obj.foo<int>()``) and Python
        subscriptions (``obj.foo[T]()``) may themselves contain ``.`` or
        ``::``.
        """
        # These are grammar field names used by Python attributes/subscripts
        # and C++ field expressions, qualified identifiers, and templates.
        for field_name in ("attribute", "field", "name", "value", "argument"):
            child = callee.child_by_field_name(field_name)
            if child is not None:
                return self.final_callee_name(child)
        return self.node_text(callee, self.source_code)

    def get_usages(self, function: str):
        query_source: str = ""
        if self.language_name == "python":
            query_source = """
                (call
                    function: (_) @callee
                    arguments: (argument_list) @arguments
                ) @call
                """
        elif self.language_name == "cpp":
            query_source = """
            (call_expression
                function: (_) @callee
                arguments: (argument_list) @arguments
            ) @call
            """

        query = Query(
            self.normalized_language,
            query_source,
        )

        cursor = QueryCursor(query)
        found = False
        for _, captures in cursor.matches(self.tree.root_node):
            call = captures["call"][0]
            callee = captures["callee"][0]
            arguments = captures["arguments"][0]

            callee_text = self.node_text(callee, self.source_code)

            final_name = self.final_callee_name(callee)

            if function not in {callee_text, final_name}:
                continue

            found = True

            logger.info(
                "Call found: %s",
                {
                    "call": self.node_text(call, self.source_code),
                    "callee": callee_text,
                    "arguments": self.node_text(
                        arguments,
                        self.source_code,
                    ),
                    "line": call.start_point.row + 1,
                    "column": call.start_point.column + 1,
                },
            )

        if not found:
            logger.info("No usages found for %r", function)

    def find_field_values(self, node_type: str, field_name: str):
        """Find the value of ``field_name`` on every node of ``node_type``.

        Unlike a query capture, this demonstrates navigating a grammar field
        directly with Node.child_by_field_name().
        """
        found = False
        nodes = [self.tree.root_node]

        while nodes:
            node = nodes.pop()
            nodes.extend(reversed(node.children))

            if node.type != node_type:
                continue

            field_value = node.child_by_field_name(field_name)
            if field_value is None:
                continue

            found = True
            logger.info(
                "%s.%s: %s (line %d, column %d)",
                node_type,
                field_name,
                self.node_text(field_value, self.source_code),
                field_value.start_point.row + 1,
                field_value.start_point.column + 1,
            )

        if not found:
            logger.info(
                "No %r fields found on %r nodes",
                field_name,
                node_type,
            )

    def read_file(self, filename: str):
        with open(filename, "rb") as file:
            self.source_code = file.read()
        self.tree = self.parser.parse(self.source_code)

    def dump_tree(self):
        with tempfile.TemporaryFile(mode="w+b", buffering=0) as output:
            self.tree.print_dot_graph(output)

            output.seek(0)
            graph = output.read().decode("utf-8", errors="replace")

        logger.info("Tree graph:\n%s", graph.rstrip())

    def run_command(self, command: str):
        global arg_parser
        args = arg_parser.parse_args()
        if command == "simple_file":
            self.read_file(args.file)
            self.dump_tree()
        tokens = command.split()
        if tokens[0] == "read_file":
            logger.debug(f"{tokens[0]} called with {tokens[1]}")
            self.read_file(tokens[1])
        if tokens[0] == "dump_tree":
            self.dump_tree()
        if tokens[0] == "get_usages":
            self.get_usages(tokens[1])
        if tokens[0] == "find_field_values":
            self.find_field_values(tokens[1], tokens[2])

    def get_commands(self) -> dict[str, Any]:
        return {
            "simple_file": {},
            "read_file": {},
            "dump_tree": {},
            "get_usages": {},
            "find_field_values": {},
        }

    def get_positional_path_args(self) -> dict[str, set[int]]:
        return {"read_file": {0}}

    def get_path_options(self) -> set[str]:
        return {
            "read_file",
        }


def simple_file():
    global arg_parser
    args = arg_parser.parse_args()

    tree_parser = SingleFileParser(args)
    tree_parser.read_file(args.file)
    tree_parser.dump_tree()