Merge pull request #38 from pganssle/rrule-generators

Add xafter method to rrule
dateutil · Mar 16, 2015 · 4e17b7f · 4e17b7f
2 parents acb07a2 + d98fb41
commit 4e17b7f
Show file tree

Hide file tree

Showing 2 changed files with 76 additions and 1 deletion.
diff --git a/dateutil/rrule.py b/dateutil/rrule.py
@@ -209,7 +209,48 @@ def after(self, dt, inc=False):
                     return i
         return None
 
-    def between(self, after, before, inc=False):
+    def xafter(self, dt, count=None, inc=False):
+        """
+        Generator which yields up to `count` recurrences after the given
+        datetime instance, equivalent to `after`.
+
+        :param dt:
+            The datetime at which to start generating recurrences.
+
+        :param count:
+            The maximum number of recurrences to generate. If `None` (default),
+            dates are generated until the recurrence rule is exhausted.
+
+        :param inc:
+            If `dt` is an instance of the rule and `inc` is `True`, it is
+            included in the output.
+
+        :yields: Yields a sequence of `datetime` objects.
+        """
+
+        if self._cache_complete:
+            gen = self._cache
+        else:
+            gen = self
+
+        # Select the comparison function
+        if inc:
+            comp = lambda dc, dtc: dc >= dtc
+        else:
+            comp = lambda dc, dtc: dc > dtc
+
+        # Generate dates
+        n = 0
+        for d in gen:
+            if comp(d, dt):
+                yield d
+
+                if count is not None:
+                    n += 1
+                    if n >= count:
+                        break
+
+    def between(self, after, before, inc=False, count=1):
         """ Returns all the occurrences of the rrule between after and before.
         The inc keyword defines what happens if after and/or before are
         themselves occurrences. With inc=True, they will be included in the

diff --git a/dateutil/test/test.py b/dateutil/test/test.py
@@ -2715,6 +2715,40 @@ def testAfterInc(self):
                                .after(datetime(1997, 9, 4, 9, 0), inc=True),
                          datetime(1997, 9, 4, 9, 0))
 
+    def testXAfter(self):
+        self.assertEqual(list(rrule(DAILY,
+                                    dtstart=datetime(1997, 9, 2, 9, 0))
+                                    .xafter(datetime(1997, 9, 8, 9, 0), count=12)),
+                                    [datetime(1997, 9, 9, 9, 0),
+                                     datetime(1997, 9, 10, 9, 0),
+                                     datetime(1997, 9, 11, 9, 0),
+                                     datetime(1997, 9, 12, 9, 0),
+                                     datetime(1997, 9, 13, 9, 0),
+                                     datetime(1997, 9, 14, 9, 0),
+                                     datetime(1997, 9, 15, 9, 0),
+                                     datetime(1997, 9, 16, 9, 0),
+                                     datetime(1997, 9, 17, 9, 0),
+                                     datetime(1997, 9, 18, 9, 0),
+                                     datetime(1997, 9, 19, 9, 0),
+                                     datetime(1997, 9, 20, 9, 0)])
+
+    def testXAfterInc(self):
+        self.assertEqual(list(rrule(DAILY,
+                                    dtstart=datetime(1997, 9, 2, 9, 0))
+                                    .xafter(datetime(1997, 9, 8, 9, 0), count=12, inc=True)),
+                                    [datetime(1997, 9, 8, 9, 0),
+                                     datetime(1997, 9, 9, 9, 0),
+                                     datetime(1997, 9, 10, 9, 0),
+                                     datetime(1997, 9, 11, 9, 0),
+                                     datetime(1997, 9, 12, 9, 0),
+                                     datetime(1997, 9, 13, 9, 0),
+                                     datetime(1997, 9, 14, 9, 0),
+                                     datetime(1997, 9, 15, 9, 0),
+                                     datetime(1997, 9, 16, 9, 0),
+                                     datetime(1997, 9, 17, 9, 0),
+                                     datetime(1997, 9, 18, 9, 0),
+                                     datetime(1997, 9, 19, 9, 0)])
+
     def testBetween(self):
         self.assertEqual(rrule(DAILY,
                                #count=5,