From 3148aaf7c43866bd672d54ac54a0e70bc71f1020 Mon Sep 17 00:00:00 2001
From: Alban Gruin
Date: Sat, 18 Nov 2017 17:14:26 +0100
Subject: Mise en commun du traitement des données avant rendu de l’emploi du
 temps

---
 models.py | 13 +++++++++----
 1 file changed, 9 insertions(+), 4 deletions(-)

(limited to 'models.py')

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() \
-- 
cgit v1.2.1