From 96d81c297adddf034c8eef6b2c76cbd9181d6d22 Mon Sep 17 00:00:00 2001 From: Keyan Pishdadian Date: Wed, 11 Feb 2015 14:53:39 -0500 Subject: [PATCH] Change wrapper logic for testing --- scripts/flaskext_migrate.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/scripts/flaskext_migrate.py b/scripts/flaskext_migrate.py index b7a1baf0..de7290e4 100644 --- a/scripts/flaskext_migrate.py +++ b/scripts/flaskext_migrate.py @@ -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]