Fix a lot of bugs

Author: cedricrupb

code_graph/__init__.py

@@ -3,12 +3,11 @@
 from .graph import CodeGraph
 
 from .pylang   import pylang_analyses
-from .javalang import javalang_analyses
+from .javalang import javalang_analyses, java_preprocess
 
 
 def codegraph(source_code, lang = "guess", analyses = None, **kwargs):
-    tokens = ctok.tokenize(source_code, lang = lang, **kwargs)
-    root_node = _root_node(tokens)
+    root_node, tokens = preprocess_code(source_code, lang, **kwargs)
 
     graph_analyses = load_lang_analyses(tokens[0].config.lang)
 
@@ -32,14 +31,31 @@ def load_lang_analyses(lang):
     if lang == 'java'  : return javalang_analyses()
 
     raise NotImplementedError("Language %s is not supported" % lang)
-    
+
+
+def preprocess_code(source_code, lang, **kwargs):
+
+    if lang == "java": return java_preprocess(source_code, **kwargs)
+
+    return default_preprocess(source_code, lang, **kwargs)
+
+
+def default_preprocess(source_code, lang, **kwargs):
+    tokens = ctok.tokenize(source_code, lang = lang, **kwargs)
+    root_node = _root_node(tokens)
+    return root_node, tokens 
 
 # Helper methods --------------------------------
 
 def _root_node(tokens):
     if len(tokens) == 0: raise ValueError("Empty program has no root node")
 
-    base_token  = tokens[0]
+    i = 0
+    while not hasattr(tokens[i], "ast_node"):
+        i += 1
+
+    base_token = tokens[i]
+
     current_ast = base_token.ast_node 
 
     while current_ast.parent is not None:

code_graph/ast.py

@@ -15,4 +15,4 @@ def visit(self, ast_node):
 
         prev_sibling = ast_node.prev_sibling
         if prev_sibling is not None:
-            graph.add_relation(prev_sibling, ast_node, "sibling")
+            graph.add_relation(prev_sibling, ast_node, "sibling", no_create=True)

code_graph/graph.py

@@ -16,11 +16,16 @@ 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.token_nodes = []
+
         prev_token = self._add_token(tokens[0])
+        self.token_nodes.append(prev_token)
+
         for token in tokens[1:]:
             token_node = self._add_token(token)
             prev_token.add_successor(token_node, "next_token")
             prev_token = token_node
+            self.token_nodes.append(token_node)
 
         self.root_node = self.add_or_get_node(root_node)
 
@@ -70,7 +75,12 @@ def add_or_get_node(self, node):
         except Exception:
             raise ValueError("Cannot add or get node %s. Only AST nodes can be indexed." % str(node))
 
-    def add_relation(self, source_node, target_node, relation = "ast"):
+    def add_relation(self, source_node, target_node, relation = "ast", no_create = False):
+
+        if no_create:
+            if not self.has_node(source_node): return
+            if not self.has_node(target_node): return
+
         source_node = self.add_or_get_node(source_node)
         target_node = self.add_or_get_node(target_node)
         source_node.add_successor(target_node, relation)
@@ -207,7 +217,8 @@ def __hash__(self):
 class TokenNode(SyntaxNode):
 
     def __init__(self, token):
-        super().__init__(token.ast_node)
+        token_node = token.ast_node if hasattr(token, 'ast_node') else None
+        super().__init__(token_node)
         self.token = token
 
     def node_name(self):
@@ -216,7 +227,7 @@ def node_name(self):
     def __hash__(self):
         if self.ast_node is not None:
             return hash(node_key(self.ast_node))
-        return self.token.text
+        return hash(self.token.text)
 
 
 

code_graph/javalang/__init__.py

@@ -3,9 +3,59 @@
 from .cfg import ControlFlowVisitor
 from .dataflow import DataFlowVisitor
 
+import code_tokenize as ctok
+
+from code_tokenize.tokens import TokenSequence
+
+
 def javalang_analyses():
     return {
         "ast": ASTRelationVisitor,
         "cfg": ControlFlowVisitor,
         "dataflow": DataFlowVisitor,
-    }
+    }
+
+
+# Preprocessor ------------------------------------------------
+
+def _try_tokenize_or_wrap(source_code, **kwargs):
+    try:
+        custom_args = {k: v for k, v in kwargs.items() if k != "syntax_error"}
+        return ctok.tokenize(
+            source_code, lang = "java", syntax_error = "raise", **custom_args
+        ), False
+    except SyntaxError:
+        source_code = "public class Test {%s}" % source_code
+        return ctok.tokenize(source_code, lang = "java", **kwargs), True
+
+
+def java_preprocess(source_code, **kwargs):
+    tokens, has_wrapped = _try_tokenize_or_wrap(source_code, **kwargs)
+
+    if not has_wrapped:
+        return _root_node(tokens), tokens
+
+    output_tokens = TokenSequence(tokens[4:-1])
+    
+    root_node = output_tokens[0].ast_node
+    while root_node.parent is not None:
+        root_node = root_node.parent
+        if root_node.type == "method_declaration": break
+    
+    return root_node, output_tokens
+        
+        
+def _root_node(tokens):
+    if len(tokens) == 0: raise ValueError("Empty program has no root node")
+
+    base_token  = tokens[0]
+    current_ast = base_token.ast_node 
+    
+    while current_ast.parent is not None:
+        current_ast = current_ast.parent
+
+    # If root only has one child skip to child
+    if current_ast.child_count == 1:
+        current_ast = current_ast.children[0]
+    
+    return current_ast

code_graph/pylang/dataflow.py

@@ -88,6 +88,9 @@ def record_write(self, node):
 
     def record_read(self, node):
         node  = self.graph.add_or_get_node(node)
+
+        assert hasattr(node, "token"), "Expected to read from a token, but got: %s" % node
+
         qname = self.qualname(node.token.text)
 
         for last_read in self._last_reads[qname]:
@@ -385,6 +388,9 @@ def visit_lambda(self, node):
 
         self._restore_rw_context(rw_context)
         return False
+    
+    def visit_string(self, node):
+        return False # Currently we do not support f-strings
 
 
 # Helper --------------------------------------------------------

code_graph/visitor.py

@@ -156,6 +156,7 @@ def __call__(self, root_node):
 # Helper --------------------------------
 
 def node_equal(n1, n2):
+    if n1 == n2: return True
     try:
         return (n1.type == n2.type 
                     and n1.start_point == n2.start_point