Python: Using a generator inside a while loop

Question

Here's my code so far:

def down_data(h):
    t = 0
    while h > 0:
        t += .0001
        h = -4.9*t**2 + 1
        v = g*t
        yield h, v
    return False

def up_data(v):
    t = 0
    h = 0
    vel = v
    while h < 1:
        t += .0001
        h = -4.9*t**2 - v*t
        vel = -9.8*t - v
        yield h, vel, t

What I'm trying to do is get my main() function to output the values of down_data(1) for as long as that function remains valid. Running this incomplete code just sets up an infinite loop and I'm not sure why.

I'm trying to combine these two functions to simulate a bouncing ball, but stitching them together has proven difficult because of the yield function. I've tried creating a main() and running while loop over one of the functions, but this creates an infinite loop. Here's the main() I've tried:

def main():
    data = [] 
    while down_data(1):
        data.append(down_data(1))
    print(data)

Is there maybe an easier way to do this? The main() function I'm after would be something like this:

Edit: I've also tried out this, but the same problem happens; the while loop doesn't close out to move to the next i:

def main():
    data = []
    for i in down_data(1):
        while i[0] > 0:
            data.append(i[0])
            continue
    print(data)

Answer 1

Generator functions in Python work a little different than the way you are trying to use them.

You should call down_data(1) only once.

Try this:

g = 9.8

def down_data(h):
    t = 0
    while h > 0:
        t += .0001
        h = -4.9*t**2 + 1
        v = g*t
        yield h, v

for h, v in down_data(1):
    print h, v

This will generate a number of output lines, but it will stop when h reaches 0.

Explanation: A Python function with a yield in it, when called, produces a generator. Each time you call next on the generator, the generator runs until it hits a yield and the yielded values are returned to the caller. If the generator hits the end of its body without yielding anything, a StopIteration exception is raised.

When you iterate over a generator in a for statement, the for-loop exits cleanly when the generator has no more data to yield.

Answer 2

Ray is right that generators can be awkward to use outside of for e in gen which hides the StopIteration exception

but no SyntaxError: 'return' with argument inside generator, at least in Python 3.5, Spyder IDE

g = 1

def down_data(h):
    t = 0
    while h > 0:
        t += .0001
        h = -4.9*t**2 + 1
        v = g*t
        yield h, v
    return 'fred' # this will be returned on StopIteration Exception

def mainly():
   data_gen =  down_data(1)
   data = [next(data_gen)]
   while data_gen:
       data.append(next(data_gen))
   #print(data)
   return len(data)        

mainly() fails because of the StopIteration Exception, but the down_data(h) return does get executed following the exception:

 mainly()
Traceback (most recent call last):

  File "<ipython-input-103-015c95441865>", line 1, in <module>
    mainly()

  File "C:/Users/john/mypy/mySE_answers/tagd.py", line 48, in mainly
    data.append(next(data_gen))

StopIteration: fred

explicit try/except is needed if you are going to the end of the generator with your own next() calls

def my_try():
   data_gen =  down_data(1)
   try:
       data = [next(data_gen)]
       while True:
           data.append(next(data_gen))
   except: StopIteration

   return len(data), data[-1]

my_try()
Out[104]: (4518, (-0.0002038759998519435, 0.45179999999996656))