3
def get_data(batchsize=None):
    conn = psycopg2.connect('database parameter')
    cursor = conn.cursor()
    query = 'select * from table'
    cursor.execute(query)
    if not batchsize:
        result = cursor.fetchall()
        return result
    else:
        while True:
            result = cursor.fetchmany(batchsize)
            if not result:
                break
            yield result

if __name__ == '__main__':
    data = get_data()

In above function type of data should have been a list as argument batchsize=None. but this function is returning generator in both cases. If I comment else part of the function then it's returning a list.

Mohit Rajpoot
  • 101
  • 10
  • Alright, so what's the question here? And what are you trying to achieve? – Ivan Vinogradov May 25 '18 at 13:20
  • I want this function to return a list when batchsize=None and a generator when batchsize is not None. but right now I'm getting generator only for both cases. Am I missing something here? – Mohit Rajpoot May 25 '18 at 13:28

1 Answers1

3

Ah, I got what you're asking about. You're asking:

Why do I still receive a generator even if I don't hit yield keyword in get_data()?

The thing is that a function, which contains at least one yield statement is a generator. So your get_data() function is a generator. Because of that you always receive a generator object out of get_data().

The meaning of return in a generator is a bit different from ordinary function. In a generator any return x statement is equivalent to raise StopIteration(x).

And when you comment out part with else, you comment out yield statement, so get_data() becomes an ordinary function. Therefore it returns a list as you expect.

Related SO post: Return and yield in the same function

Ivan Vinogradov
  • 4,269
  • 6
  • 29
  • 39