Python has a special primitive called a generator. A generator is an iterable that could be backed by a data structure, or some other form of deferred computation. Python has a toolkit in the standard library called itertools that provides a set of high-level functions for producing and combining iterables in special ways. This year’s Advent of Code day 16 problem involved multiplying a set of numbers by an sequence of numbers following a specific pattern:

  1. Start from a repeating pattern of [0, 1, 0, -1, …]
  2. For the given digit you are calculating, repeat each element that many times. So for the first digit, use [0, 1, 0, -1, …], for the second digit use [0, 0, 1, 1, 0, 0, -1, -1, …], for the third, [0, 0, 0, 1, 1, 1, 0, 0, 0, -1, -1, -1, …]
  3. Shift the sequence by one by discarding the first element. So the first pattern would start with [1, 0, 1, -1, …], the second with [0, 1, 1, 0, 0, -1, -1, …] and so on.

In Python, you can use itertools to implement the above requirements fairly concisely:

import itertools

def base_pattern(i):
    base = [0, 1, 0, -1]
    expanded = itertools.chain.from_iterable([e] * i for e in base)
    repeated = itertools.cycle(expanded)
    first_skipped = itertools.islice(repeated, 1, None)
    return first_skipped

list(itertools.islice(base_pattern(1), 10))
# [1, 0, -1, 0, 1, 0, -1, 0, 1, 0]
list(itertools.islice(base_pattern(2), 10))
# [0, 1, 1, 0, 0, -1, -1, 0, 0, 1]
list(itertools.islice(base_pattern(3), 10))
# [0, 0, 1, 1, 1, 0, 0, 0, -1, -1]

In Python, all the intermediate steps of composing our requirements in code take the form of generators. To view the output of our base_pattern(..) results, we have to wrap it in a list(itertools.islice(...)) in order to grab the first few elements of what is actually an infinitely-long sequence of elements.

Occasionally, you need to write code that is more low-level than the high-level tools provided by itertools and generator-comprehension syntax. For this Python lets you specify a generator much in the same way you’d write a function, but instead of return‘ing your data, you yield it.

Where a return statement in a function results in the code of the function halting and bailing out, a yield pauses the control flow in the code of a function, and will resume it when the generator needs to produce the next value in a sequence.

Let’s convert the above high-level code into something more direct by writing the generator ourselves.

# translate
def base_pattern(i):
    base = [0, 1, 0, -1]
    expanded = itertools.chain.from_iterable([e] * i for e in base)
    repeated = itertools.cycle(expanded)
    first_skipped = itertools.islice(repeated, 1, None)
    return first_skipped

Something else we can do with generators, is rely on an interactive API to create a body of code that responds to novel inputs. Passing a generator to the next(..) built-in function is how we pull an element out of a generator, and calling the .send(..) method on a generator is how we pass an element to it, causing the deferred computation inside of it to change.

Let’s begin by defining a small generator that just returns one element, and then use the Python interpreter to interact with it.

def toy(n):
    yield n

Instantiate the generator and try to pull things out of it

>>> def toy(n):
...     yield n
...
>>> t = toy(3)
>>> t
<generator object toy at 0x7f3dad6badd0>
>>> next(t)
3
>>> next(t)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
StopIteration
>>>

The StopIteration exception just means that the generator cannot advance forwards via the normal iterable API.


>>> def toy(n):
...   while True:
...     if n == 0:
...        break
...     for i in range(n):
...       yield i
...     n = yield
...
>>> t = toy(2)
>>> next(t)
0
>>> next(t)
1
>>> next(t)
>>> next(t)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "<stdin>", line 5, in toy
TypeError: 'NoneType' object cannot be interpreted as an integer
>>> next(t)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
StopIteration
>>> next(t)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
StopIteration
>>> t.send(3)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
StopIteration

wow that’s interesting – it is possible to call next() on it again and break the little toy. maybe there is some way to make a command/control convention so that you can inspect if it is asking for a next value or not.