Runtime tests / Debug and prepare for packaging

Author: cedricrupb

code_graph/graph.py

@@ -16,13 +16,13 @@ def __init__(self, root_node, tokens, lang = "python"):
         self._ast_nodes       = {} # Nodes indexed by an AST node
         self._anonymous_nodes = [] # Unindexed nodes, can only be indexed by traversal
 
-        self.root_node = self.add_node(root_node)
-
         prev_token = self._add_token(tokens[0])
         for token in tokens[1:]:
             token_node = self._add_token(token)
             prev_token.add_successor(token_node, "next_token")
             prev_token = token_node
+        
+        self.root_node = self.add_or_get_node(root_node)
 
     # Helper methods --------------------------------
 
@@ -61,6 +61,9 @@ def add_node(self, node):
         return self._add_ast_node(node)
 
     def add_or_get_node(self, node):
+        if isinstance(node, SyntaxNode):
+            return self.add_or_get_node(node.ast_node)
+
         if isinstance(node, Node): return node
         try:
             return self._add_ast_node(node)
@@ -74,6 +77,9 @@ def add_relation(self, source_node, target_node, relation = "ast"):
 
     # API GET methods-----------------------------------------
 
+    def has_node(self, ast_node):
+        return node_key(ast_node) in self._ast_nodes
+
     def nodes(self):
         return chain(self._ast_nodes.values(), self._anonymous_nodes)
 
@@ -88,6 +94,17 @@ def todot(self, file_name = None, edge_colors = None):
                 dotwriter.run(f)
                 f.seek(0)
                 return f.read()
+
+    def tokens_only(self):
+        """
+        Computes a graph containing only tokens
+
+        Any edges of inner nodes will be propagated down to leaves.
+        The first token related to an inner node acts as an representant
+
+        """
+        return graph_to_tokens_only(self)
+         
 
     # Internal GET methods -----------------------------------
 
@@ -290,4 +307,67 @@ def escape(token):
 
         # Cleanup
         for src_node in self.graph:
-            del src_node._dot_node_id
+            del src_node._dot_node_id
+
+
+# Propagate to leaves ----------------------------------------------------------------
+
+def _compute_representer(graph):
+    representer = {}
+
+    root_node = graph.root_node
+    queue     = [root_node]
+
+    while len(queue) > 0:
+        current_node = queue.pop(-1)
+        
+        path = []
+        while not hasattr(current_node, "token"):
+            path.append(current_node)
+            syntax_node    = current_node.ast_node
+            children       = [graph.add_or_get_node(c)
+                                for c in syntax_node.children 
+                                if graph.has_node(c)]
+            if len(children) == 0: break
+            first, *others = children
+            queue.extend(others)
+            current_node = first
+
+        for r in path: representer[r] = current_node
+
+    return representer
+
+
+SYNTAX_TYPES = {"child", "sibling"}
+
+def graph_to_tokens_only(graph):
+    representers = _compute_representer(graph)
+    tokens       = graph.tokens
+
+    output = CodeGraph(tokens[0].ast_node, tokens, lang = graph.lang)
+
+    for token in graph.tokens:
+        if not hasattr(token, "ast_node"): continue
+        token_node  = graph.add_or_get_node(token.ast_node)
+        output_node = output.add_or_get_node(token.ast_node)
+
+        for _, edge_type, successor in token_node.successors():
+            if edge_type in SYNTAX_TYPES: continue
+            if not hasattr(successor, "token"): continue
+            output_succ = output.add_or_get_node(successor.ast_node)
+            output.add_relation(output_node, output_succ, edge_type)
+            
+
+    for node, representer in representers.items():
+        output_representer = output.add_or_get_node(representer.ast_node)
+        
+        for _, edge_type, successor in node.successors():
+            if edge_type in SYNTAX_TYPES: continue
+            if successor not in representers: continue
+            successor_representer = representers[successor]
+            output_successor_representer = output.add_or_get_node(successor_representer.ast_node)
+            output.add_relation(output_representer, 
+                                output_successor_representer, 
+                                edge_type)
+
+    return output

code_graph/visitor.py

@@ -4,6 +4,24 @@ class ASTVisitor:
     def __init__(self):
         self._ast_handler = None
 
+    # Error handling ------------------------------------------------
+
+    def visit_ERROR(self, node):
+        """
+        An ERROR node is introduced if the parser reacts to an syntax error.
+
+        The subtree rooted an ERROR node might node be conventional
+        or a tree (might include cycles).
+        The walk function however assumes a tree as input and will
+        run in an infinite loop for errors.
+
+        Therefore, default strategy is to skip error nodes.
+        Can be overriden by subclasses.
+        
+        """
+        return False
+
+
     # Custom handler  ------------------------------------------------
 
     @staticmethod
