from django.core.management.base import BaseCommand, CommandError from django.db import transaction from django.utils import timezone from edt.models import Timetable, Group, Room, Course from bs4 import BeautifulSoup from ._private import delete_courses_in_week, get_events, get_weeks, get_xml, Week import datetime import requests @transaction.atomic def process_timetable_week(timetable, year, week, soup, weeks_in_soup): delete_courses_in_week(timetable, year, week) for name, type_, groups, rooms, notes, begin, end in get_events(soup, weeks_in_soup, year, week, timetable): course = Course.objects.create(timetable=timetable, begin=begin, end=end) course.name = name course.type = type_ course.notes = notes course.groups.add(*groups) if rooms is not None: course.rooms.add(*rooms) course.save() def process_timetable(timetable, year, weeks): soup = get_xml(timetable.url) weeks_in_soup = get_weeks(soup) for week in weeks: process_timetable_week(timetable, year, week, soup, weeks_in_soup) class Command(BaseCommand): help = "Fetches registered celcat timetables" def add_arguments(self, parser): parser.add_argument("--week", type=int, choices=range(1, 54), nargs="+") parser.add_argument("--year", type=int, nargs=1) def handle(self, *args, **options): year = None if options["week"] is None: _, week, day = timezone.now().isocalendar() if day >= 6: year, week, _ = (timezone.now() + datetime.timedelta(weeks=1)).isocalendar() weeks = [week] else: weeks = options["week"] if options["year"] is None and year is None: year = timezone.now().year elif year is None: year = options["year"][0] for timetable in Timetable.objects.all(): self.stdout.write("Processing {0}".format(timetable)) try: process_timetable(timetable, year, weeks) except Exception as e: self.stderr.write(self.style.error("Failed to process {0}: {1}".format(timetable, e))) self.stdout.write(self.style.SUCCESS("Done."))