Change parsing format to use AST via RedBaron

Moved away from using manual parsing. Source is parsed using RedBaron to make the FST. Some import formats not yet implemented, and still some bugs to work out. However, the current script works well for the two cases in the file comments.

scripts/fix_ext_import.py

@@ -1,20 +0,0 @@
-from lib2to3.fixer_base import BaseFix
-from lib2to3.fixer_util import Name, syms
-
-
-class FixExtImport(BaseFix):
-
-    PATTERN = "fixnode='oldname'"
-
-    def transform(self, node, results):
-        fixnode = results['fixnode']
-        fixnode.replace(Name('newname', prefix=fixnode.prefix))
-
-        if node.type == syms.import_from and \
-                getattr(results['imp'], 'value', None) == 'flask.ext':
-            return 0
-            # TODO: Case 2
-
-
-# CASE 1 - from flask.ext.foo import bam --> from flask_foo import bam
-# CASE 2 - from flask.ext import foo --> import flask_foo as foo

scripts/flaskext_migrate.py

@@ -5,10 +5,41 @@ from redbaron import RedBaron
 import sys
 
 
-with open("test.py", "r") as source_code:
+def read_source(input_file):
-    red = RedBaron(source_code.read())
+    with open(input_file, "r") as source_code:
+        red = RedBaron(source_code.read())
+    return red
 
-print red.dumps()
 
-# with open("code.py", "w") as source_code:
+def write_source(red, input_file):
-#     source_code.write(red.dumps())
+    with open(input_file, "w") as source_code:
+        source_code.write(red.dumps())
+
+
+def fix_imports(red):
+    from_imports = red.find_all("FromImport")
+    for x in range(len(from_imports)):
+        values = from_imports[x].value
+        if (values[0].value == 'flask') and (values[1].value == 'ext'):
+            # Case 1
+            if len(from_imports[x].value) == 3:
+                package = values[2].value
+                modules = from_imports[x].modules()
+                r = "{}," * len(modules)
+                print modules
+                from_imports[x].replace("from flask_%s import %s"
+                                        % (package, r.format(*modules)[:-1]))
+            # Case 2
+            else:
+                module = from_imports[x].modules()[0]
+                from_imports[x].replace("import flask_%s as %s"
+                                        % (module, module))
+
+    return red
+
+
+if __name__ == "__main__":
+    input_file = sys.argv[1]
+    ast = read_source(input_file)
+    new_ast = fix_imports(ast)
+    write_source(new_ast, input_file)

scripts/test.py

@@ -1,6 +1,18 @@
-from flask.ext.foo import bam
+from flask.ext.foo import \
+    bam, \
+    crackle
+
+from flask.ext import foo
+
+from flask.ext.foo import (bam,
+                           a,
+                           b
+)
+
+
 from flask.ext import foo
 
+import sys
 
 def migrate(old_file):
     new_file = open("temp.py", "w")
@@ -16,5 +28,6 @@ def migrate(old_file):
 
         new_file.write(line)
 
+
 if __name__ == "__main__":
     old_file = open(sys.arv[1])