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