Code cleanup

generate.py

@@ -1,9 +1,11 @@
 #!/usr/bin/env python
+import os
 import re
 import MySQLdb as mysql
-from singplur import singularize
 from itertools import combinations
+from argparse import ArgumentParser
 
+from singplur import singularize
 
 PRIMARY_KEY = 'id'
 TAB = '\t'
@@ -291,10 +293,7 @@ def partition(callback, iterable):
     return a, b
 
 
-if __name__ == '__main__':  # pragma: nocover
-    import os
-    from argparse import ArgumentParser
-
+def parse_options():
     parser = ArgumentParser(description='Generate PHPActiveRecord models.')
     parser.add_argument('dbname', help='database name')
     parser.add_argument('-H', '--host', metavar='ADDRESS', default='localhost',
@@ -333,6 +332,11 @@ if __name__ == '__main__':  # pragma: nocover
     if args.spaces:
         TAB = args.spaces * ' '
 
+    return args, options
+
+
+if __name__ == '__main__':  # pragma: nocover
+    args, options = parse_options()
     conn = mysql.connect(args.host, args.user, args.password, args.dbname)
     models = create_models(conn, options)
     conn.close()