Dataclassy w Pythonie

Jakiś czas temu, w ramach zadania rekrutacyjnego, miałem zaimplementować prosty system bankowy. Zadanie było podzielone na etapy i wykonywane na żywo, w ograniczonym czasie. Po poprawnym napisaniu obsługi dodawania kont, wpłacania i wypłacania z nich pieniędzy, otrzymałem kolejne polecenie.

Następną funkcjonalnością miały być przelewy zlecone. Miałem stworzyć kolejkę, do której użytkownik będzie mógł wrzucić zdefiniowany przelew - podając kwotę, konta wysyłającego i adresata, oraz timestamp, w jakim przelew ma być wykonany.

W prawym rogu ekranu widziałem stoper, który bezlitośnie odliczał minuty pozostałe do końca. Postanowiłem więc oszczędzić sobie czasu i, zamiast definiować jakąś strukturę danych, przechowywać dane tych operacji w zwykłej krotce (tuple). I to był poważny błąd.

Nie zliczę ile razy w ciągu kolejnego kwadransa zadawałem sobie pytania “czy timestamp jest na pierwszym, czy ostatnim miejscu?”, albo “najpierw jest kwota czy id odbiorcy?”. Finalnie, straciłem pewnie więcej czasu, tyle że podzielonego na mniejsze, za to wypełnione wkurwieniem, odcinki.

Co mogłem zrobić, żeby tego uniknąć? Kiedy pisałem o protokołach, wspominałem o obiektach służących głównie do przechowywania danych. I o jednym z nich - dataclass będzie traktował ten wpis.

Zacznijmy jednak od definicji problemu. Jedną jego stronę już znamy- użycie krotki jest rozwiązaniem szybkim, ale wysoce niesatysfakcjonującym. Można spróbować zamienić krotkę na słownik. Wtedy nie będę musiał pamiętać o kolejności, mogę sobie ponazywać pola jak chcę i się do nich odwoływać. Zatem zamiast tego:

sender = 121
receiver = 144
amount = 1000
timestamp = 11

delayed_transfer_tuple = (sender, receiver, amount, timestamp)

Miałbym to:

delayed_transfer_dict = {
    "sender": sender,
    "receiver": receiver,
    "amount": amount,
    "timestamp": timestamp
}

Czy jest to lepsze? Nadal muszę pamiętać, czy kwota to ammount czy sum, IDE nie podpowie mi jaka jest nazwa, a nazwy pól muszę podawać używając gołych stringów. Co przy pobieraniu może wywołać błędy, ale przy operacjach przypisania może pozostać niezauważone. Jak w poniższym przykładzie.

delayed_transfer_dict["amount"] = 11

Suma pozostanie niezmieniona, za to słownik wzbogaci się o kolejną parę klucz-wartość, a ja będę się głowił dlaczego przypisanie się nie powiodło¹.

Z drugiej strony, definiowanie całej klasy, wydaje się nieco nadmiarowe, choćby ze względu na ilość kodu, którą trzeba wygenerować. Zacznijmy od rzeczy najprostszych:

class DelayedTransfer:
    def __init__(self, sender, receiver, amount, timestamp):
        self.sender = sender
        self.receiver = receiver
        self.amount = amount
        self.timestamp = timestamp

transfer = DelayedTransfer(sender, receiver, ammount, timestamp)

Niby niewiele, ale każdą z wartości już trzeba było wymienić trzykrotnie. Możliwe jednak, że będziemy potrzebowali tego kodu więcej. Wyobraźcie sobie choćby sprawdzenie, czy dwa przelewy są identyczne. W krotce dostajemy to za darmo. W przypadku klasy, niestety już nie:

delayed_transfer_tuple = (sender, receiver, amount, timestamp)
delayed_transfer_tuple_2 = (sender, receiver, amount, timestamp)
print(delayed_transfer_tuple == delayed_transfer_tuple_2) # True

transfer = DelayedTransfer(sender, receiver, amount, timestamp)
transfer2 = DelayedTransfer(sender, receiver, amount, timestamp)
print(transfer == transfer2)  # False

Podobnie z porównaniami, potrzebnymi do sortowania. Metody __lt__, __gt__ i tak dalej będziemy musieli zaimplementować sami.

Dlatego też powstał dokument PEP-557, i w wersji 3.7 Pythona otrzymaliśmy dataclassy- struktury przeznaczone do przechowywania danych, pozwalające na łatwe stworzenie klasy dającej nam powyższe funkcjonalności - właśnie Dataclass.

Dzięki nim, wystarczy poniższa deklaracja:

@dataclass
class DelayedTransfer:
    sender: int
    receiver: int
    amount: int
    timestamp: int

Abyśmy otrzymali klasę “uzbrojoną” w konstruktor oraz metody __repr__, __eq__, i wykorzystywaną przy pattern matchingu krotkę __match_args__ ² Ale to nie wszystko. O tym, jakie dokładnie metody zostaną wygenerowane, możemy decydować, podając następujące argumenty:

Oczywiście, metody te są bardzo generyczne. Konstruktor jedynie przypisuje wartości pól. Wygenerowane operatory porównania używają krotki z wartościami pól. Jednak w znacznej większości przypadków, jest to logika całkowicie wystarczająca.

