diff --git a/cycler.py b/cycler.py index 8400444..f9191eb 100644 --- a/cycler.py +++ b/cycler.py @@ -558,3 +558,32 @@ def _cycler(label, itr): itr = (v[lab] for v in itr) return Cycler._from_iter(label, itr) + + +def cycler_with_tail(cyl, tail): + '''After the cycle is exhausted continue to yield tail + + + Parameters + ---------- + cyl : Cycler + The cycler to iterate through + + tail : dict + The dictionary to yield. Must have the same keys as ``cyl``. + + Yields + ------ + sty : dict + ''' + tk = set(tail.keys()) + if cyl.keys != tk: + raise RuntimeError('The cycler has keys {} and the tail {}. The ' + 'different keys are {}'.format(cyl.keys, + tk, tk ^ cyl.keys)) + + for sty in cyl: + yield sty + + while True: + yield tail diff --git a/test_cycler.py b/test_cycler.py index 52f65ec..b86f7d6 100644 --- a/test_cycler.py +++ b/test_cycler.py @@ -2,7 +2,7 @@ import six from six.moves import zip, range -from cycler import cycler, Cycler, concat +from cycler import cycler, Cycler, concat, cycler_with_tail import pytest from itertools import product, cycle, chain from operator import add, iadd, mul, imul @@ -341,3 +341,18 @@ def test_contains(): assert 'a' in ab assert 'b' in ab + + +def test_tail(): + a = cycler('a', range(3)) + tail = {'a': 4} + + cy_with_tail = cycler_with_tail(a, tail) + for j in range(3): + next(cy_with_tail) + + for j in range(5): + assert {'a': 4} == next(cy_with_tail) + + with pytest.raises(RuntimeError): + next(cycler_with_tail(a, {'b': 1}))