CFG and dataflow pass finished for Python; Do we need symbols, docs, etc?

Author: cedricrupb

code_graph/__init__.py

@@ -0,0 +1,51 @@
+import code_tokenize as ctok
+
+from .graph import CodeGraph
+
+from .ast import ASTRelationVisitor
+from .cfg import ControlFlowVisitor
+from .dataflow import DataFlowVisitor
+
+
+GRAPH_ANALYSES = {
+    "ast": ASTRelationVisitor,
+    "cfg": ControlFlowVisitor,
+    "dataflow": DataFlowVisitor,
+}
+
+def codegraph(source_code, lang = "guess", analyses = None, **kwargs):
+    tokens = ctok.tokenize(source_code, lang = lang, **kwargs)
+    root_node = _root_node(tokens)
+
+    if analyses is None:
+        analyses = GRAPH_ANALYSES.keys()
+    else:
+        assert all(a in GRAPH_ANALYSES for a in analyses), \
+                "Not all analyses are supported. Available analyses are: %s" % ", ".join(GRAPH_ANALYSES.keys())
+    
+    graph = CodeGraph(root_node, tokens, lang = lang)
+    
+    for analysis in analyses:
+        analysis_visitor = GRAPH_ANALYSES[analysis]
+        analysis_visitor(graph)(root_node)
+
+    return graph
+
+    
+
+# Helper methods --------------------------------
+
+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/ast.py

@@ -0,0 +1,18 @@
+from .visitor import ASTVisitor
+
+
+class ASTRelationVisitor(ASTVisitor):
+
+    def __init__(self, graph):
+        super().__init__()
+        self.graph = graph
+
+    def visit(self, ast_node):
+        graph = self.graph
+
+        for child in ast_node.children:
+            graph.add_relation(ast_node, child, "child")
+        
+        prev_sibling = ast_node.prev_sibling
+        if prev_sibling is not None:
+            graph.add_relation(prev_sibling, ast_node, "sibling")

code_graph/cfg.py

@@ -0,0 +1,161 @@
+from .visitor import ASTVisitor
+
+class ControlFlowVisitor(ASTVisitor):
+    
+    def __init__(self, graph):
+        super().__init__()
+        self.graph = graph
+        self._last_stmts = tuple()
+
+        self._break_from    = []
+        self._continue_from = []
+        self._returns_from  = []
+        self._yields_from   = []
+
+
+    def _add_next(self, stmt_node):
+        for last_stmt_node in self._last_stmts:
+            self.graph.add_relation(last_stmt_node, stmt_node, "controlflow")
+        self._last_stmts = (stmt_node,)
+
+    def _reset_last_stmts(self, reset_target):
+        last_stmts = self._last_stmts
+        self._last_stmts = (reset_target,)
+        return last_stmts
+
+    def visit_block(self, node):
+        
+        for stmt in node.children:
+            self.walk(stmt)
+
+        return False
+
+    def visit_function_definition(self, node):
+        outside_last, self._last_stmts = self._last_stmts, (node,)
+        outside_returns, outside_yields = self._returns_from, self._yields_from
+        self._returns_from, self._yields_from = [], []
+
+        self.walk(
+            node.child_by_field_name("body")
+        )
+
+        for stmt in self._last_stmts:
+            self.graph.add_relation(stmt, node, "return_from")
+
+        for stmt in self._returns_from:
+            self.graph.add_relation(stmt, node, "return_from")
+
+        for stmt in self._yields_from:
+            self.graph.add_relation(stmt, node, "yield_from")
+
+        self._returns_from, self._yields_from = outside_returns, outside_yields
+        self._last_stmts = outside_last
+        return False
+
+    def visit_if_statement(self, node):
+        self._add_next(node)
+
+        # Consequences
+        self.walk(node.child_by_field_name("consequence"))
+        left_last_stmts = self._reset_last_stmts(node)
+
+        # Alternative
+        self.walk(node.child_by_field_name("alternative"))
+        right_last_stmts = self._reset_last_stmts(node)
+
+        self._last_stmts = left_last_stmts + right_last_stmts
+        return False
+
+    def visit_return_statement(self, node):
+        self._add_next(node)
+        self._returns_from.append(node)
+        self._last_stmts = tuple()
+        return False
+
+    def visit_yield_statement(self, node):
+        self._add_next(node)
+        self._yields_from.append(node)
+        self._last_stmts = tuple()
+        return False
+
+    def visit_break_statement(self, node):
+        self._add_next(node)
+        self._break_from.append(node)
+        self._last_stmts = tuple()
+        return False
+
+    def visit_continue_statement(self, node):
+        self._add_next(node)
+        self._continue_from.append(node)
+        self._last_stmts = tuple()
+        return False
+
+    def visit_for_statement(self, node):
+
+        prev_break, prev_continue = self._break_from, self._continue_from
+        self._break_from, self._continue_from = [], []
+
+        self._add_next(node)
+        self.walk(node.child_by_field_name("body"))
+        self._last_stmts += tuple(self._continue_from)
+        self._add_next(node)
+
+        self.walk(node.child_by_field_name("alternative"))
+
+        self._last_stmts += tuple(self._break_from)
+
+        self._break_from, self._continue_from = prev_break, prev_continue
+        return False
+
+    def visit_while_statement(self, node):
+
+        prev_break, prev_continue = self._break_from, self._continue_from
+        self._break_from, self._continue_from = [], []
+
+        self._add_next(node)
+        self.walk(node.child_by_field_name("body"))
+        self._last_stmts += tuple(self._continue_from)
+        self._add_next(node)
+
+        self.walk(node.child_by_field_name("alternative"))
+
+        self._last_stmts += tuple(self._break_from)
+
+        self._break_from, self._continue_from = prev_break, prev_continue
+        return False
+    
+    def visit_try_statement(self, node):
+        self._add_next(node)
+        starting_stmt = self._last_stmts
+
+        self.walk(node.child_by_field_name("body"))
+        self.walk(node.child_by_field_name("alternative"))
+
+        exception_starting_stmts = self._last_stmts + starting_stmt
+        self._last_stmts = exception_starting_stmts
+        out_last_stmts = tuple()
+
+        finally_clauses = []
+        for possible_exception in node.children:
+            if possible_exception.type == "except_clause":
+                self.walk(possible_exception)
+                out_last_stmts += self._last_stmts
+                self._last_stmts = exception_starting_stmts
+            if possible_exception.type == "finally_clause":
+                finally_clauses.append(possible_exception)
+
+        self._last_stmts += out_last_stmts
+
+        for finally_clause in finally_clauses:
+            self.walk(finally_clause)
+
+        return False
+
+    def visit(self, node):
+        # All statement node type end with statement
+        # in tree-sitter
+        # Therefore, we can savely do this hack.
+        if node.type.endswith("statement"):
+            self._add_next(node)
+            return False
+        return True