Yield only as many are required from a generator

Question

I wish to yield from a generator only as many items are required.

In the following code

a, b, c = itertools.count()

I receive this exception:

ValueError: too many values to unpack

I've seen several related questions, however I have zero interest in the remaining items from the generator, I only wish to receive as many as I ask for, without providing that quantity in advance.

It seems to me that Python determines the number of items you want, but then proceeds to try to read and store more than that number (generating ValueError).

How can I yield only as many items as I require, without passing in how many items I want?

Update0

To aid in understanding the approximate behaviour I believe should be possible, here's a code sample:

def unpack_seq(ctx, items, seq):
    for name in items:
        setattr(ctx, name, seq.next())

import itertools, sys
unpack_seq(sys.modules[__name__], ["a", "b", "c"], itertools.count())
print a, b, c

If you can improve this code please do.

Alex Martelli's answer suggests to me the byte op UNPACK_SEQUENCE is responsible for the limitation. I don't see why this operation should require that generated sequences must also exactly match in length.

Note that Python 3 has different unpack syntaxes which probably invalidate technical discussion in this question.

Answer 1

You need to make sure that the number of items on both sides matches. One way is to use islice from the itertools module:

from itertools import islice
a, b, c = islice(itertools.count(), 3)

Answer 2

Python does not work the way you desire. In any assignment, like

a, b, c = itertools.count()

the right-hand side is evaluated first, before the left-hand side. The right-hand side can not know how many items are on the left-hand side unless you tell it.

Answer 3

You need a deep bytecode hack - what you request cannot be done at Python source code level, but (if you're willing to commit to a specific version and release of Python) may be doable by post-processing the bytecode after Python has compiled it. Consider, e.g.:

>>> def f(someit):
...   a, b, c = someit()
... 
>>> dis.dis(f)
  2           0 LOAD_FAST                0 (someit)
              3 CALL_FUNCTION            0
              6 UNPACK_SEQUENCE          3
              9 STORE_FAST               1 (a)
             12 STORE_FAST               2 (b)
             15 STORE_FAST               3 (c)
             18 LOAD_CONST               0 (None)
             21 RETURN_VALUE        
>>> 

As you see, the UNPACK_SEQUENCE 3 bytecode occurs without the iterator someit being given the least indication of it (the iterator has already been called!) -- so you must prefix it in the bytecode with a "get exactly N bytes" operation, e.g.:

>>> def g(someit):
...   a, b, c = limit(someit(), 3)
... 
>>> dis.dis(g)
  2           0 LOAD_GLOBAL              0 (limit)
              3 LOAD_FAST                0 (someit)
              6 CALL_FUNCTION            0
              9 LOAD_CONST               1 (3)
             12 CALL_FUNCTION            2
             15 UNPACK_SEQUENCE          3
             18 STORE_FAST               1 (a)
             21 STORE_FAST               2 (b)
             24 STORE_FAST               3 (c)
             27 LOAD_CONST               0 (None)
             30 RETURN_VALUE        

where limit is your own function, easily implemented (e.g via itertools.slice). So, the original 2-bytecode "load fast; call function" sequence (just before the unpack-sequence bytecode) must become this kind of 5-bytecode sequence, with a load-global bytecode for limit before the original sequence, and the load-const; call function sequence after it.

You can of course implement this bytecode munging in a decorator.

Alternatively (for generality) you could work on altering the function's original source, e.g. via parsing and alteration of the AST, and recompiling to byte code (but that does require the source to be available at decoration time, of course).

Is this worth it for production use? Absolutely not, of course -- ridiculous amount of work for a tiny "syntax sugar" improvement. However, it can be an instructive project to pursue to gain mastery in bytecode hacking, ast hacking, and other black magic tricks that you will probably never need but are surely cool to know when you want to move beyond the language of mere language wizard to that of world-class guru -- indeed I suspect that those who are motivated to become gurus are typically people who can't fail to yield to the fascination of such "language internals' mechanics"... and those who actually make it to that exalted level are the subset wise enough to realize such endeavors are "just play", and pursue them as a spare-time activity (mental equivalent of weight lifting, much like, say, sudoku or crosswords;-) without letting them interfere with the important tasks (delivering value to users by deploying solid, clear, simple, well-performing, well-tested, well-documented code, most often without even the slightest hint of black magic to it;-).

Answer 4

Or use a list comprehension, since you know the desired number of items:

ic = itertools.count()
a,b,c = [ic.next() for i in range(3)]

or even simpler:

a,b,c = range(3)