initdb.py 2.98 KB
#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import csv
import argparse
import re
import string
import sys
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker

from models import Base, Student, Test, Question

# SIIUE names have alien strings like "(TE)" and are sometimes capitalized
# We remove them so that students dont keep asking what it means
def fix(name):
    return string.capwords(re.sub('\(.*\)', '', name).strip())

# ===========================================================================
#   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('--demo', default='students.db', type=str, help='add demo students')
argparser.add_argument('csvfile', nargs='?', type=str, default='', help='CSV filename')
args = argparser.parse_args()

# ===========================================================================
engine = create_engine('sqlite:///{}'.format(args.db), echo=False)

# Criate schema if needed
Base.metadata.create_all(engine)

Session = sessionmaker(bind=engine)

# --- start session ---
try:
    session = Session()

    # add administrator
    session.add(Student(id='0', name='Professor', password=''))

    # add students
    if args.csvfile:
        # from csv file if available
        try:
            csvreader = csv.DictReader(open(args.csvfile, encoding='iso-8859-1'), delimiter=';', quotechar='"', skipinitialspace=True)
        except EnvironmentError:
            print('Error: CSV file "{0}" not found.'.format(args.csvfile))
            session.rollback()
            sys.exit(1)
        else:
            session.add_all([Student(id=r['N.º'], name=fix(r['Nome']), password='') for r in csvreader])
    else:
        # otherwise add a few fake students
        fakes = [
            ['1888', 'Fernando Pessoa'],
            ['1799', 'Almeida Garrett'],
            ['1845', 'José Maria Eça de Queirós'],
            ['1465', 'Gil Vicente'],
            ['1825', 'Camilo Castelo Branco'],
            ['1842', 'Antero de Quental'],
            ['1907', 'Miguel Torga'],
            ['1810', 'Alexandre Herculano'],
            ['1924', 'Alexandre O\'Neill'],
        ]
        session.add_all([Student(id=i, name=name, password='') for i,name in fakes])

    session.commit()

except Exception:
    print('Error: Database already exists.')
    session.rollback()
    sys.exit(1)

else:
    n = session.query(Student).count()
    print('New database created: {0}\n{1} users inserted:'.format(args.db, n))

    users = session.query(Student).order_by(Student.id).all()
    print('    {0:8} - {1} (administrator)'.format(users[0].id, users[0].name))
    if n > 1:
        print('    {0:8} - {1}'.format(users[1].id, users[1].name))
    if n > 3:
        print('    .\n    .\n    .')
    if n > 2:
        print('    {0:8} - {1}'.format(users[-1].id, users[-1].name))

# --- end session ---