Simplify canonicalize.canonicalize

1 files changed, 12 insertions, 14 deletions

diff --git a/canonicalize/__init__.py b/canonicalize/__init__.py
index 8b3fda0..7c4ef79 100644
--- a/canonicalize/__init__.py
+++ b/canonicalize/__init__.py
@@ -27,7 +27,13 @@ def giveIds(a, idCounter=0):
                     setattr(a, field, child)
                 giveIds(child, idCounter)
 
-def getCanonicalForm(tree, problem_name=None, given_names=None, argTypes=None):
+def canonicalize(code, problem_name=None, given_names=None):
+    if given_names is None:
+        given_names = []
+
+    tree = ast.parse(code)
+    giveIds(tree)
+
     transformations = [
         constantFolding,
 
@@ -50,12 +56,11 @@ def getCanonicalForm(tree, problem_name=None, given_names=None, argTypes=None):
         deadCodeRemoval
     ]
 
-    imports = getAllImportStatements(tree)
-
     varmap = {}
     tree = propogateMetadata(tree, {'foo': ['int', 'int']}, varmap, [0])
-
     tree = simplify(tree)
+
+    imports = getAllImportStatements(tree)
     tree = anonymizeNames(tree, given_names, imports)
 
     oldTree = None
@@ -64,20 +69,13 @@ def getCanonicalForm(tree, problem_name=None, given_names=None, argTypes=None):
         helperFolding(tree, problem_name, imports)
         for t in transformations:
             tree = t(tree)
-    return tree
-
-def canonicalize(code):
-    tree = ast.parse(code)
-    giveIds(tree)
-
-    new_tree = getCanonicalForm(tree, given_names=['bmi'])
-    new_code = printFunction(new_tree)
 
+    new_code = printFunction(tree)
     return new_code
         
 
 if __name__ == '__main__':
-        code = '''
+    code = '''
 def isPositive(x):
     return x > 0
 
@@ -91,4 +89,4 @@ def foo(x, y):
         return x
     return -x
 '''
-        print(canonicalize(code))
+    print(canonicalize(code))