Декоратор @dataclass в Python: создание data-классов

В этом руководстве вы узнаете о декораторе @dataclass в Python и о том, как его эффективно использовать.

Содержание

Что такое dataclass в Python?

Python представил dataclass в версии 3.7(PEP 557). Он позволяет вам определять data-классы с меньшим количеством кода и большей функциональностью.

Ниже определяется обычный класс Person с двумя атрибутами экземпляра name и age:

class Person:
    def __init__(self, name, age):
        self.name = name
        self.age = age

Этот класс Person имеет метод __init__, который инициализирует атрибуты имени и возраста.

Если вы хотите иметь строковое представление объекта Person, вам необходимо реализовать метод __str__ или __repr__. Кроме того, если вы хотите сравнить два экземпляра класса Person по атрибуту, вам необходимо реализовать метод __eq__.

Однако если вы используете dataclass, вы получите все эти функции (и даже больше) без реализации этих методов.

Чтобы сделать класс Person data-классом, выполните следующие действия:

  • Сначала импортируйте декоратор dataclass из модуля dataclasses:
from dataclasses import dataclass
  • Во-вторых, примените к классу Person декоратор @dataclass и объявите атрибуты:
@dataclass
class Person:
    name: str
    age: int

В этом примере класс Person имеет два атрибута name типа str и age типа int. Сделав это, декоратор @dataclass неявно создает метод __init__ следующим образом:

def __init__(name: str, age: int)

Обратите внимание, что порядок атрибутов, объявленных в классе, будет определять порядок параметров в методе __init__.

И вы можете создать объект Person:

p1 = Person('John', 25)

Распечатав объект Person, вы получите читаемый формат:

print(p1)

Выход:

Person(name='John', age=25)

Кроме того, если вы сравните два объекта Person с одинаковым значением атрибута, он вернет True. Например:

p1 = Person('John', 25)
p2 = Person('John', 25)
print(p1 == p2)

Выход:

True

Ниже обсуждаются другие функции, предоставляемые dataclass.

Значения по умолчанию

При использовании обычного класса вы можете определить значения атрибутов по умолчанию. Например, следующий класс Person имеет параметр iq со значением по умолчанию 100.

class Person:
    def __init__(self, name, age, iq=100):
        self.name = name
        self.age = age
        self.iq = iq

Чтобы определить значение по умолчанию для атрибута в data-классе, вы присваиваете его атрибуту следующим образом:

from dataclasses import dataclass


@dataclass
class Person:
    name: str
    age: int
    iq: int = 100


print(Person('John Doe', 25))

Как и правила параметров, атрибуты со значениями по умолчанию должны появляться после атрибутов без значений по умолчанию. Поэтому следующий код не будет работать:

from dataclasses import dataclass


@dataclass
class Person:
    iq: int = 100
    name: str
    age: int

Преобразование в кортеж или словарь

Модуль dataclasses имеет функции astuple() и asdict(), которые преобразуют экземпляр data-класса в кортеж и словарь. Например:

from dataclasses import dataclass, astuple, asdict


@dataclass
class Person:
    name: str
    age: int
    iq: int = 100


p = Person('John Doe', 25)

print(astuple(p))
print(asdict(p))

Выход:

('John Doe', 25, 100)
{'name': 'John Doe', 'age': 25, 'iq': 100}

Создание неизменяемых объектов

Чтобы создать объекты только для чтения из dataclass, вы можете установить для аргумента frozen декоратора dataclass значение True. Например:

from dataclasses import dataclass, astuple, asdict


@dataclass(frozen=True)
class Person:
    name: str
    age: int
    iq: int = 100

Если вы попытаетесь изменить атрибуты объекта после его создания, вы получите сообщение об ошибке. Например:

p = Person('Jane Doe', 25)
p.iq = 120

Ошибка:

dataclasses.FrozenInstanceError: cannot assign to field 'iq'

Настройка поведения атрибутов

Если вы не хотите инициализировать атрибут в методе __init__, вы можете использовать функцию field() из модуля dataclasses.

В следующем примере определяется атрибут can_vote, который инициализируется с помощью метода __init__:

from dataclasses import dataclass, field


class Person:
    name: str
    age: int
    iq: int = 100
    can_vote: bool = field(init=False)

Функция field() имеет несколько интересных параметров, таких как воспроизведение, хэш, сравнение и метаданные. Если вы хотите инициализировать атрибут, который зависит от значения другого атрибута, вы можете использовать метод __post_init__. Как следует из названия, Python вызывает метод __post_init__ после метода __init__.

Следующие примеры используют метод __post_init__ для инициализации атрибута can_vote на основе атрибута age:

from dataclasses import dataclass, field


@dataclass
class Person:
    name: str
    age: int
    iq: int = 100
    can_vote: bool = field(init=False)

    def __post_init__(self):
        print('called __post_init__ method')
        self.can_vote = 18 <= self.age <= 70


p = Person('Jane Doe', 25)
print(p)

Выход:

called the __post_init__ method
Person(name='Jane Doe', age=25, iq=100, can_vote=True)

Сортировка объектов

По умолчанию dataclass реализует метод __eq__.

Чтобы разрешить различные типы сравнений, такие как __lt__, __lte__, __gt__, __gte__, вы можете установить для аргумента order декоратора @dataclass значение True:

@dataclass(order=True)

При этом dataclass будет сортировать объекты по каждому полю, пока не найдет неравное значение.

На практике часто требуется сравнивать объекты по определенному атрибуту, а не по всем атрибутам. Для этого вам необходимо определить поле с именем sort_index и установить его значение для атрибута, который вы хотите отсортировать.

Предположим, что у вас есть список объектов Person и вы хотите отсортировать их по возрасту:

members = [
    Person('John', 25),
    Person('Bob', 35),
    Person('Alice', 30)
]

Для этого вам необходимо:

  • Сначала передать параметр order=True декоратору @dataclass.
  • Во-вторых, определить атрибут sort_index и установить для его параметра init значение False.
  • В-третьих, установить для sort_index атрибут age в методе __post_init__, чтобы отсортировать объект Person по возрасту.

Ниже показан код сортировки объектов Person по возрасту:

from dataclasses import dataclass, field


@dataclass(order=True)
class Person:
    sort_index: int = field(init=False, repr=False)

    name: str
    age: int
    iq: int = 100
    can_vote: bool = field(init=False)

    def __post_init__(self):
        self.can_vote = 18 <= self.age <= 70
        # sort by age
        self.sort_index = self.age


members = [
    Person(name='John', age=25),
    Person(name='Bob', age=35),
    Person(name='Alice', age=30)
]

sorted_members = sorted(members)
for member in sorted_members:
    print(f'{member.name}(age={member.age})')

Выход:

John(age=25)
Alice(age=30)
Bob(age=35)
Похожие посты
Добавить комментарий

Ваш адрес email не будет опубликован. Обязательные поля помечены *