diff --git a/initdb_from_csv.py b/initdb_from_csv.py index 41d0622..124fbb5 100755 --- a/initdb_from_csv.py +++ b/initdb_from_csv.py @@ -1,6 +1,7 @@ #!/usr/bin/env python3 # -*- coding: utf-8 -*- +import sys import sqlite3 import csv import argparse @@ -10,6 +11,7 @@ import string import re # SIIUE names have alien strings like "(TE)" and are sometimes capitalized +# We remove them so that students dont keep asking what it means def fixname(s): return string.capwords(re.sub('\(.*\)', '', s).strip()) @@ -47,25 +49,47 @@ sql_cmd = '''PRAGMA foreign_keys = ON; # --------- Parse command line options ----------- argparser = argparse.ArgumentParser(description='Create new database from a CSV file (SIIUE format)') argparser.add_argument('--db', default='students.db', type=str, help='database filename') -argparser.add_argument('csvfile', type=str, help='CSV filename') +argparser.add_argument('csvfile', nargs='?', type=str, default='', help='CSV filename') args = argparser.parse_args() +db_exists = os.path.exists(args.db) -# --------- Parse CSV -------- -with open(args.csvfile, encoding='iso-8859-1') as csvfile: # SIIUE format - reader = csv.DictReader(csvfile, delimiter=';', quotechar='"') - db_exists = os.path.exists(args.db) +with sqlite3.connect(args.db) as c: + # use existing or create new database schema + if db_exists: + print('-> Using previous database "{}"...'.format(args.db)) + else: + print('-> Creating new database "{}"...'.format(args.db)) + c.executescript(sql_cmd) - with sqlite3.connect(args.db) as c: - if not db_exists: - print('Creating new database "{}"...'.format(args.db)) - c.executescript(sql_cmd) - c.execute('INSERT INTO students VALUES (?,?,?)', ('0', 'Professor', '')) - else: - print('Database "{}" already exists.'.format(args.db)) + # get students + if args.csvfile: + csvfile = open(args.csvfile, encoding='iso-8859-1') + print('-> Using students from CSV file "{}"...'.format(args.csvfile)) + students = genstudent(csv.DictReader(csvfile, delimiter=';', quotechar='"')) + else: + print('-> Creating fake students numbered 1 to 5...'.format(args.csvfile)) + students = [ + ('1', 'Student1', ''), + ('2', 'Student2', ''), + ('3', 'Student3', ''), + ('4', 'Student4', ''), + ('5', 'Student5', '') + ] - print('Inserting students into database... ', end='') - try: - c.executemany('INSERT INTO students VALUES (?,?,?)', genstudent(reader)) - except sqlite3.IntegrityError: - print('\rStudents already in the database. Aborting!!!') + # insert students into database + print('-> Inserting students into database... ') + try: + c.executemany('INSERT INTO students VALUES (?,?,?)', students) + except sqlite3.IntegrityError: + print('** ERROR ** Students already exist. Aborted!') + sys.exit(1) + + # insert professor into database + print('-> Inserting professor (id=0)...') + try: + c.execute('INSERT INTO students VALUES (?,?,?)', ('0', 'Professor', '')) + except sqlite3.IntegrityError: + print('** WARNING ** Professor already exists.') + +print('Done.') -- libgit2 0.21.2