1 files changed, 9 insertions, 4 deletions
diff --git a/models.py b/models.py
index 547c0e1..369a4a6 100644
--- a/models.py
+++ b/models.py
@@ -157,10 +157,15 @@ class Room(models.Model):
 
 
 class CourseManager(Manager):
-    def get_courses_for_group(self, group, **criteria):
-        return self.get_queryset() \
-                   .filter(groups__in=Group.objects.get_parents(group), **criteria) \
-                   .order_by("begin").prefetch_related("rooms")
+    def get_courses(self, obj, **criteria):
+        qs = self.get_queryset()
+        if isinstance(obj, Group):
+            qs = qs.filter(groups__in=Group.objects.get_parents(obj), **criteria) \
+                   .prefetch_related("rooms")
+        elif isinstance(obj, Room):
+            qs = qs.filter(rooms__in=(obj,), **criteria)
+
+        return qs.order_by("begin")
 
     def get_weeks(self, **criteria):
         return self.get_queryset() \