add: preprocessor for comments and tokens corpuses

Author: msaidow

python150k/ast_conversion.py

@@ -0,0 +1,41 @@
+def convert(ast):
+    increase_by = {}  # count of how many idx to increase the new idx by:
+    # each time there is a value node
+    cur = 0
+    for i, node in enumerate(ast):
+        increase_by[i] = cur
+        if "value" in node:
+            cur += 1
+
+    new_dp = []
+    for i, node in enumerate(ast):
+        inc = increase_by[i]
+        if "value" in node:
+            child = [i + inc + 1]
+            if "children" in node:
+                child += [n + increase_by[n] for n in node["children"]]
+            new_dp.append({"type": node["type"], "children": child})
+            new_dp.append({"value": node["value"]})
+        else:
+            if "children" in node:
+                node["children"] = [n + increase_by[n] for n in node["children"]]
+            new_dp.append(node)
+
+    # sanity check
+    children = []
+    for node in new_dp:
+        if "children" in node:
+            children += node["children"]
+    assert len(children) == len(set(children))
+    return new_dp
+
+
+def get_dfs(ast, only_leaf=False):
+    dp = []
+    for node in ast:
+        if "value" in node:
+            dp.append(str(node["value"]))
+        else:
+            if not only_leaf:
+                dp.append("<"+node["type"]+">")
+    return dp

python150k/parse_python3.py

@@ -20,9 +20,12 @@ def read_file_to_string(filename):
     f.close()
     return s
 
-def parse_file(filename):
+def parse_file(filename, input="file"):
     global c, d
-    tree = ast.parse(read_file_to_string(filename), filename)
+    if input == "file":
+        tree = ast.parse(read_file_to_string(filename), filename)
+    else:
+        tree = ast.parse(filename)
 
     json_tree = []
     def gen_identifier(identifier, node_type = 'identifier'):
@@ -56,7 +59,7 @@ def traverse(node):
         elif isinstance(node, ast.Num):
             json_node['value'] = unicode(node.n)
         elif isinstance(node, ast.Str):
-            json_node['value'] = node.s#.decode('utf-8')
+            json_node['value'] = "string" #node.s #.decode('utf-8')
         elif isinstance(node, ast.alias):
             json_node['value'] = unicode(node.name)
             if node.asname:
@@ -73,6 +76,8 @@ def traverse(node):
                 children.append(gen_identifier(n))
         elif isinstance(node, ast.keyword):
             json_node['value'] = unicode(node.arg)
+        elif isinstance(node, ast.arg):
+            children.append(gen_identifier(node.arg))
 
 
         # Process children.
@@ -88,10 +93,12 @@ def traverse(node):
             if node.orelse:
                 children.append(traverse_list(node.orelse, 'orelse'))
         elif isinstance(node, ast.With):
+            children.append(traverse_list(node.items, 'items'))
+            children.append(traverse_list(node.body, 'body'))
+        elif isinstance(node, ast.withitem):
             children.append(traverse(node.context_expr))
             if node.optional_vars:
                 children.append(traverse(node.optional_vars))
-            children.append(traverse_list(node.body, 'body'))
         elif isinstance(node, ast.Try): # Except
             children.append(traverse_list(node.body, 'body'))
             children.append(traverse_list(node.handlers, 'handlers'))
@@ -106,14 +113,14 @@ def traverse(node):
             children.append(traverse_list(node.args, 'args'))
             children.append(traverse_list(node.defaults, 'defaults'))
             if node.vararg:
-                children.append(gen_identifier(node.vararg, 'vararg'))
+                children.append(traverse(node.vararg))
             if node.kwarg:
-                children.append(gen_identifier(node.kwarg, 'kwarg'))
+                children.append(traverse(node.kwarg))
         elif isinstance(node, ast.ExceptHandler):
             if node.type:
                 children.append(traverse_list([node.type], 'type'))
             if node.name:
-                children.append(traverse_list([node.name], 'name'))
+                children.append(gen_identifier(node.name, 'name'))
             children.append(traverse_list(node.body, 'body'))
         elif isinstance(node, ast.ClassDef):
             children.append(traverse_list(node.bases, 'bases'))