refactor on occurence generator

llazzaro · May 4, 2016 · c57215f · c57215f
1 parent 3ed432d
commit c57215f
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 13 deletions.
diff --git a/schedule/models/events.py b/schedule/models/events.py
@@ -231,15 +231,14 @@ def _occurrences_after_generator(self, after=None):
         if rule is None:
             if self.end > after:
                 yield self._create_occurrence(self.start, self.end)
-            raise StopIteration
+            return
         date_iter = iter(rule)
         difference = self.end - self.start
         loop_counter = 0
-        while True:
-            o_start = next(date_iter)
+        for o_start in date_iter:
             o_start = tzinfo.localize(o_start)
-            if self.end_recurring_period and o_start > self.end_recurring_period:
-                raise StopIteration
+            if self.end_recurring_period and self.end_recurring_period and o_start > self.end_recurring_period:
+                break
             o_end = o_start + difference
             if o_end > after:
                 yield self._create_occurrence(o_start, o_end)
@@ -257,15 +256,11 @@ def occurrences_after(self, after=None, max_occurences=None):
         occ_replacer = OccurrenceReplacer(self.occurrence_set.all())
         generator = self._occurrences_after_generator(after)
         trickies = list(self.occurrence_set.filter(original_start__lte=after, start__gte=after).order_by('start'))
-        while True:
-            try:
-                nxt = next(generator)
-            except StopIteration:
-                nxt = None
+        for index, nxt in enumerate(generator):
+            if max_occurences and index > max_occurences - 1:
+                break
             if (len(trickies) > 0 and (nxt is None or nxt.start > trickies[0].start)):
                 yield trickies.pop(0)
-            if (nxt is None):
-                raise StopIteration
             yield occ_replacer.get_occurrence(nxt)
 
     @property

diff --git a/tests/test_event.py b/tests/test_event.py
@@ -372,7 +372,7 @@ def test_occurences_with_recurrent_event_no_end_recurring_period_max_loop(self):
                             'Non recurring event test get_occurrence',
                             start,
                             start + datetime.timedelta(hours=1),
-                            None,
+                            start + datetime.timedelta(hours=10),
                             rule,
                             cal)
         event.save()