example-code-2e/17-it-generator/aritprog_v3.py
2021-02-14 20:58:46 -03:00

12 lines
302 B
Python

# tag::ARITPROG_ITERTOOLS[]
import itertools
def aritprog_gen(begin, step, end=None):
first = type(begin + step)(begin)
ap_gen = itertools.count(first, step)
if end is not None:
ap_gen = itertools.takewhile(lambda n: n < end, ap_gen)
return ap_gen
# end::ARITPROG_ITERTOOLS[]