1 files changed, 7 insertions, 5 deletions
diff --git a/management/commands/cleancourses.py b/management/commands/cleancourses.py
index f6041ef..246cfcc 100644
--- a/management/commands/cleancourses.py
+++ b/management/commands/cleancourses.py
@@ -15,22 +15,24 @@
 
 from django.core.management.base import BaseCommand
 from django.db import transaction
-from edt.models import Course, Group
+
+from ...models import Course, Group
 
 
 class Command(BaseCommand):
     help = "Remove all courses and groups from the database"
 
     def add_arguments(self, parser):
-        parser.add_argument("--timetable", type=int, nargs="+")
+        parser.add_argument("--source", type=int, nargs="+")
 
     def handle(self, *args, **options):
         with transaction.atomic():
-            if options["timetable"] is None:
+            if options["source"] is None:
                 Course.objects.all().delete()
                 Group.objects.all().delete()
             else:
-                Course.objects.filter(timetable__id__in=options["timetable"]).delete()
-                Group.objects.filter(timetable__id__in=options["timetable"]).delete()
+                Course.objects.filter(source__id__in=options["source"]) \
+                              .delete()
+                Group.objects.filter(source__id__in=options["source"]).delete()
 
         self.stdout.write(self.style.SUCCESS("Done."))