import tree_sitter from tree_sitter import Language, Parser Language.build_library("./build/my-languages.so", ['./tree-sitter-glsl']) GLSL_LANGUAGE = Language('./build/my-languages.so', 'glsl') parser = Parser() parser.set_language(GLSL_LANGUAGE) def replace_function(old_func_node, new_func_node): """ replaces the old function node with the new function node """ tree = give_tree(old_func_node) old_func_start, old_func_end = node_str_idx(old_func_node) # new_func_start, new_func_end = node_str_idx(new_func_node) new_code = tree.text[:old_func_start] + new_func_node.text + tree.text[old_func_end:] return new_code def get_root(node): """ returns the root node the tree of the given node (recursively) """ if node.parent is None: return node else: return get_root(node.parent) def node_str_idx(node): """ returns the character index of start and end of a node """ whole_text = get_root(node).text.decode() start_idx = line_chr2char(whole_text, node.start_point[0], node.start_point[1]) end_idx = line_chr2char(whole_text, node.end_point[0], node.end_point[1]) return start_idx, end_idx def give_tree(func_node): """ return the tree where this function node is in """ return parser.parse(func_node.parent.text) #really no better way? def parse_functions(in_code): """ returns all functions in the code as their actual nodes. includes any comment made directly after the function definition or diretly after #copilot trigger """ tree = parser.parse(bytes(in_code, "utf8")) funcs = [n for n in tree.root_node.children if n.type == "function_definition"] return funcs def get_docstrings(func_node): """ returns the docstring of a function node """ docstring = "" for node in func_node.child_by_field_name("body").children: if node.type == "comment" or node.type == "{": docstring += node.text.decode() + "\n" else: return docstring return docstring def grab_before_comments(func_node): """ returns the comments that happen just before a function node """ precomment = "" last_comment_line = 0 for node in func_node.parent.children: #could you optimize where to iterated from? directon? if node.start_point[0] != last_comment_line + 1: precomment = "" if node.type == "comment": precomment += node.text.decode() + "\n" last_comment_line = node.start_point[0] elif node == func_node: return precomment return precomment def line_chr2char(text, line_idx, chr_idx): """ returns the character index at the given line and character index. """ lines = text.split("\n") char_idx = 0 for i in range(line_idx): char_idx += len(lines[i]) + 1 char_idx += chr_idx return char_idx