example-code-2e/17-it-generator/aritprog_v3.py
2021-09-15 22:48:08 -03:00

12 lines
296 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 None:
return ap_gen
return itertools.takewhile(lambda n: n < end, ap_gen)
# end::ARITPROG_ITERTOOLS[]