Python zip magic for classes instead of tuples

Question

Python zip function is its own inverse (in a way), thus we can do this:

points = [(1,2), (3,4), (5,6), (7,8)]
xs, ys = zip(*points)

and now xs=[1,3,5,7] and ys=[2,4,6,8].

I wonder if something similar can be done with data class instances instead of tuples:

from dataclasses import dataclass

@dataclass
class XY:
    "2d point"
    x: float | int
    y: float | int

points = [XY(1,2), XY(3,4), XY(5,6), XY(7,8)]
xs, ys = zip(*[(p.x,p.y) for p in points])

but without an explicit list comprehension.

Of course, the result would not be a tuple (xs,ys) but a dict with keys x and y because, without an explicit list comprehension, we would be collecting all fields.

Answer 1

You can define custom __iter__ magic function in your dataclass:

from dataclasses import dataclass

@dataclass
class XY:
    "2d point"
    x: float | int
    y: float | int

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

points = [XY(1,2), XY(3,4), XY(5,6), XY(7,8)]

xs, ys = zip(*points)
print(xs)
print(ys)

Prints:

(1, 3, 5, 7)
(2, 4, 6, 8)

Answer 2

With astuple:

from dataclasses import dataclass, astuple

@dataclass
class XY:
    "2d point"
    x: float | int
    y: float | int
    def __iter__(self):
        return iter(astuple(self))

points = [XY(1,2), XY(3,4), XY(5,6), XY(7,8)]
xs, ys = zip(*points)

Or instead map it:

xs, ys = zip(*map(astuple, points))