@@ -18,6 +36,9 @@ def _parse_visit_pattern(name):
         # or function is the node type
         # and definition is the edge type
         # Therefore, we register both
+        # Assumption:
+        #  - No node type function with edge definition
+        #  - or no node type function_definition
 
         atomic_name = "_".join(parts[1:])
         node_name   = "_".join(parts[1:-1])
@@ -113,18 +134,36 @@ def walk(self, root_node):
             if next_node is None:
                 next_node = self._next_sibling(current_node)
 
+            previous_node = current_node
+
             # Step 4: Go up until sibling exists
             while next_node is None and current_node.parent is not None:
                 current_node = current_node.parent
-                if current_node == root_node: break
+                if node_equal(current_node, root_node): break
                 next_node    = self._next_sibling(current_node)
 
+                if node_equal(previous_node, next_node): 
+                    # A loop can occur if the program is not syntactically correct
+                    # Is this enough?
+                    next_node = None
+
             current_node = next_node
 
     def __call__(self, root_node):
         return self.walk(root_node)
 
 
+# Helper --------------------------------
+
+def node_equal(n1, n2):
+    try:
+        return (n1.type == n2.type 
+                    and n1.start_point == n2.start_point
+                    and n1.end_point   == n2.end_point)
+    except AttributeError:
+        return n1 == n2
+
+
 # Compositions ----------------------------------------------------------------
 
 class VisitorComposition(ASTVisitor):

requirements.txt

@@ -0,0 +1 @@
+code_tokenize>=0.1.0

run_test.py

@@ -0,0 +1,58 @@
+import os
+import argparse
+import json
+
+from time import time
+
+from tqdm import tqdm
+
+from glob import glob
+
+import code_graph as cg
+
+def load_examples(files):
+    for n, file in enumerate(files):
+        name = os.path.basename(file)
+        desc = "File %d / %d: %s" % (n+1, len(files), name)
+        total = sum(1 for _ in open(file, "r"))
+        with open(file, "r") as lines:
+            for line in tqdm(lines, total = total, desc = desc):
+                yield json.loads(line)
+
+def tokens_to_text(tokens):
+    method_text =  " ".join(tokens)
+    return "public class Test{\n%s\n}" % method_text
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("input_dir")
+    parser.add_argument("result_file")
+    args = parser.parse_args()
+
+    if os.path.isfile(args.input_dir):
+        files = [args.input_dir]
+    else:
+        files = glob(os.path.join(args.input_dir, "*.jsonl"))
+
+    run_times = open(args.result_file, "w")
+
+    for example in load_examples(files):
+        tokens = example["tokens"]
+        length = len(tokens)
+        code   = tokens_to_text(tokens)
+
+        start_time = time()
+
+        try:
+            graph = cg.codegraph(code, lang = "java", syntax_error = "raise")
+        except Exception as e:
+            print(e)
+            continue
+
+        end_time = time() - start_time
+        
+        run_times.write(json.dumps([length, len(graph), end_time]) + "\n")
+
+if __name__ == '__main__':
+    main()

setup.cfg

@@ -0,0 +1,3 @@
+[metadata]
+description-file = README.md
+ong_description_content_type = text/markdown

setup.py

@@ -0,0 +1,36 @@
+from setuptools import setup
+
+with open("README.md", "r") as f:
+    long_description = f.read()
+
+setup(
+  name = 'code_graph',
+  packages = ['code_graph'], 
+  version = '0.0.1', 
+  license='MIT',     
+  description = 'Fast program graph generation in Python',
+  long_description = long_description,
+  long_description_content_type="text/markdown",
+  author = 'Cedric Richter',                   
+  author_email = 'cedricr.upb@gmail.com',    
+  url = 'https://github.com/cedricrupb/code_graph',  
+  download_url = '', 
+  keywords = ['code', 'graph', 'program', 'language processing'], 
+  install_requires=[          
+          'tree_sitter',
+          'GitPython',
+          'requests',
+          'code_tokenize'
+      ],
+  classifiers=[
+    'Development Status :: 3 - Alpha',    
+    'Intended Audience :: Developers',  
+    'Topic :: Software Development :: Build Tools',
+    'License :: OSI Approved :: MIT License',
+    'Programming Language :: Python :: 3', 
+    'Programming Language :: Python :: 3.6',
+    'Programming Language :: Python :: 3.7',
+    'Programming Language :: Python :: 3.8',
+    'Programming Language :: Python :: 3.9',
+  ],
+)