Refactor python files

wgmaster/import.py

@@ -11,6 +11,8 @@ import datetime
 import csv
 
 
+DIR: str
+
 TABLES = ['branch', 'wb_gl', 'wb_kg', 'wb_rd', 'wine_attribute', 'wine_cultivation',
           'member', 'member_billing_address', 'contract', 'area_commitment',
           'season', 'modifier', 'delivery', 'delivery_part', 'delivery_part_modifier', ]
@@ -44,7 +46,7 @@ def sqlite_regexp(pattern: str, value: Optional[str]) -> Optional[bool]:
 
 
 def import_csv(cur: sqlite3.Cursor, table_name: str) -> None:
-    rows = csv.parse(f'{args.dir}/{table_name}.csv')
+    rows = csv.parse(f'{DIR}/{table_name}.csv')
     names = next(rows)
 
     sql = f'INSERT INTO {table_name} ({", ".join(names)}) VALUES ({", ".join(["?"] * len(names))})'
@@ -81,7 +83,9 @@ def check_foreign_keys(cur: sqlite3.Cursor) -> bool:
     return len(rows) == 0
 
 
-if __name__ == '__main__':
+def main() -> None:
+    global DIR
+
     parser = argparse.ArgumentParser()
     parser.add_argument('dir', type=str, metavar='DIR',
                         help='The directory where the migrated csv files are stored')
@@ -91,6 +95,8 @@ if __name__ == '__main__':
                         help='Whether the database file should be overwritten or kept')
     args = parser.parse_args()
 
+    DIR = args.dir
+
     if not args.keep:
         try:
             os.remove(args.db)
@@ -100,28 +106,32 @@ if __name__ == '__main__':
     sqlite3.register_adapter(datetime.date, lambda d: d.strftime('%Y-%m-%d'))
     sqlite3.register_adapter(datetime.time, lambda t: t.strftime('%H:%M:%S'))
 
-    DB_CNX = sqlite3.connect(args.db)
-    DB_CNX.create_function('REGEXP', 2, sqlite_regexp)
+    cnx = sqlite3.connect(args.db)
+    cnx.create_function('REGEXP', 2, sqlite_regexp)
 
     if not args.keep:
         for file_name in get_sql_files():
             with open(file_name, encoding='utf-8') as sql_file:
                 print(f'Executing {file_name}')
-                DB_CNX.executescript(sql_file.read())
+                cnx.executescript(sql_file.read())
 
     try:
-        DB_CNX.isolation_level = None
+        cnx.isolation_level = None
         # Member predecessors may refer to a higher MgNr
-        DB_CNX.execute("PRAGMA foreign_keys = OFF")
-        DB_CNX.execute("BEGIN")
+        cnx.execute("PRAGMA foreign_keys = OFF")
+        cnx.execute("BEGIN")
         for table in TABLES:
-            import_csv(DB_CNX.cursor(), table)
-        if not check_foreign_keys(DB_CNX.cursor()):
+            import_csv(cnx.cursor(), table)
+        if not check_foreign_keys(cnx.cursor()):
             raise RuntimeError('foreign key constraint failed')
-        DB_CNX.execute("COMMIT")
+        cnx.execute("COMMIT")
     except Exception as err:
-        DB_CNX.execute("ROLLBACK")
+        cnx.execute("ROLLBACK")
         raise err
     finally:
-        DB_CNX.execute("PRAGMA foreign_keys = ON")
-        DB_CNX.close()
+        cnx.execute("PRAGMA foreign_keys = ON")
+        cnx.close()
+
+
+if __name__ == '__main__':
+    main()