itertools.repeat(object[, times]) How do I optionally set the times option at runtime?

Question

I have a loop:

for x in itertools.repeat(1)

Depending on the cmd line parameters I need this to be infinite or x number of times.

So this

for x in itertools.repeat(1)

or this

for x in itertools.repeat(1, x)

How can I do this?

Answer 1

itertools.repeat returns an iterable, so you can do this:

import sys
import itertools

if len(sys.argv) > 1:
    repeater = itertools.repeat(1, int(sys.argv[1]))
else:
    repeater = itertools.repeat(1)

for x in repeater:
    print x

sys.argv is the cli inputs; argv[0] will always be the script name, so I'm assuming argv[1] will be your input.

Answer 2

It's Python, so you can do this: if myCondition is the thing you want to test (eg a == 2 ), then:

myIterator = itertools.repeat(1) if myCondition else itertools.repeat(1, x)

for x in myIterator:
    do-something

Answer 3

You can use islice instead of the second argument to limit the number of items returned.

import argparse
import itertools

p = argparse.ArgumentParser()
p.add_argument("--count", type=int, default=None)
args = p.parse_args()

ones = itertools.repeat(1)
for x in itertools.islice(ones, args.count):
    ...

If the stop argument to islice is None , then there is no upper limit on the number of items in the slice.

---

I'm surprised that repeat itself does not accept None to explicitly trigger the default behavior. I submitted a bug report ; we'll see if the maintainers consider it a bug as well.

Answer 4

Here's a quick one-liner solution that's nice compact which I used, I wanted to repeat an infinite number of times or only once, you can use this:

itertools.repeat(my_object, *([] if my_condition else [1]))

The argument expansion won't happen for the [] case and thus you get the effect of None you would expect based on the documentation and (in my opinion) a reasonable reading of the functions purpose.

Otherwise [1] will expand to times=1 , you could replace 1 with whatever value you wanted.

times=None would have been cleaner, but c'est las vie, the maintainers didn't want it ( https://bugs.python.org/issue34169 ).