import itertools from django.db.models import Count, Min from symposion.schedule.models import Room, Slot, SlotRoom class TimeTable(object): def __init__(self, day): self.day = day self.slots_qs = Slot.objects.filter(day=day)\ .select_related('kind', 'content_ptr__speaker__user')\ .prefetch_related('content_ptr__additional_speakers') self._times = sorted( set(itertools.chain(*self.slots_qs.values_list("start", "end"))) ) def slots(self): if not hasattr(self, '_slots'): filters = { "room_count": Count("slotroom"), "order": Min("slotroom__room__order") } self._slots = self.slots_qs.annotate(**filters).order_by("order") return self._slots def rooms(self): if not hasattr(self, '_rooms'): self._rooms = Room.objects\ .filter(slotroom__slot__in=self.slots_qs)\ .distinct()\ .order_by("order") return self._rooms def __iter__(self): slots = self.slots() row = [] total_room_count = self.rooms().count() for time, next_time in pairwise(self._times): row = {"time": time, "end": next_time, "slots": []} row_slots = [ slot for slot in slots if slot.start == time] for slot in row_slots: slot.rowspan = TimeTable.rowspan(self._times, slot.start, slot.end) slot.colspan = slot.room_count if not slot.exclusive else total_room_count row["slots"].append(slot) if row["slots"] or next_time is None: yield row @staticmethod def rowspan(times, start, end): return times.index(end) - times.index(start) def pairwise(iterable): a, b = itertools.tee(iterable) next(b) return itertools.zip_longest(a, b)