AST based code diff

Author: cedricrupb

code_diff/__init__.py

@@ -0,0 +1,51 @@
+from .ast import parse_ast
+
+
+# Main method --------------------------------------------------------
+
+def difference(source, target, lang = "guess", **kwargs):
+    
+    source_ast = parse_ast(source, lang = lang, **kwargs)
+    target_ast = parse_ast(target, lang = lang, **kwargs)
+
+    # Concretize Diff
+    source_ast, target_ast = diff_search(source_ast, target_ast)
+
+    return ASTDiff(source_ast, target_ast)
+
+
+# Diff Search --------------------------------------------------------
+# Run BFS until we find a node with at least two diffs
+
+def diff_search(source_ast, target_ast):
+    if source_ast is None or source_ast.isomorph(target_ast): return None, None
+
+    queue = [(source_ast, target_ast)]
+    while len(queue) > 0:
+        source_node, target_node = queue.pop(0)
+
+        if len(source_node.children) != len(target_node.children):
+            return (source_node, target_node)
+        
+        next_children = []
+        for i, source_child in enumerate(source_node.children):
+            target_child = target_node.children[i]
+
+            if not source_child.isomorph(target_child): 
+                next_children.append((source_child, target_child))
+        
+        if len(next_children) == 1:
+            queue.append(next_children[0])
+        else:
+            return (source_node, target_node)
+
+
+
+
+# AST Difference --------------------------------------------------------
+
+class ASTDiff:
+
+    def __init__(self, source_ast, target_ast):
+        self.source_ast = source_ast
+        self.target_ast = target_ast

code_diff/ast.py

@@ -0,0 +1,131 @@
+import code_tokenize as ct
+
+from collections import defaultdict
+
+# AST Node ----------------------------------------------------------------
+
+
+class ASTNode(object):
+
+    def __init__(self, type, text = None, parent = None, children = None):
+
+        # Basic node attributes
+        self.type = type
+        self.children = children if children is not None else []
+        self.parent   = parent
+        self.text     = text   # If text is not None, then leaf node
+
+        # Tree based attributes
+        self.subtree_hash      = None
+        self.subtree_height    = 0
+        self.subtree_weight    = 1
+
+    def isomorph(self, other):
+        return ((self.subtree_hash, self.type, self.subtree_height, self.subtree_weight) == 
+                    (other.subtree_hash, other.type, other.subtree_height, other.subtree_weight))
+
+
+    def sexp(self):
+        name = self.text if self.text is not None else self.type
+
+        child_sexp = []
+        for child in self.children:
+            text = child.sexp()
+            text = ["  " + t for t in text.splitlines()]
+            child_sexp.append("\n".join(text))
+        
+        if len(child_sexp) == 0:
+            return name
+
+        return "%s {\n%s\n}" % (name, " ".join(child_sexp))
+        
+    def __repr__(self):
+        attrs = {"type": self.type, "text": self.text}
+        return "ASTNode(%s)" % (", ".join(["%s=%s" % (k, v) for k, v in attrs.items() if v is not None]))
+
+
+def default_create_node(type, children, text = None):
+    new_node = ASTNode(type, text = text, children = children)
+
+    # Subtree metrics
+    height = 0
+    weight = 1
+    hash_str = []
+
+    for child in children:
+        child.parent = new_node # Set parent relation
+        height       = max(child.subtree_height + 1, height)
+        weight      += child.subtree_weight
+        hash_str.append(str(child.subtree_hash))
+    
+    new_node.subtree_height = height
+    new_node.subtree_weight = weight
+
+    # WL hash subtree representation
+    base_str = new_node.type if new_node.text is None else new_node.text
+    hash_str.insert(0, base_str)
+    hash_str = "_".join(hash_str)
+    new_node.subtree_hash = hash(hash_str)
+
+    return new_node
+
+
+def _node_key(node):
+    return (node.type, node.start_point, node.end_point)
+
+
+class TokensToAST:
+
+    def __init__(self, create_node_fn):
+        self.create_node_fn = create_node_fn
+
+        self.root_node = None
+        self.waitlist = []
+        self.node_index = {}
+        self.child_count = defaultdict(int)
+
+    def _create_node(self, ast_node, text = None):
+        node_key = _node_key(ast_node)
+        children = [self.node_index[_node_key(c)] for c in ast_node.children
+                     if _node_key(c) in self.node_index]
+
+        current_node = self.create_node_fn(ast_node.type, children, text = text)
+        self.node_index[node_key] = current_node
+
+        # Add parent if ready
+        if ast_node.parent:
+            parent_ast = ast_node.parent
+            parent_key = _node_key(parent_ast)
+            self.child_count[parent_key] += 1
+
+            if len(parent_ast.children) == self.child_count[parent_key]:
+                self.waitlist.append(parent_ast)
+
+        else:
+            self.root_node = current_node
+
+
+    def __call__(self, tokens):
+        
+        token_nodes = ((t.text, t.ast_node) for t in tokens if hasattr(t, "ast_node"))
+        for token_text, token_ast in token_nodes:
+            self._create_node(token_ast, text = token_text)
+
+        while len(self.waitlist) > 0:
+            self._create_node(self.waitlist.pop(0))
+    
+        return self.root_node
+
+    
+
+# Interface ----------------------------------------------------------------
+
+def parse_ast(source_code, lang = "guess", **kwargs):
+    
+    # Parse AST 
+    kwargs["lang"] = lang
+    kwargs["syntax_error"] = "ignore"
+
+    ast_tokens = ct.tokenize(source_code, **kwargs)
+    
+    return TokensToAST(default_create_node)(ast_tokens)

