bundle.py

#!/usr/bin/env python3

from concurrent.futures import ProcessPoolExecutor
from multiprocessing import cpu_count
from argparse import ArgumentParser
import functools
import os
from sqlite_utils import Database

from lib.json_folder_map import json_folder_map
from lib.calculate_terms import calculate_terms
from lib.regress_course import regress_course
from lib.load_courses import load_some_courses
from lib.save_term import save_term
from lib.database import create_schema
from lib.database import insert_course
from lib.database import tracer
from lib.paths import COURSE_DATA
from lib.log import log
from lib.paths import term_clbid_mapping_path


def list_all_course_index_files():
    for file in os.listdir(term_clbid_mapping_path):
        if file.startswith('.'):
            continue
        yield int(file.split('.')[0])


def one_term(args, term):
    pretty_term = f'{str(term)[:4]}:{str(term)[4]}'

    log(pretty_term, 'Loading courses')
    courses = list(load_some_courses(term))

    if args.legacy:
        [regress_course(c) for c in courses]

    log(pretty_term, 'Saving term')
    for f in args.format:
        if f == 'sqlite':
            continue
        save_term(term, courses, kind=f, root_path=args.out_dir)


def strip_build_indexes(db):
    """Drop text indexes used only for deduplication during build."""
    for idx in ['idx_description_text_text', 'idx_name_text_text',
                'idx_title_text_text', 'idx_notes_text_text']:
        db.execute(f'DROP INDEX IF EXISTS {idx}')
    db.execute('VACUUM')


def build_database(path, courses, should_trace=False):
    """Rebuild the catalog from scratch at `path`."""
    if os.path.exists(path):
        os.remove(path)

    db = Database(path, tracer=tracer if should_trace else None)
    create_schema(db)
    with db.conn:
        for course in courses:
            insert_course(db, course)
    strip_build_indexes(db)
    return db



def resolve_terms(term_or_year):
    """The terms to bundle, as a list.

    Both sources are generators, and the sqlite pass walks the terms a second
    time after the format passes have already consumed them once.
    """
    if term_or_year:
        return list(calculate_terms(term_or_year))
    return list(list_all_course_index_files())


def run(args):
    terms = resolve_terms(args.term_or_year)
    edit_one_term = functools.partial(one_term, args)

    if args.workers > 1:
        with ProcessPoolExecutor(max_workers=args.workers) as pool:
            list(pool.map(edit_one_term, terms))
    else:
        list(map(edit_one_term, terms))

    if 'sqlite' in args.format:
        from datetime import date
        current_year = date.today().year

        log('sqlite', 'Building catalog.db (all data)')
        courses = (c for term in terms for c in load_some_courses(term))
        build_database(os.path.join(args.out_dir, 'catalog.db'),
                       courses,
                       should_trace=args.trace)

        recent_cutoff = current_year - 5
        recent_terms = [t for t in terms if int(str(t)[:4]) >= recent_cutoff]
        log('sqlite', f'Building catalog-recent.db ({recent_cutoff}+)')
        courses = (c for term in recent_terms for c in load_some_courses(term))
        build_database(os.path.join(args.out_dir, 'catalog-recent.db'),
                       courses,
                       should_trace=args.trace)

    if set(args.format) & {'json', 'csv', 'xml'}:
        json_folder_map(root=args.out_dir, folder='terms', name='info')


def main():
    argparser = ArgumentParser(description='Bundle SIS term data.')
    argparser.allow_abbrev = False

    argparser.add_argument('term_or_year',
                           type=int,
                           nargs='*',
                           help='Terms (or entire years) for which to request data from the SIS')
    argparser.add_argument('-w',
                           metavar='WORKERS',
                           type=int,
                           default=cpu_count(),
                           dest='workers',
                           help='The number of operations to perform in parallel')
    argparser.add_argument('--legacy',
                           action='store_true',
                           help="Use legacy mode (you don't need this)")
    argparser.add_argument('--out-dir',
                           nargs='?',
                           action='store',
                           default=COURSE_DATA,
                           help='Where to put info.json and terms/')
    argparser.add_argument('--format',
                           action='append',
                           nargs='?',
                           choices=['json', 'csv', 'xml', 'sqlite'],
                           help='Change the output filetype')
    argparser.add_argument('--trace',
                           action='store_true',
                           help="Verbose tracing of sqlite queries")    

    args = argparser.parse_args()
    args.format = ['json'] if not args.format else args.format

    run(args)


if __name__ == '__main__':
    try:
        main()
    except KeyboardInterrupt:
        pass