Enable croniter_range() to use derived class

- Add new optional 'croniter' argument to croniter_range(), making it useable
  with an alternate or derived croniter class.
- Add unittest demonstrating a simple subclass that accepts 5-part expression
  (no seconds) which can now be used with croniter_range().

2 files changed, 27 insertions, 1 deletions

diff --git a/src/croniter/croniter.py b/src/croniter/croniter.py
index d6167b8..9588647 100644
--- a/src/croniter/croniter.py
+++ b/src/croniter/croniter.py
@@ -698,7 +698,8 @@ class croniter(object):
         return (max(tdp, tdt) - min(tdp, tdt)).total_seconds() < 60
 
 
-def croniter_range(start, stop, expr_format, ret_type=None, day_or=True, exclude_ends=False):
+def croniter_range(start, stop, expr_format, ret_type=None, day_or=True, exclude_ends=False,
+                   croniter=croniter):
     """
     Generator that provides all times from start to stop matching the given cron expression.
     If the cron expression matches either 'start' and/or 'stop', those times will be returned as
diff --git a/src/croniter/tests/test_croniter.py b/src/croniter/tests/test_croniter.py
index 6eb5991..121099e 100755
--- a/src/croniter/tests/test_croniter.py
+++ b/src/croniter/tests/test_croniter.py
@@ -1258,6 +1258,31 @@ class CroniterRangeTest(base.TestCase):
         matches = list(croniter_range(datetime(2020, 9, 30), datetime(2020, 10, 30), cron, day_or=False))
         self.assertEqual(len(matches), 0)
 
+    def test_croniter_range_derived_class(self):
+        # trivial example extending croniter
+
+        class croniter_nosec(croniter):
+            """ Like croniter, but it forbids second-level cron expressions. """
+            @classmethod
+            def expand(cls, expr_format):
+                if len(expr_format.split()) == 6:
+                    raise CroniterBadCronError("Expected 'min hour day mon dow'")
+                return croniter.expand(expr_format)
+
+        cron = "0 13 8 1,4,7,10 wed"
+        matches = list(croniter_range(datetime(2020, 1, 1), datetime(2020, 12, 31), cron, day_or=False, croniter=croniter_nosec))
+        self.assertEqual(len(matches), 3)
+
+        cron = "0 1 8 1,15,L wed 15,45"
+        with self.assertRaises(CroniterBadCronError):
+            # Should fail using the custom class that forbids the seconds expression
+            croniter_nosec(cron)
+
+        with self.assertRaises(CroniterBadCronError):
+            # Should similiarly fail because it's using the custom classs too
+            i = croniter_range(datetime(2020, 1, 1), datetime(2020, 12, 31), cron, croniter=croniter_nosec)
+            next(i)
+
     def test_explicit_year_forward(self):
         start = datetime(2020, 9, 24)
         cron = "0 13 8 1,4,7,10 wed"