Search code examples
pythoniteratorgeneratorpython-3.6iterable-unpacking

How to unpack an object as it was a tuple in a for loop?


I tried create the following code:

class Test(object):

    def __init__(self, arg):
        self.arg1 = arg + 1
        self.arg2 = arg + 2
        self.arg3 = arg + 3

    def __iter__(self):
        return self

    def __next__(self):
        yield self.arg1, self.arg2, self.arg3


test_list = [Test(0), Test(1), Test(2)]

for arg1, arg2, arg3 in test_list:
    print('arg1', arg1, 'arg2', arg2, 'arg3', arg3)

But when I try to run, python says:

Traceback (most recent call last):
  File "test.py", line 20, in <module>
    for arg1, arg2, arg3 in test_list:
ValueError: too many values to unpack (expected 3)

I can work around it by unpacking by hand:

class Test(object):

    def __init__(self, arg):
        self.arg1 = arg + 1
        self.arg2 = arg + 2
        self.arg3 = arg + 3

test_list = [Test(0), Test(1), Test(2)]

for test in test_list:
    arg1, arg2, arg3 = test.arg1, test.arg2, test.arg3
    print('arg1', arg1, 'arg2', arg2, 'arg3', arg3)

How can we unpack objects in a python list without doing it explicitly as the workaround demonstrated? For the last example, the result would be like:

arg1 1 arg2 2 arg3 3
arg1 2 arg2 3 arg3 4
arg1 3 arg2 4 arg3 5

Solution

  • You are close, however, you need to yield the values in the __iter__ method, not the __next__ method:

    class Test:
      def __init__(self, arg):
        self.arg1 = arg + 1
        self.arg2 = arg + 2
        self.arg3 = arg + 3
      def __iter__(self):
        yield from [self.arg1, self.arg2, self.arg3]
    
    for a, b, c in [Test(0), Test(1), Test(2)]:
      pass
    

    yield self.arg1, self.arg2, self.arg3 will give a tuple result (1, 2, 3) which, when iterating over the list, requires additional unpacking i.e:

    for [(a, b, c)] in [Test(0), Test(1), Test(2)]:
      pass
    

    Thus, in order to avoid the additional unpacking in the loop, you have to create a stream of generated values by looping over the attributes and yielding each, one at a time.