Add cleaning script + Improve AST parsing

Author: Cedric Richter

code_diff/__init__.py

@@ -15,6 +15,9 @@ def difference(source, target, lang = "guess", **kwargs):
     source_ast = parse_ast(source, lang = lang, **kwargs)
     target_ast = parse_ast(target, lang = lang, **kwargs)
 
+    if source_ast is None or target_ast is None:
+        raise ValueError("Source / Target AST seems to be empty: %s" % source)
+
     # Concretize Diff
     source_ast, target_ast = diff_search(source_ast, target_ast)
 

code_diff/ast.py

@@ -164,6 +164,100 @@ def __call__(self, tokens):
 
         return self.root_node
 
+
+
+class BottomUpParser:
+
+    def __init__(self, create_node_fn):
+        
+        self.create_node_fn = create_node_fn
+        self.waitlist    = [] # Invariant: All children have been processed
+        self.open_index  = {} 
+        self.node_index  = {} # Nodes that have been processed
+
+    def _should_ignore(self, node):
+        return node.type == "comment"
+
+    def _add_to_waitlist(self, node):
+        if self._should_ignore(node): return
+
+        node_key = _node_key(node)
+
+        if node_key not in self.node_index and node_key not in self.open_index:
+            self.open_index[node_key] = node
+            self.waitlist.append(node)
+
+
+    def _init_lists(self, tokens):
+        
+        for token in tokens:
+            if hasattr(token, 'ast_node'):
+                ast_node = token.ast_node
+                if self._should_ignore(ast_node): continue
+                self.open_index[_node_key(ast_node)] = ast_node
+                self._create_node(ast_node, token.text)
+
+        if ast_node is None: return
+
+        # Get to root
+        root = ast_node
+        while root.parent is not None:
+            root = root.parent
+
+        self._open_descandents(root)
+
+        return root
+
+    def _open_descandents(self, node):
+
+        queue = [node]
+        while len(queue) > 0:
+            current_node = queue.pop(0)
+
+            has_opened = False
+            for child in current_node.children:
+                if _node_key(child) not in self.node_index:
+                    has_opened = True
+                    queue.append(child)
+            
+            if not has_opened: 
+                self._add_to_waitlist(current_node)
+
+
+    def _open_parent(self, ast_node):
+        parent = ast_node.parent
+
+        if all(_node_key(c) in self.node_index for c in parent.children if not self._should_ignore(c)):
+            self._add_to_waitlist(parent)
+
+    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]
+
+        position = (ast_node.start_point, ast_node.end_point)
+        current_node = self.create_node_fn(ast_node.type, children, text = text, position = position)
+        current_node.backend = ast_node
+
+        self.node_index[node_key] = current_node
+        del self.open_index[node_key]
+        
+        if ast_node.parent: self._open_parent(ast_node)
+
+
+    def __call__(self, tokens):
+        root_node = self._init_lists(tokens)
+
+        while len(self.waitlist) > 0:
+            self._create_node(self.waitlist.pop(0))
+        
+        if _node_key(root_node) not in self.node_index:
+            return None
+
+        return self.node_index[_node_key(root_node)]
+
+
 
 
 # Interface ----------------------------------------------------------------
@@ -177,4 +271,4 @@ def parse_ast(source_code, lang = "guess", **kwargs):
 
     ast_tokens = ct.tokenize(source_code, **kwargs)
 
-    return TokensToAST(default_create_node)(ast_tokens)
+    return BottomUpParser(default_create_node)(ast_tokens)

code_diff/diff_utils.py

@@ -86,3 +86,55 @@ def parse_hunks(diff):
 
     return hunks
 
+
+# Diff cleaning --------------------------------
+
+def _has_incomplete_comment(lines):
+    is_incomplete2 = False
+    is_incomplete1 = False
+
+    for line in lines:
+        count2 = line.count("\"\"\"")
+        if count2 % 2 == 1: is_incomplete2 = not is_incomplete2
+        
+        count1 = line.count("\'\'\'")
+        if count1 % 2 == 1: is_incomplete1 = not is_incomplete1
+
+    return is_incomplete1 or is_incomplete2
+
+
+def _determine_incomplete_comment(lines):
+    last_incomplete2 = -1
+    last_incomplete1 = -1
+
+    for i, line in enumerate(lines):
+        count2 = line.count("\"\"\"")
+        if count2 % 2 == 1:
+            last_incomplete2 = i if last_incomplete2 == -1 else -1
+        
+        count1 = line.count("\'\'\'")
+        if count1 % 2 == 1:
+            last_incomplete1 = i if last_incomplete1 == -1 else -1
+
+    assert last_incomplete1 != -1 or last_incomplete2 != -1
+
+    last_incomplete = last_incomplete2 if last_incomplete2 != -1 else last_incomplete1
+
+    dist_to_end = len(lines) - last_incomplete
+
+    if last_incomplete < dist_to_end:
+        return last_incomplete + 1, len(lines)
+    else:
+        return 0, last_incomplete
+
+
+def clean_hunk(hunk):
+    if not _has_incomplete_comment(hunk.lines): return hunk
+    start, end = _determine_incomplete_comment(hunk.lines)
+
+    new_lines = hunk.lines[start:end]
+    added_lines = [l - start for l in hunk.added_lines if l >= start and l < end]
+    rm_lines    = [l - start for l in hunk.rm_lines if l >= start and l < end]
+
+    return Hunk(new_lines, added_lines, rm_lines)
+

code_diff/sstubs.py

@@ -38,8 +38,8 @@ class SStubPattern(Enum):
     ADD_ATTRIBUTE_ACCESS           = 21 
 
     # If condition
-    MORE_SPECIFIC_IF               = 21
-    LESS_SPECIFIC_IF               = 22
+    MORE_SPECIFIC_IF               = 22
+    LESS_SPECIFIC_IF               = 23
 
 
 # SStub classification -------------------------------
@@ -143,7 +143,7 @@ def _to_plain_constant(text):
 
 
 def change_constant_type(source_ast, target_ast):
-    
+
     if source_ast.type == "identifier": return False
     if target_ast.type == "identifier": return False