1 files changed, 39 insertions, 25 deletions
diff --git a/management/commands/timetables.py b/management/commands/timetables.py
index 37a5201..9882759 100644
--- a/management/commands/timetables.py
+++ b/management/commands/timetables.py
@@ -10,43 +10,57 @@ from ._private import delete_courses_in_week, get_events, get_weeks, get_xml, We
 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 __get_weeks(self, weeks):
-        if weeks is None:
-            year, week, day = timezone.now().isocalendar()
+    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()
-
-            yield year, week
+            weeks = [week]
         else:
-            year = timezone.now().year
-            for week in weeks:
-                yield year, week
-
-    @transaction.atomic
-    def handle(self, *args, **options):
-        for year, week in self.__get_weeks(options["week"]):
-            delete_courses_in_week(year, week)
-
-            for timetable in Timetable.objects.all():
-                self.stdout.write("Processing {0}, week {1}".format(timetable, week))
+            weeks = options["week"]
 
-                soup = get_xml(timetable.url)
-                weeks = get_weeks(soup)
-
-                for name, type_, groups, rooms, notes, begin, end in get_events(soup, weeks, week, timetable):
-                    course = Course.objects.create(name=name, type=type_, notes=notes, timetable=timetable, begin=begin, end=end)
+        if options["year"] is None and year is None:
+            year = timezone.now().year
+        else:
+            year = options["year"][0]
 
-                    course.groups.add(*groups)
-                    if rooms is not None:
-                        course.rooms.add(*rooms)
+        for timetable in Timetable.objects.all():
+            self.stdout.write("Processing {0}".format(timetable))
 
-                    course.save()
+            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."))