improved docstring extraction

utils/__init__.py

@@ -1,8 +1,8 @@
-from .tree_utils import (parse_functions, get_docstrings, grab_before_comments, line_chr2char, replace_function, get_root, node_str_idx, give_tree)
+from .tree_utils import (parse_functions, get_docstrings, grab_before_comments, line_chr2char, replace_function, get_root, node_str_idx, give_tree, full_func_head, has_docstrings)
 from .html_utils import (make_iframe, make_script, construct_embed)
 from .generation import (combine_generation_kwargs, stream_generation, construct_model_context)
 
-tree_funcs = ["parse_functions", "get_docstrings", "grab_before_comments", "line_chr2char", "replace_function", "get_root", "node_str_idx", "give_tree"]
+tree_funcs = ["parse_functions", "get_docstrings", "grab_before_comments", "line_chr2char", "replace_function", "get_root", "node_str_idx", "give_tree", "full_func_head", "has_docstrings"]
 html_funcs = ["make_iframe", "make_script", "construct_embed"]
 gen_funcs = ["combine_generation_kwargs", "stream_generation", "construct_model_context"]
 

utils/generation.py

@@ -1,6 +1,6 @@
 from transformers import TextIteratorStreamer
 from threading import Thread
-from .tree_utils import get_docstrings, grab_before_comments
+from .tree_utils import full_func_head, grab_before_comments
 
 def combine_generation_kwargs(temperature=2.0, max_new_tokens=512, top_p=0.95, repetition_penalty=1.2):
     """
@@ -47,11 +47,7 @@ def construct_model_context(func_node, prompt="") -> str:
     """
     Constructs the model context from a function node.
     """
-    model_context = func_node.child_by_field_name("type").text.decode() + " " + func_node.child_by_field_name("declarator").text.decode() #func_start_idx:body_start_idx?
-    docstring = get_docstrings(func_node) #might be empty?
-    if docstring:
-        model_context = model_context + "\n" + docstring
-    model_context = grab_before_comments(func_node) + model_context #prepend comments
+    model_context = grab_before_comments(func_node) + full_func_head(func_node) # (identifier + docstrings)
     if prompt != "":
         model_context = "//Title: " + prompt + "\n" + model_context #prepend user prompt/title
         model_context = "//Language: Shadertoy GLSL fragment shader\n" + model_context #prepend system prompt, language hint

utils/tree_utils.py

@@ -56,13 +56,28 @@ 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
+    for node in func_node.children:
+        if node.type == "comment": #comment in like the declarator
+            docstring += node.text.decode()
+        elif node.type == "compound_statement": #body below here
+            for body_node in node.children:
+                if body_node.type == "comment" or body_node.type == "{":
+                    docstring += " " * body_node.start_point[1] #add in indentation
+                    docstring += body_node.text.decode() + "\n" 
+                else:
+                    return docstring
     return docstring
 
+def full_func_head(func_node):
+    """
+    returns function head including docstrings before any real body code
+    """
+    cursor = func_node.child_by_field_name("body").walk()
+    cursor.goto_first_child()
+    while cursor.node.type == "comment" or cursor.node.type == "{":
+        cursor.goto_next_sibling()
+    end = cursor.node.start_point
+    return "\n".join(func_node.text.decode().split("\n")[:(end[0]-func_node.start_point[0])])
 
 def grab_before_comments(func_node):
     """
@@ -80,6 +95,13 @@ def grab_before_comments(func_node):
             return precomment
     return precomment
 
+def has_docstrings(func_node):
+    """
+    returns whether a function node has a docstring 
+    """
+    return get_docstrings(func_node).strip() != "{" or grab_before_comments(func_node) != ""
+
+
 def line_chr2char(text, line_idx, chr_idx):
     """
     returns the character index at the given line and character index.