Add ability to clone schedule data from another week

tailbone/views/shifts/schedule.py

@@ -29,7 +29,7 @@ from __future__ import unicode_literals, absolute_import
 import datetime
 
 from rattail.db import model
-from rattail.time import localtime, make_utc
+from rattail.time import localtime, make_utc, get_sunday
 
 from pyramid_simpleform import Form
 
@@ -54,6 +54,12 @@ class ScheduleView(TimeSheetView):
             self.request.session.flash("Removed {} shifts from current schedule.".format(count))
             return self.redirect(self.request.route_url('schedule.edit'))
 
+        # okay then, check if we should copy data from another week
+        if self.request.method == 'POST' and self.request.POST.get('copy-week'):
+            sunday, copied = self.copy_schedule()
+            self.request.session.flash("Copied {} shifts from week of {}".format(copied, sunday.strftime('%m/%d/%Y')))
+            return self.redirect(self.request.route_url('schedule.edit'))
+
         # okay then, process filters; redirect if any were received
         form = Form(self.request, schema=ShiftFilter)
         self.process_filter_form(form)
@@ -129,6 +135,44 @@ class ScheduleView(TimeSheetView):
                 deleted += 1
         return deleted
 
+    def copy_schedule(self):
+        """
+        Clear current schedule, then copy shift data from another week.
+        """
+        try:
+            sunday = datetime.datetime.strptime(self.request.POST['copy-week'], '%m/%d/%Y').date()
+        except ValueError as error:
+            self.request.session.flash("Invalid date specified: {}: {}".format(type(error), error), 'error')
+            raise self.redirect(self.request.route_url('schedule.edit'))
+        sunday = get_sunday(sunday)
+        context = self.get_timesheet_context()
+        if sunday == context['date']:
+            self.request.session.flash("Cannot copy schedule from same week; please specify a different week.", 'error')
+            raise self.redirect(self.request.route_url('schedule.edit'))
+
+        self.clear_schedule()
+
+        copied = 0
+        if context['employees']:
+            offset = context['date'] - sunday
+            sunday = datetime.datetime.combine(sunday, datetime.time(0))
+            start_time = localtime(self.rattail_config, sunday)
+            end_time = localtime(self.rattail_config, sunday + datetime.timedelta(days=7))
+            shifts = Session.query(model.ScheduledShift)\
+                            .filter(model.ScheduledShift.employee_uuid.in_([e.uuid for e in context['employees']]))\
+                            .filter(model.ScheduledShift.start_time >= make_utc(start_time))\
+                            .filter(model.ScheduledShift.end_time < make_utc(end_time))
+            for shift in shifts:
+                Session.add(model.ScheduledShift(
+                    employee_uuid=shift.employee_uuid,
+                    store_uuid=shift.store_uuid,
+                    start_time=shift.start_time + offset,
+                    end_time=shift.end_time + offset,
+                ))
+                copied += 1
+
+        return sunday, copied
+
     @classmethod
     def defaults(cls, config):
         cls._defaults(config)