How to implement comprehension in a class

user_12040751 0
Viewed 42

As far as I know, comprehensions in Python only work with lists, dictionaries, tuples and sets. There is nothing on comprehension in the Python data model.

The syntax for tuples is quite interesting

>>> tuple(i for i in range(3))
(0, 1, 2)

Following the same pattern, I'd like to write a comprehension for my custom class.

>>> MySequence(i for i in range(3))
MySequence(0, 1, 2)
>>> MyMapping{str(i): i for i in range(3)}
MyMapping({'0': 0, '1': 1, '2': 2})

How can I achieve that?

dictionary-comprehension list-comprehension python
2 Answers

You are mixing up two things that are not related: class instantiation and (generator-)comprehensions.

tuple(i for i in range(3))

is equivalent to

tuple((i for i in range(3)))

which is equivalent to

generator = (i for i in range(3))
tuple(generator)

The generator-comprehension is evaluated before tuple.__init__ (or __new__) is called. In general, all arguments in Python are evaluated before being passed to a callable.

Any class can accept an iterable (such as generators) for instantiation if you code __init__ accordingly, e.g.

class AcceptIterable:
    def __init__(self, it):
        for x in it:
            print(x)

Demo:

>>> a = AcceptIterable(i for i in range(3))
0
1
2
>>> a = AcceptIterable([i for i in range(3)])
0
1
2
>>> a = AcceptIterable(range(3))
0
1
2
>>> a = AcceptIterable([0, 1, 2])
0
1
2
user_36200030

I managed to achieve this by subclassing and extending the __init__ method. Meet AbsList, which inherits from list and takes the absolute value of every number when initialising.

class AbsList(list):
    def __init__(self, li):
        super().__init__(abs(l) for l in li)
    def __repr__(self):
        return f"{self.__class__.__name__}({super().__repr__()})"

>>> AbsList(-i for i in range(-3))
AbsList([0, 1, 2])

AbsDict is quite similar.

class AbsDict(dict):
    def __init__(self, di):
        super().__init__({key: abs(value) for key, value in di.items()})
    def __repr__(self):
        return f"{self.__class__.__name__}({super().__repr__()})"

>>> AbsDict({str(i): -i for i in range(3)})
AbsDict({'0': 0, '1': 1, '2': 2})

I would have liked to remove the () brackets but I couldn't manage.

user_120407510
Related