canonicalize: don’t rename imports

1 files changed, 5 insertions, 6 deletions

diff --git a/canonicalize/__init__.py b/canonicalize/__init__.py
index 481b6e2..8b3fda0 100644
--- a/canonicalize/__init__.py
+++ b/canonicalize/__init__.py
@@ -4,7 +4,7 @@ import uuid
 
 from .transformations import *
 from .display import printFunction
-from .astTools import deepcopy
+from .astTools import deepcopy, getAllImportStatements
 
 def giveIds(a, idCounter=0):
     if isinstance(a, ast.AST):
@@ -27,10 +27,7 @@ def giveIds(a, idCounter=0):
                     setattr(a, field, child)
                 giveIds(child, idCounter)
 
-def getCanonicalForm(tree, problem_name=None, given_names=None, argTypes=None, imports=None):
-    if imports == None:
-        imports = []
-
+def getCanonicalForm(tree, problem_name=None, given_names=None, argTypes=None):
     transformations = [
         constantFolding,
 
@@ -53,6 +50,8 @@ def getCanonicalForm(tree, problem_name=None, given_names=None, argTypes=None, i
         deadCodeRemoval
     ]
 
+    imports = getAllImportStatements(tree)
+
     varmap = {}
     tree = propogateMetadata(tree, {'foo': ['int', 'int']}, varmap, [0])
 
@@ -71,7 +70,7 @@ def canonicalize(code):
     tree = ast.parse(code)
     giveIds(tree)
 
-    new_tree = getCanonicalForm(tree, problem_name='foo', given_names=['foo'])
+    new_tree = getCanonicalForm(tree, given_names=['bmi'])
     new_code = printFunction(new_tree)
 
     return new_code