old_console/rotation/services/cron_utils.py

import datetime
import logging

from celery.schedules import crontab
from django.conf import settings
from django_celery_beat.models import CrontabSchedule

from rotation.enums import Period, WeekDay, Month

_log = logging.getLogger(__name__)


def get_enum_indexes(enum_choises, serializer_value, start=1):
    """ Iterate thru enum and get all index numbers that are used in serializer_value
    So if we have a enum apple = a, banana = b and serializer value is ['a', 'b'], then
    in result we will have a '1,2' string

    :param enum_choises: Enum to iterate
    :param serializer_value: Values to find
    :param start: Start counter
    :return: String like '1,2'
    """
    index = start
    result = ''
    for value, label in enum_choises:
        if value in serializer_value:
            if len(result) > 0:
                result += ',' + str(index)
            else:
                result = str(index)

        index += 1

    return result


def create_crontab(crontab_pattern=None,
                   *,
                   minute=None,
                   hour=None,
                   day_of_week=None,
                   day_of_month=None,
                   month_of_year=None):
    """ Create new crontab filled with values from crontab_pattern overriden by explicit arguments """
    if not crontab_pattern:
        crontab_pattern = crontab()
    return crontab(
        minute=minute if minute is not None else crontab_pattern._orig_minute,
        hour=hour if hour is not None else crontab_pattern._orig_hour,
        day_of_week=day_of_week if day_of_week is not None else crontab_pattern._orig_day_of_week,
        day_of_month=day_of_month if day_of_month is not None else crontab_pattern._orig_day_of_month,
        month_of_year=month_of_year if month_of_year is not None else crontab_pattern._orig_month_of_year,
    )


def crontab_to_schedule(crontab=None,
                        schedule=None,
                        *,
                        minute=None,
                        hour=None,
                        day_of_week=None,
                        day_of_month=None,
                        month_of_year=None):
    """ Fill schedule with crontab info or explicit args,
        if schedule is None new CrontabSchedule will be created """
    if not schedule:
        schedule = CrontabSchedule()
    crontab = create_crontab(crontab,
                             minute=minute,
                             hour=hour,
                             day_of_week=day_of_week,
                             day_of_month=day_of_month,
                             month_of_year=month_of_year)
    schedule.minute = crontab._orig_minute
    schedule.hour = crontab._orig_hour
    schedule.day_of_week = crontab._orig_day_of_week
    schedule.day_of_month = crontab._orig_day_of_month
    schedule.month_of_year = crontab._orig_month_of_year
    return schedule


def serializer_to_model(serializer):
    """ Convert serializer values to celery beat schedule
    :param serializer: RotationSettingsSerializer filled with data
    :return: CrontabSchedule filled with data but not saved
    """
    schedule = CrontabSchedule()
    if serializer.validated_data['schedule']['period'] == Period.DAY:
        cur = serializer.validated_data['schedule']['time']
        crontab_to_schedule(schedule=schedule, minute=str(cur.minute), hour=str(cur.hour))
    elif serializer.validated_data['schedule']['period'] == Period.WEEK:
        day_of_week = get_enum_indexes(WeekDay.choices,
                                       serializer.validated_data['schedule']['week_day'],
                                       start=0)
        # also, need to set time
        weeklycb = getattr(settings, 'WEEKLY_CRONTAB', crontab(minute='0', hour='0'))
        crontab_to_schedule(weeklycb, schedule, day_of_week=day_of_week)
    elif serializer.validated_data['schedule']['period'] == Period.MONTH:
        month_of_year = get_enum_indexes(Month.choices,
                                         serializer.validated_data['schedule']['month'])
        # also, need to set day and time
        monthlycb = getattr(settings, 'MONTHLY_CRONTAB', crontab(minute='0', hour='0', day_of_month='1'))
        crontab_to_schedule(monthlycb, schedule, month_of_year=month_of_year)

    return schedule


def schedule_to_dict(schedule):
    """ Convert CrontabSchedule model to dict that can be used for ScheduleSerializer
    :param schedule:
    :return:
    """
    data = {
        'period': '',
        'time': '',
        'week_day': [],
        'month': []
    }

    if schedule.month_of_year != '*':
        index = 1
        values = schedule.month_of_year.split(',')
        data['period'] = Period.MONTH
        for key, value in Month.choices:
            if str(index) in values:
                data['month'].append(key)
            index += 1
        return data
    if schedule.day_of_week != '*':
        index = 0
        values = schedule.day_of_week.split(',')
        data['period'] = Period.WEEK
        for key, value in WeekDay.choices:
            if str(index) in values:
                data['week_day'].append(key)
            index += 1
        return data
    # If we get here - we have a time record
    data['period'] = Period.DAY

    # we can have these options:
    # - it's a int -> just convert to int
    # - it's a * -> set 0
    # - it contains - -> set 0
    # - it contains , -> set 0
    if schedule.minute == '*' or '-' in schedule.minute or ',' in schedule.minute:
        minute = 0
    else:
        minute = int(schedule.minute)

    if schedule.hour == '*' or '-' in schedule.hour or ',' in schedule.hour:
        hour = 0
    else:
        hour = int(schedule.hour)

    data['time'] = datetime.time(minute=minute, hour=hour)
    return data