Większość tych parametrów jest dość oczywista. Warto jednak wspomnieć więcej o kilku. I tak, frozen sprawia, że nasza klasa staje się niezmienna (immutable), więc zachowaniem przypomina krotkę na sterydach a jakakolwiek próba przypisania wartości zakończy się rzuceniem FrozenInstanceError.

Unsafe hash z kolei wymusza stworzenie funkcji __hash__, nawet jeśli wartości pól naszej klasy mogą się zmieniać. To dośc wyspecjalizowany przypadek, stosowany głównie wtedy, gdy definicja klasy zawiera mutowalne wartości, ale programista wie, że logika programu nie pozwoli na ich zmianę. Natomiast w przypadku, gdy parametry eq i frozen ustawione są na True, metoda __hash__ wygenerowana zostanie automatycznie.

Stworzona w ten sposób klasa w żaden sposób nie różni się od innych klas. Tak samo, możemy zdefiniować w niej swoje własne metody. Czasem jest to nawet nieodzowne.

Sięgnijcie bowiem pamięcią do wszystkich metod __init__ jakie w życiu napisaliście. Zapewne większość z nich zawiera jedynie proste przypisania, w stylu self.foo = foo.

Zdarza się jednak, że oczekujemy od konstruktora, wykonania bardziej zawiłej logiki. Wyobraźcie sobie, że modelujemy samochód. Dla uproszczenia, nasze auto będzie określane przez kolor i model silnika. Chcielibyśmy jednak określić, jaki dźwięk wydaje nasz samochód- jeśli jest to elektryk, będzie wydawał z siebie bzyczenie, jeśli zaś diesel albo benzyniak- będzie warczał.

Możemy oczywiście ten dźwięk podać jako kolejny argument dla konstruktora, ale to złe rozwiązanie. Po pierwsze, dopuszczamy pomieszanie dźwięków, a nie chcemy żeby nasze TDI wydawało z siebie brzęczenie komara, zamiast prawdziwie męskich pomruków.

@dataclass
class Car:
    engine: str
    color: str
    sound: str = None

Dodatkowo, przy każdej instancjalizacji klasy, programista musi zadbać o podanie dodatkowego argumentu, w dodatku zgodnego z dokumentacją.

Rozwiązania są dwa- albo zaszyjemy tę logikę w metodzie z dekoratorem property, i za każdym razem będziemy wyznaczać ten dźwięk na nowo, albo użyjemy __post_init__.

Jest to specjalna metoda, która zostaje wywołana na końcu automatycznie wygenerowanego konstruktora, i pozwala na dopisanie właśnie takiej logiki. Zatem nasza klasa może wyglądać tak:

@dataclass
class Car:
    engine: str
    color: str
    
    def __post_init__(self):
        if self.engine.startswith('EL'):
            self.sound = 'Wzium'
        else:
            self.sound = 'Brum'

Dzięki czemu, nie musimy się już więcej kłopotać dźwiękiem (przynajmniej dopóki coś nie zacznie stukać w nadkolu):

In [22]: a= Car('Red', 'EL-123')

In [23]: a.sound
Out[23]: 'Wzium'

In [24]: a= Car('Red', 'TDI-123')

In [25]: a.sound
Out[25]: 'Brum'

Ale czy korzyść z dataclass ogranicza się do automatycznego generowania metod? No nie. Dzięki klasie Field zyskujemy kontrolę nad tym, jakie wartości atrybutów są dopuszczalne itp. Ale o tym będzie w kolejnym wpisie.

Na koniec jednak muszę dodać jedną ciekawostkę. PEP-484, a więc dokument wprowadzający adnotacje typów, zakładał, że zawsze będą one opcjonalne. Co za tym idzie, ich obecność, albo nie, nie powinna w żadnym wypadku zmieniać zachowania interpretera w runtime:

It should also be emphasized that Python will remain a dynamically typed language, and the authors have no desire to ever make type hints mandatory, even by convention.

Dataclassy jednak są przypadkiem, kiedy adnotacje zmieniają zachowanie. Weźmy taki przyład, zaczerpnięty ze znakomitej książki Luciano Ramalho “Fluent Python”.

from dataclasses import dataclass
@dataclass
class DemoDataClass:
    a: int
    b: float = 1.1
    c = 'spam'

Jest to prosta klasa z trzema argumentami. Dwa z nich mają adnotacje, a jeden nie. Dwa z nich mają wartości domyślne, a jeden nie. Co zatem dzieje się w runtime?

>>> from demo_dc import DemoDataClass
>>> DemoDataClass.__annotations__
{'a': <class 'int'>, 'b': <class 'float'>}
>>> DemoDataClass.__doc__
'DemoDataClass(a: int, b: float = 1.1)'
>>> DemoDataClass.a
Traceback (most recent call last):
File "<stdin>", line 1, in <module>
AttributeError: type object 'DemoDataClass' has no attribute 'a'
>>> DemoDataClass.b
1.1
>>> DemoDataClass.c
'spam'

Argumenty b i c widoczne są z poziomu klasy. Natomiast pytanie, co zawiera, generowana automatycznie, metoda __init__?

In [16]: inspect.signature(DemoDataClass.__init__)
Out[16]: <Signature (self, a: int, b: float = 1.1) -> None>

Jak widać, tylko te atrybuty, które opatrzono adnotacją typu, stają się atrybutami instancji. A zatem, wbrew intencjom autorów, adnotacja typu zmienia zachowanie kodu!