From 7a8b602aae0328e1d14692c214d1094c9fbcb926 Mon Sep 17 00:00:00 2001
From: Simon Drabble <sdrabble@scriptainsights.com>
Date: Tue, 28 Nov 2023 13:57:15 -0700
Subject: [PATCH] Allow IntervalSchedule to use a custom schedule type.

---
 django_celery_beat/models.py | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/django_celery_beat/models.py b/django_celery_beat/models.py
index 4f3e1578..3da655a1 100644
--- a/django_celery_beat/models.py
+++ b/django_celery_beat/models.py
@@ -146,6 +146,8 @@ class IntervalSchedule(models.Model):
     >>> every=2, period=DAYS
     """
 
+    _schedule_cls_ = schedules.schedule
+
     DAYS = DAYS
     HOURS = HOURS
     MINUTES = MINUTES
@@ -176,7 +178,7 @@ class Meta:
 
     @property
     def schedule(self):
-        return schedules.schedule(
+        return self._schedule_cls_(
             timedelta(**{self.period: self.every}),
             nowfun=lambda: make_aware(now())
         )
@@ -185,7 +187,9 @@ def schedule(self):
     def from_schedule(cls, schedule, period=SECONDS):
         every = max(schedule.run_every.total_seconds(), 0)
         try:
-            return cls.objects.get(every=every, period=period)
+            obj = cls.objects.get(every=every, period=period)
+            obj._schedule_cls_ = type(schedule)
+            return obj
         except cls.DoesNotExist:
             return cls(every=every, period=period)
         except MultipleObjectsReturned: