Change wrapper logic for testing

This commit is contained in:
Keyan Pishdadian 2015-02-11 14:53:39 -05:00
parent 4a9e34a574
commit 96d81c297a
1 changed files with 3 additions and 4 deletions

View File

@ -83,10 +83,9 @@ def fix_standard_imports(red):
return red
def fix(input_file):
ast = read_source(input_file)
new_ast = fix_imports(ast)
write_source(new_ast, input_file)
def fix(ast):
"""Wrapper which allows for testing when not running from shell"""
return fix_imports(ast).dumps()
if __name__ == "__main__":
input_file = sys.argv[1]