from django.core.exceptions import ValidationError from django.db import models from django.utils.translation import gettext_lazy as _ from .program import Program __all__ = ( "Rerun", "RerunQuerySet", ) class RerunQuerySet(models.QuerySet): """Queryset for Rerun (sub)classes.""" def station(self, station=None, id=None): return ( self.filter(program__station=station) if id is None else self.filter(program__station__id=id) ) def program(self, program=None, id=None): return ( self.filter(program=program) if id is None else self.filter(program__id=id) ) def rerun(self): return self.filter(initial__isnull=False) def initial(self): return self.filter(initial__isnull=True) class Rerun(models.Model): """Abstract model offering rerun facilities. Assume `start` is a datetime field or attribute implemented by subclass. """ program = models.ForeignKey( Program, models.CASCADE, db_index=True, verbose_name=_("related program"), ) initial = models.ForeignKey( "self", models.SET_NULL, related_name="rerun_set", verbose_name=_("rerun of"), limit_choices_to={"initial__isnull": True}, blank=True, null=True, db_index=True, ) objects = RerunQuerySet.as_manager() class Meta: abstract = True @property def is_initial(self): return self.initial is None @property def is_rerun(self): return self.initial is not None def get_initial(self): """Return the initial schedule (self or initial)""" return self if self.initial is None else self.initial.get_initial() def clean(self): super().clean() if ( hasattr(self, "start") and self.initial is not None and self.initial.start >= self.start ): raise ValidationError( {"initial": _("rerun must happen after original")} ) def save_rerun(self): self.program = self.initial.program def save_initial(self): pass def save(self, *args, **kwargs): if self.initial is not None: self.initial = self.initial.get_initial() if self.initial == self: self.initial = None if self.is_rerun: self.save_rerun() else: self.save_initial() super().save(*args, **kwargs)