Using treesitter with Npp and PythonScript plugin

query_lang.py

from Npp import notepad, editor
import ctypes
from ctypes import wintypes

from WinDialog import (
    Dialog, WS, DS,
    ListBox, TextBox
)
from WinDialog.controls.textbox import ES
from WinDialog.controls.listbox import LBS
from WinDialog.win_helper import GetWindowRect

from tree_sitter_languages import get_language, get_parser

class QueryLanguage(Dialog):
    def __init__(self, title=''):
        super().__init__(title)

        self.size = (240, 94)
        self.style = DS.SETFONT | WS.POPUP
        self.pointsize = 11
        self.typeface = 'Segoe UI'
        self.textbox = TextBox('', (240, 11), (0, 0))
        self.textbox.style = ES.LEFT | WS.CHILD | WS.VISIBLE | WS.BORDER | WS.TABSTOP
        self.textbox.onChange = self.on_change
        self.listbox = ListBox('', (240, 80), (0, 11))
        self.listbox.style = WS.VISIBLE | WS.TABSTOP |LBS.NOINTEGRALHEIGHT

        rect = self.get_scintilla_rect()
        self.position = ((((rect.right-rect.left) // 2) - self.size[0]) //2, 0)  # top centered
        # self.center = True

        self.queries = {
            'string_select': lambda: self.select_current_string(),
            'string_copy': lambda: self.select_current_string(True),
            'function_select': lambda: self.select_current_function(),
            'function_copy': lambda: self.select_current_function(True),
            'add_docstring': self.add_docstring_to_current_function,
        }
        self.filtered_queries = list(self.queries.keys())
        language__ = notepad.getLanguageName(notepad.getLangType()).lower()
        self.language = get_language(language__)
        self.tree = get_parser(language__).parse(editor.getText().encode())

        self.show()

    def get_scintilla_rect(self):
        rect = wintypes.RECT()
        GetWindowRect(editor.hwnd, ctypes.byref(rect))
        return rect

    def filter_queries(self, text):
        for query_ in self.queries:
            if query_.startswith(text):
                yield query_

    def initialize(self):
        self.listbox.addStrings(self.filtered_queries)

    def onIdOk(self):
        self.queries[self.filtered_queries[0]]()
        self.terminate()

    def on_change(self):
        text = self.textbox.getText()
        self.listbox.clear()
        if text:
            self.filtered_queries = list(self.filter_queries(text))
            self.listbox.addStrings(self.filtered_queries)
        else:
            self.listbox.addStrings(self.queries.keys())

    def select_current_string(self, copy=False):
        query = self.language.query("(string) @_")
        captures = query.captures(self.tree.root_node)
        current_pos = editor.getCurrentPos()

        for node, _ in captures:
            if node.start_byte <= current_pos <= node.end_byte:
                if copy:
                    string = editor.getText()[node.start_byte:node.end_byte]
                    editor.copyText(string)
                else:
                    editor.setSel(node.start_byte, node.end_byte)

    def select_current_function(self, copy=False):
        query = self.language.query("(function_definition) @_")
        captures = query.captures(self.tree.root_node)
        current_pos = editor.getCurrentPos()

        for node, _ in captures:
            if node.start_byte <= current_pos <= node.end_byte:
                if copy:
                    string = editor.getText()[node.start_byte:node.end_byte]
                    editor.copyText(string)
                else:
                    editor.setSel(node.start_byte, node.end_byte)

    def add_docstring_to_current_function(self):
        docstring_template = '''"""
{0}
{0}    Args:
{0}        {1}
{0}    Returns:

{0}"""
{0}'''
        indent = 4

        query = self.language.query(
        """
        (function_definition
            name: (identifier) @function
            parameters: (parameters) @function.parameters
            body: (block) @function.body)
        """
        )

        captures = query.captures(self.tree.root_node)

        functions = []
        for node, capture in captures:
            match capture:
                case "function":
                    function_start_position = editor.positionFromLine(node.start_point[0])
                    function_name = node.text.decode()
                case "function.parameters":
                    parameters = node.text.decode()
                case "function.body":
                   function_info = {
                       "name": function_name,
                       "function_start_position": function_start_position,
                       "parameters": parameters,
                       "body_start_position": editor.positionFromLine(node.start_point[0]) + node.start_point[1],
                       "indent_level": node.start_point[1] // indent,
                       "body_end_position": editor.positionFromLine(node.end_point[0]) + node.end_point[1],
                   }
                   functions.append(function_info)
                case _ :
                    continue

        current_pos = editor.getCurrentPos()
        for func in functions:
            if func['function_start_position'] <= current_pos <= func['body_end_position']:
                editor.insertText(func['body_start_position'], f'{docstring_template.format(" "*indent*func["indent_level"], func["parameters"])}')
                break


QueryLanguage()