code_diff/diff_utils.py

@@ -0,0 +1,88 @@
+import re
+
+
+# Diff parsing -----------------------------------------------------------------
+
+class Hunk:
+    
+    def __init__(self, lines, added_lines, rm_lines):
+        self.lines       = lines
+        self.added_lines = set(added_lines)
+        self.rm_lines    = set(rm_lines)
+        
+        
+    @property
+    def after(self):
+        
+        alines = []
+        
+        for i, line in enumerate(self.lines):
+            if i in self.rm_lines: continue
+            if i in self.added_lines:
+                alines.append(" " + line[1:])
+            else:
+                alines.append(line)
+                
+        return "".join(alines)
+        
+        
+    @property
+    def before(self):
+        
+        alines = []
+        
+        for i, line in enumerate(self.lines):
+            if i in self.added_lines: continue
+            if i in self.rm_lines:
+                alines.append(" " + line[1:])
+            else:
+                alines.append(line)
+                
+        return "".join(alines)
+        
+    def __repr__(self):
+        return "".join(self.lines)
+
+    
+def _parse_hunk(lines, start, end):
+    
+    hunk_lines = lines[start + 1:end]
+     
+    added_lines = []
+    rm_lines    = []
+    
+    for i, hline in enumerate(hunk_lines):
+        if hline.startswith("+"): added_lines.append(i)
+        if hline.startswith("-"): rm_lines.append(i)
+    
+    return Hunk(hunk_lines, added_lines, rm_lines)
+    
+
+hunk_pat = re.compile("@@ -(\d+)(,\d+)? \+(\d+)(,\d+)? @@.*")
+        
+def parse_hunks(diff):
+    lines = diff.splitlines(True)
+    
+    hunks = []
+    
+    start_ix = -1
+    end_ix   = -1
+    
+    for line_ix, line in enumerate(lines):
+        
+        if hunk_pat.match(line):
+            
+            end_ix = line_ix - 1
+            
+            if start_ix >= 0 and start_ix < end_ix: 
+                hunks.append(_parse_hunk(lines, start_ix, end_ix))
+            
+            start_ix = line_ix
+    
+    end_ix = len(lines)
+    
+    if start_ix >= 0 and start_ix < end_ix: 
+        hunks.append(_parse_hunk(lines, start_ix, end_ix))
+                
+    return hunks
+