Assigning a condition to a variable in Python

Question

Is there a way to assign a condition to a variable in Python?

Specifically, I am writing a generator that is infinite by default, limited by choice. This code works but the duplication is ugly:

def generate(start=0, stop=None):
    i = start

    if stop == None:
        while True:
            yield i
            i += 1
    else:
        while i <= stop:
            yield i
            i += 1

Specifically, I would like to express something like this:

def generate(start=0, stop=None):
    i = start

    if stop == None:
        condition = True
    else:
        condition = 'i <= stop' # Of course, this will throw an error

    while condition:
        yield i
        i += 1

What is the best way to accomplish this?

Thanks

Answer 1

It looks like you are reimplementing itertools.count and itertools takewhile.

Count forever forever:

itertools.count()

Count upto stop:

itertools.takewhile(lambda a: a <= stop, itertools.count())

But maybe this is a dumbed down example, and maybe making the condition a function is the answer you are looking for:

def generate(start=0, stop=None):
    i = start

    if stop == None:
        condition = lambda _: True
    else:
        condition = lambda x: x <= stop

    while condition(i):
        yield i
        i += 1

Answer 2

 def gener(x):
     for item in range(x):
             if item == 5:
                     raise(StopIteration)
             yield(item)

this is an example of a generator that wont go past 5.

>>> a = gener(7)
>>> [x for x in a]
>>>
[0, 1, 2, 3, 4]

This works, because StopIteration is a built in exception that is raised by the generator object when you reach the end of a generators output

>>> def gen(): #the most basic generator
...     for item in range(3): yield item
...
>>>
>>> a = gen()
>>> a.next()
0
>>> a.next()
1
>>> a.next()
2
>>> a.next()
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
StopIteration

when it reaches the end it naturally raises StopIteration as you can see... so raise it yourself, if you want to stop iterating.... I am unsure as to how pythonic this is.... to me it kinda feels alright though, lol