Home > Mobile >  Unpack two variables at a time in a list comprehension?
Unpack two variables at a time in a list comprehension?

Time:09-03

What is the simplest way to do this?

class a:
  def __init__(self):
    self.x = 1
    self.y = 2

a_list = [a(),a(),a()]

l = [a.x, a.y for a in a_list]

As a result, I want to get [1,2,1,2,1,2].

CodePudding user response:

Do this:

>>> [i for a in a_list for i in (a.x, a.y)]
[1, 2, 1, 2, 1, 2]

CodePudding user response:

this is not list comprehension, however, is simple.

import itertools
l = list(itertools.chain.from_iterable((a.x, a.y) for a in a_list))

CodePudding user response:

If you don't actually need the comprehension, don't care about defining __iter__ and only ever going to have 2 attributes:

class a:
    def __init__(self):
        self.x = 1
        self.y = 2

    def __iter__(self):
        yield self.x
        yield self.y


a_list = [*a(), *a(), *a()]

print(a_list)

outputs

[1, 2, 1, 2, 1, 2]

This removes the need for a nested loop or itertools usage and is arguably more readable.

Also, the only thing that needs to change if more attributes are added to a is the __iter__ implementation, and even that can be generalized:

class a:
    def __init__(self):
        self.x = 1
        self.y = 2
        self.z = 3

    def __iter__(self):
        yield from self.__dict__.values()
        # or return iter(self.__dict__.values()) which is a bit faster


a_list = [*a(), *a(), *a()]

print(a_list)

will output

[1, 2, 3, 1, 2, 3, 1, 2, 3]
  • Related