proc-rota/shifts.py

import datetime
import itertools
from typing import List, Tuple

ShiftName = str
DayStr = str
WeekInt = int

days = ["Mon", "Tue", "Wed", "Thu", "Fri", "Sat", "Sun"]

# weeks = [i for i in range(1,weeks_to_rota+1)]

# weeks_days_product = list(itertools.product(weeks, days))

sites = ("truro", "exeter", "torquay", "barnstaple", "plymouth")


class SingleShift(object):
    """Class to hold all details for a shift"""

    def __init__(
        self,
        sites,
        name,
        length,
        shift_days,
        balance_by_site=True,
        workers_required=1,
        rota_on_nwds=False,
    ):
        self.site = sites
        self.name = name
        self.length = length
        self.shift_days = shift_days
        self.balance_by_site = balance_by_site
        self.workers_required = workers_required
        self.rota_on_nwds = rota_on_nwds


class ShiftCollection(object):
    """Class to hold and manipulate shifts"""

    def __init__(self, weeks_to_rota=26):
        self.shifts = []  # type List[SingleShift]

        self.weeks = [i for i in range(1, weeks_to_rota + 1)]

        self.weeks_days_product = list(itertools.product(self.weeks, days))

        self.start_date = datetime.date(2020, 9, 7)
        self.rota_days_length = len(self.weeks) * 7
        self.rota_end_date = self.start_date + datetime.timedelta(self.rota_days_length)

    def add_shift(self, shift):
        """Add a shift to the collection

        :param SingleShift shift: Shift object

        """
        self.shifts.append(shift)
        self.BuildShifts()

    def add_shifts(self, *shifts: SingleShift):
        """Add multiple shifts
        
        Returns:
            None
        """
        self.shifts.extend(shifts)
        self.BuildShifts()

    def GetShiftNamesByDay(self, day: DayStr):
        """Returns the shifts required for a specific day

        Returns:
            set: Set of ShiftName

        """
        return self.day_shifts_dict[day]

    def GetShiftByName(self, name):
        """

        :param name: 

        """
        return self.shifts_by_name[name]

    def GetShiftLengthByName(self, name):
        """Get the length of a shift
        
        Args:
            name (str): Name of a shift
        
        Returns:
            int: Length of the shift
        """
        return self.shifts_by_name[name].length

    def BuildShifts(self):
        """ """
        self.shifts_by_name = {}
        self.shift_names = []  # type: List[ShiftName]

        self.day_shifts_dict = {day: set() for day in days}

        for s in self.shifts:
            self.shifts_by_name[s.name] = s
            self.shift_names.append(s.name)

            for day in s.shift_days:
                self.day_shifts_dict[day].add(s.name)

        self.week_day_shift_product = []
        self.week_day_shiftclass_product = []
        for week, day in self.weeks_days_product:
            for s in self.shifts:
                if day in s.shift_days:
                    self.week_day_shift_product.append((week, day, s.name))
                    self.week_day_shiftclass_product.append((week, day, s))

        self.day_shift_product = []
        self.day_shiftclass_product = []
        for s in self.shifts:
            for day in days:
                if day in s.shift_days:
                    self.day_shift_product.append((day, s.name))
                    self.day_shiftclass_product.append((day, s))

    def GetAllDayShiftsAsNames(self):
        """Returns a list of all required day / shift combinations
        
        Returns:
            list: list of in the following format (str->day, str->shift)
        """
        return self.day_shift_product

    def GetAllDayShiftsAsClass(self) -> List[Tuple[DayStr, SingleShift]]:
        """Returns a list of all required day / shift combinations
        
        Returns:
            list: list of in the following format (str->day, shiftObject->shift)
        """
        return self.day_shiftclass_product

    def GetAllShiftsAsNames(self) -> List[Tuple[WeekInt, DayStr, ShiftName]]:
        """Returns a list of all possible week / day / shift combinations
        
        Returns:
            list: contains tuple of all possible week / day / shift combinations
        """
        return self.week_day_shift_product

    def GetAllShiftsAsClass(self) -> List[SingleShift]:
        """ """
        return self.week_day_shiftclass_product

    def GetAllWeeksDays(self) -> list:
        """Returns a list of all week / day tuple combinations

        Returns:
            list: list of possible week day combinations
                    [(1, "Mon"), (1, "Tue"), ... (n, "Fri")]
        """
        return self.weeks_days_product

    def GetShiftOptions(self) -> List[SingleShift]:
        """Returns a list of all the registered shifts
        
        Returns:
            List[SingleShift]: list of registered shifts (as SingleShift)
        """
        return self.shifts

    def GetShiftNames(self) -> List[ShiftName]:
        """Returns a list of all the registered shift names
        
        Returns:
            List[ShiftName]: List of names of all available shifts
        """ """ """
        return self.shift_names

    def GetWeeksInBlocks(self, block_length):
        """Gets a two dimensional list of week blocks in specified length

        e.g. block_length = 4
            [
                [ 1, 2, 3, 4],
                [ 2, 3, 4, 5],
                ...
            ]
        
        Args:
            block_length (int): length of blocks to create
        
        Returns:
            list: two dimensional list containing weeks in blocks
        """
        blocks = []
        for i in range(len(self.weeks)):
            blocks.append(self.weeks[i : i + block_length])
        return blocks

    def GetRequiredShiftsWorkersAndSites(self):
        """Returns a list of all possible shifts, the workers required and site
        
        Returns:
            list: list (week, day, shift name, workers required, shift site)
        """
        l = []
        for week, day, shift in self.week_day_shiftclass_product:
            if day in shift.shift_days:
                l.append((week, day, shift.name, shift.workers_required, shift.site))

        return l

    def GetNotRequiredShifts(self):
        """Returns a set of all possible shifts combinations that are not required

        Includes those on days
        
        Returns:
            list: (week, day, shift)
        """
        l = set()

        for week in self.weeks:
            for day in days:
                for shift in self.shifts:
                    l.add((week, day, shift))

        return l - set(self.GetAllShiftsAsClass())