Перегрузка операторов и функций в классах Python

Оглавление

Если вы использовали оператор + или * на объекте str в Python, то наверняка заметили, что он ведет себя иначе, чем объекты int или float:

>>> # Adds the two numbers
>>> 1 + 2
3

>>> # Concatenates the two strings
>>> 'Real' + 'Python'
'RealPython'


>>> # Gives the product
>>> 3 * 2
6

>>> # Repeats the string
>>> 'Python' * 3
'PythonPythonPython'

Вы, возможно, задавались вопросом, как один и тот же встроенный оператор или функция показывает разное поведение для объектов разных классов. Это называется перегрузкой оператора или перегрузкой функции соответственно. Эта статья поможет вам понять этот механизм, чтобы вы могли сделать то же самое в своих Python-классах и сделать свои объекты более питонистыми.

Вы узнаете следующее:

  • API для работы с операторами и встроенными модулями в Python
  • "Секрет", скрывающийся за len() и другими встроенными операторами
  • Как сделать классы способными использовать операторы
  • Как сделать классы совместимыми со встроенными функциями Python

Модель данных Python

Допустим, у вас есть класс, представляющий онлайн-заказ, в котором есть корзина (list) и клиент (str или экземпляр другого класса, который представляет клиента).

Примечание: Если вам нужно освежить в памяти ООП в Python, посмотрите этот учебник на Real Python: Объектно-ориентированное программирование (ООП) в Python 3

В таком случае вполне естественным будет желание получить длину списка корзин. Новичок в Python может решить реализовать в своем классе метод get_cart_len() для этого. Но вы можете настроить встроенный len() таким образом, чтобы он возвращал длину списка тележек, когда нам дан объект.

В другом случае мы можем захотеть добавить что-то в корзину. Опять же, новичок в Python подумает о реализации метода append_to_cart(), который берет элемент и добавляет его в список корзины. Но вы можете настроить оператор + таким образом, чтобы он добавлял новый элемент в корзину.

Python делает все это с помощью специальных методов. Эти специальные методы имеют соглашение об именовании, в котором имя начинается с двух подчеркиваний, затем следует идентификатор и заканчивается еще одной парой подчёркиваний.

По сути, каждая встроенная функция или оператор имеет соответствующий ей специальный метод. Например, __len__(), соответствует len(), а __add__() - оператору +.

По умолчанию большинство встроенных модулей и операторов не будут работать с объектами ваших классов. Вы должны добавить соответствующие специальные методы в определение вашего класса, чтобы сделать ваш объект совместимым со встроенными модулями и операторами.

Когда вы это делаете, поведение связанной с ним функции или оператора меняется в соответствии с тем, что определено в методе.

Именно в этом вам поможет Модель данных (раздел 3 документации Python). В ней перечислены все доступные специальные методы и предоставлены средства перегрузки встроенных функций и операторов, чтобы вы могли использовать их на своих собственных объектах.

Давайте посмотрим, что это значит.

Забавный факт: Из-за соглашения о названиях, используемого для этих методов, их также называют dunder methods, что является сокращением для double underscore methods. Иногда их также называют специальными методами или магическими методами. Однако мы предпочитаем подземные методы!

Внутреннее устройство таких операций, как len() и []

Каждый класс в Python определяет свое поведение для встроенных функций и методов. Когда вы передаете экземпляр какого-либо класса во встроенную функцию или используете оператор над экземпляром, это фактически эквивалентно вызову специального метода с соответствующими аргументами.

Если существует встроенная функция и соответствующий func() ей специальный метод __func__(), Python интерпретирует вызов этой функции как obj.__func__(), где obj - объект. В случае с операторами, если у вас есть оператор opr и соответствующий специальный метод для него - __opr__(), Python интерпретирует что-то вроде obj1 <opr> obj2 как obj1.__opr__(obj2).

Таким образом, когда вы вызываете len() на объекте, Python обрабатывает вызов как obj.__len__(). Когда вы используете оператор [] на итерируемом объекте для получения значения по индексу, Python обрабатывает его как itr.__getitem__(index), где itr - объект итерируемого объекта, а index - индекс, который вы хотите получить.

Поэтому, когда вы определяете эти специальные методы в своем классе, вы переопределяете поведение связанной с ними функции или оператора, потому что за кулисами Python вызывает ваш метод. Давайте разберемся в этом подробнее:

>>> a = 'Real Python'
>>> b = ['Real', 'Python']
>>> len(a)
11
>>> a.__len__()
11
>>> b[0]
'Real'
>>> b.__getitem__(0)
'Real'

Как видите, при использовании функции или соответствующего ей специального метода вы получаете один и тот же результат. На самом деле, когда вы получаете список атрибутов и методов объекта str с помощью dir(), вы увидите эти специальные методы в списке в дополнение к обычным методам, доступным для объектов str:

>>> dir(a)
['__add__',
 '__class__',
 '__contains__',
 '__delattr__',
 '__dir__',
 ...,
 '__iter__',
 '__le__',
 '__len__',
 '__lt__',
 ...,
 'swapcase',
 'title',
 'translate',
 'upper',
 'zfill']

Если поведение встроенной функции или оператора не определено в классе специальным методом, то вы получите TypeError.

Итак, как вы можете использовать специальные методы в ваших классах?

Перегрузка встроенных функций

Многие из специальных методов, определенных в Модели данных, можно использовать для изменения поведения таких функций, как len, abs, hash, divmod и так далее. Для этого достаточно определить соответствующий специальный метод в своем классе. Давайте рассмотрим несколько примеров:

Придание длины объектам с помощью len()

Чтобы изменить поведение len(), вам нужно определить специальный метод __len__() в вашем классе. Всякий раз, когда вы передаете объект вашего класса в len(), для получения результата будет использоваться ваше пользовательское определение __len__(). Давайте реализуем len() для класса Order, о котором мы говорили в начале:

>>> class Order:
...     def __init__(self, cart, customer):
...         self.cart = list(cart)
...         self.customer = customer
...
...     def __len__(self):
...         return len(self.cart)
...
>>> order = Order(['banana', 'apple', 'mango'], 'Real Python')
>>> len(order)
3

Как видите, теперь вы можете использовать len(), чтобы напрямую получить длину тележки. Более того, интуитивно понятнее сказать "длина заказа", а не вызывать что-то вроде order.get_cart_len(). Ваш вызов и питоничен, и более интуитивен. Когда у вас не определен метод __len__(), но вы все равно вызываете len() на своем объекте, вы получаете TypeError:

>>> class Order:
...     def __init__(self, cart, customer):
...         self.cart = list(cart)
...         self.customer = customer
...
>>> order = Order(['banana', 'apple', 'mango'], 'Real Python')
>>> len(order)  # Calling len when no __len__
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
TypeError: object of type 'Order' has no len()

Но, перегружая len(), следует помнить, что Python требует, чтобы функция возвращала целое число. Если ваш метод будет возвращать не целое число, то вы получите TypeError. Скорее всего, это сделано для того, чтобы сохранить соответствие с тем фактом, что len() обычно используется для получения длины последовательности, которая может быть только целым числом:

>>> class Order:
...     def __init__(self, cart, customer):
...         self.cart = list(cart)
...         self.customer = customer
...
...     def __len__(self):
...         return float(len(self.cart))  # Return type changed to float
...
>>> order = Order(['banana', 'apple', 'mango'], 'Real Python')
>>> len(order)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
TypeError: 'float' object cannot be interpreted as an integer

Заставляя ваши объекты работать с abs()

Вы можете диктовать поведение abs() встроенного для экземпляров вашего класса, определив __abs__() специальный метод в классе. Нет никаких ограничений на возвращаемое значение abs(), и вы получаете TypeError, если специальный метод отсутствует в определении вашего класса.

В классе, представляющем вектор в двумерном пространстве, abs() можно использовать для получения длины вектора. Давайте посмотрим на это в действии:

>>> class Vector:
...     def __init__(self, x_comp, y_comp):
...         self.x_comp = x_comp
...         self.y_comp = y_comp
...
...     def __abs__(self):
...         return (self.x_comp ** 2 + self.y_comp ** 2) ** 0.5
...
>>> vector = Vector(3, 4)
>>> abs(vector)
5.0

Интуитивно понятнее сказать "абсолютное значение вектора", а не называть что-то вроде vector.get_mag().

Красивая печать объектов с помощью str()

Встроенная функция str() используется для приведения экземпляра класса к объекту str или, что более уместно, для получения удобного строкового представления объекта, которое может быть прочитано обычным пользователем, а не программистом. Вы можете определить формат строки, в котором должен отображаться ваш объект при передаче в str(), определив метод __str__() в вашем классе. Более того, __str__() - это метод, который используется Python, когда вы вызываете print() на вашем объекте.

Давайте реализуем это в классе Vector для форматирования объектов Vector в виде xi+yj. Отрицательная y-компонента будет обрабатываться с помощью мини-языка формата :

>>> class Vector:
...     def __init__(self, x_comp, y_comp):
...         self.x_comp = x_comp
...         self.y_comp = y_comp
...
...     def __str__(self):
...         # By default, sign of +ve number is not displayed
...         # Using `+`, sign is always displayed
...         return f'{self.x_comp}i{self.y_comp:+}j'
...
>>> vector = Vector(3, 4)
>>> str(vector)
'3i+4j'
>>> print(vector)
3i+4j

Необходимо, чтобы __str__() возвращал объект str, и мы получим TypeError, если возвращаемый тип не является строкой.

Представление объектов с помощью repr()

Встроенная функция repr() используется для получения разборчивого строкового представления объекта. Если объект поддается разбору, это означает, что Python должен быть в состоянии воссоздать объект из этого представления, когда repr используется в сочетании с такими функциями, как eval(). Чтобы определить поведение repr(), можно использовать специальный метод __repr__().

Это также метод, который Python использует для отображения объекта в сессии REPL. Если метод __repr__() не определен, то при попытке посмотреть на объект в сеансе REPL вы получите что-то вроде <__main__.Vector object at 0x...>. Давайте посмотрим, как это работает в классе Vector:

>>> class Vector:
...     def __init__(self, x_comp, y_comp):
...         self.x_comp = x_comp
...         self.y_comp = y_comp
...
...     def __repr__(self):
...         return f'Vector({self.x_comp}, {self.y_comp})'
...

>>> vector = Vector(3, 4)
>>> repr(vector)
'Vector(3, 4)'

>>> b = eval(repr(vector))
>>> type(b), b.x_comp, b.y_comp
(__main__.Vector, 3, 4)

>>> vector  # Looking at object; __repr__ used
'Vector(3, 4)'

Примечание: В случаях, когда метод __str__() не определен, Python использует метод __repr__() для печати объекта, а также для представления объекта при вызове str() на нем. Если оба метода отсутствуют, то по умолчанию используется метод <__main__.Vector ...>. Но __repr__() - это единственный метод, который используется для отображения объекта в интерактивной сессии. Отсутствие его в классе приводит к <__main__.Vector ...>.

Кроме того, хотя различие между __str__() и __repr__() является рекомендуемым, многие популярные библиотеки игнорируют это различие и используют эти два метода как взаимозаменяемые.

Вот рекомендуемая статья о __repr__() и __str__() нашего собственного Дэна Бейдера: Python String Conversion 101: Why Every Class Needs a "repr".

Делаем предметы правдивыми или ложными с помощью bool()

Для получения истинностного значения объекта можно использовать встроенную функцию bool(). Чтобы определить его поведение, вы можете использовать специальный метод __bool__() (__nonzero__() в Python 2.x).

Определенное здесь поведение будет определять истинностное значение экземпляра во всех контекстах, требующих получения истинностного значения, например, в утверждениях if.

Например, для класса Order, который был определен выше, экземпляр может считаться истинным, если длина списка корзин ненулевая. Это можно использовать для проверки того, следует ли обрабатывать заказ или нет:

>>> class Order:
...     def __init__(self, cart, customer):
...         self.cart = list(cart)
...         self.customer = customer
...
...     def __bool__(self):
...         return len(self.cart) > 0
...
>>> order1 = Order(['banana', 'apple', 'mango'], 'Real Python')
>>> order2 = Order([], 'Python')

>>> bool(order1)
True
>>> bool(order2)
False

>>> for order in [order1, order2]:
...     if order:
...         print(f"{order.customer}'s order is processing...")
...     else:
...         print(f"Empty order for customer {order.customer}")
Real Python's order is processing...
Empty order for customer Python

Примечание: Когда специальный метод __bool__() не реализован в классе, в качестве истинностного значения используется значение, возвращаемое __len__(), где ненулевое значение указывает на True, а нулевое - на False. В случае если оба метода не реализованы, все экземпляры класса считаются True.

Существует множество других специальных методов, которые перегружают встроенные функции. Вы можете найти их в документации. Рассмотрев некоторые из них, перейдем к операторам.

Перегрузка встроенных операторов

Изменить поведение операторов так же просто, как и поведение функций. Вы определяете соответствующие специальные методы в своем классе, и операторы работают в соответствии с поведением, определенным в этих методах.

Они отличаются от описанных выше специальных методов тем, что должны принимать в определении другой аргумент, отличный от self, обычно обозначаемый именем other. Давайте рассмотрим несколько примеров.

Создание объектов, которые можно добавлять с помощью +

Специальным методом, соответствующим оператору +, является метод __add__(). Добавление пользовательского определения __add__() изменяет поведение оператора. Рекомендуется, чтобы __add__() возвращал новый экземпляр класса, а не изменял сам вызывающий экземпляр. Такое поведение довольно часто встречается в Python:

>>> a = 'Real'
>>> a + 'Python'  # Gives new str instance
'RealPython'
>>> a  # Values unchanged
'Real'
>>> a = a + 'Python'  # Creates new instance and assigns a to it
>>> a
'RealPython'

Вы можете видеть выше, что использование оператора + на объекте str фактически возвращает новый экземпляр str, сохраняя значение вызывающего экземпляра (a) неизменным. Чтобы изменить его, нам нужно явно присвоить новый экземпляр объекту a.

Давайте реализуем возможность добавления новых товаров в корзину в классе Order с помощью оператора. Мы будем следовать рекомендуемой практике и сделаем так, чтобы оператор возвращал новый экземпляр Order, содержащий требуемые изменения, вместо того чтобы вносить изменения непосредственно в наш экземпляр:

>>> class Order:
...     def __init__(self, cart, customer):
...         self.cart = list(cart)
...         self.customer = customer
...
...     def __add__(self, other):
...         new_cart = self.cart.copy()
...         new_cart.append(other)
...         return Order(new_cart, self.customer)
...
>>> order = Order(['banana', 'apple'], 'Real Python')

>>> (order + 'orange').cart  # New Order instance
['banana', 'apple', 'orange']
>>> order.cart  # Original instance unchanged
['banana', 'apple']

>>> order = order + 'mango'  # Changing the original instance
>>> order.cart
['banana', 'apple', 'mango']

Аналогично, у вас есть __sub__(), __mul__() и другие специальные методы, которые определяют поведение -, * и так далее. Эти методы также должны возвращать новый экземпляр класса.

Короткие пути: оператор +=

Оператор += служит сокращением для выражения obj1 = obj1 + obj2. Соответствующий ему специальный метод - __iadd__(). Метод __iadd__() должен вносить изменения непосредственно в аргумент self и возвращать результат, который может быть, а может и не быть self. Такое поведение существенно отличается от __add__(), поскольку последний создает новый объект и возвращает его, как вы видели выше.

Грубо говоря, любое использование += на двух объектах эквивалентно следующему:

>>> result = obj1 + obj2
>>> obj1 = result

Здесь result - это значение, возвращаемое __iadd__(). Второе присваивание выполняется Python автоматически, то есть вам не нужно явно присваивать obj1 результату, как в случае с obj1 = obj1 + obj2.

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

>>> class Order:
...     def __init__(self, cart, customer):
...         self.cart = list(cart)
...         self.customer = customer
...
...     def __iadd__(self, other):
...         self.cart.append(other)
...         return self
...
>>> order = Order(['banana', 'apple'], 'Real Python')
>>> order += 'mango'
>>> order.cart
['banana', 'apple', 'mango']

Как видно, любое изменение вносится непосредственно в self, а затем возвращается обратно. Что произойдет, если вы вернете случайное значение, например строку или целое число?

>>> class Order:
...     def __init__(self, cart, customer):
...         self.cart = list(cart)
...         self.customer = customer
...
...     def __iadd__(self, other):
...         self.cart.append(other)
...         return 'Hey, I am string!'
...
>>> order = Order(['banana', 'apple'], 'Real Python')
>>> order += 'mango'
>>> order
'Hey, I am string!'

Несмотря на то, что соответствующий товар был добавлен в корзину, значение order изменилось на то, которое вернул __iadd__(). Python неявно выполнил присваивание за вас. Это может привести к неожиданному поведению, если вы забудете вернуть что-то в своей реализации:

>>> class Order:
...     def __init__(self, cart, customer):
...         self.cart = list(cart)
...         self.customer = customer
...
...     def __iadd__(self, other):
...         self.cart.append(other)
...
>>> order = Order(['banana', 'apple'], 'Real Python')
>>> order += 'mango'
>>> order  # No output
>>> type(order)
NoneType

<<<Поскольку все функции (или методы) Python возвращают None неявно, order переназначается на None, и сессия REPL не показывает никакого вывода при проверке order. Посмотрев на тип order, вы увидите, что теперь это NoneType. Поэтому всегда убеждайтесь, что в вашей реализации __iadd__() вы что-то возвращаете и что это результат операции, а не что-то другое.

Подобно __iadd__(), у вас есть __isub__(), __imul__(), __idiv__() и другие специальные методы, которые определяют поведение -=, *=, /= и других подобных методов.

Примечание: Когда __iadd__() или его друзья отсутствуют в определении вашего класса, но вы все равно используете их операторы на своих объектах, Python использует __add__() и его друзей для получения результата операции и присваивает его вызывающему экземпляру. Вообще говоря, можно не реализовывать __iadd__() и его друзей в ваших классах до тех пор, пока __add__() и его друзья работают правильно (возвращают что-то, что является результатом операции).

В документации Python documentation есть хорошее объяснение этих методов. Также посмотрите на пример this, который показывает, какие предостережения связаны с += и другими методами при работе с immutable типами.

Индексация и нарезка объектов с помощью []

Оператор [] называется оператором индексации и используется в различных контекстах в Python, например, для получения значения по индексу в последовательностях, для получения значения, связанного с ключом в словарях, или для получения части последовательности с помощью нарезки. Вы можете изменить его поведение с помощью специального метода __getitem__().

Давайте настроим наш класс Order так, чтобы мы могли напрямую использовать объект и получать товар из корзины:

>>> class Order:
...     def __init__(self, cart, customer):
...         self.cart = list(cart)
...         self.customer = customer
...
...     def __getitem__(self, key):
...         return self.cart[key]
...
>>> order = Order(['banana', 'apple'], 'Real Python')
>>> order[0]
'banana'
>>> order[-1]
'apple'

Обратите внимание, что выше имя аргумента __getitem__() не index, а key. Это связано с тем, что аргумент может иметь в основном три формы: целочисленное значение, в этом случае это либо индекс, либо ключ словаря, строковое значение, в этом случае это ключ словаря, и объект slice, в этом случае он будет нарезать последовательность, используемую классом. Хотя существуют и другие возможности, эти встречаются чаще всего.

Поскольку наша внутренняя структура данных представляет собой список, мы можем использовать оператор [] для нарезки списка, так как в этом случае аргумент key будет представлять собой объект slice. Это одно из самых больших преимуществ наличия определения __getitem__() в вашем классе. Если вы используете структуры данных, которые поддерживают нарезку (списки, кортежи, строки и так далее), вы можете настроить свои объекты на непосредственную нарезку структуры:

>>> order[1:]
['apple']
>>> order[::-1]
['apple', 'banana']

Примечание: Существует аналогичный специальный метод __setitem__(), который используется для определения поведения obj[x] = y. Этот метод принимает два аргумента в дополнение к self, обычно называемых key и value, и может быть использован для изменения значения в key на value.

Обратные операторы: Создание математически корректных классов

Хотя определение специальных методов __add__(), __sub__(), __mul__() и подобных позволяет использовать операторы, когда экземпляр класса является операндом левой стороны, оператор не будет работать, если экземпляр класса является операндом правой стороны:

>>> class Mock:
...     def __init__(self, num):
...         self.num = num
...     def __add__(self, other):
...         return Mock(self.num + other)
...
>>> mock = Mock(5)
>>> mock = mock + 6
>>> mock.num
11

>>> mock = 6 + Mock(5)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
TypeError: unsupported operand type(s) for +: 'int' and 'Mock'

Если ваш класс представляет математическую сущность, такую как вектор, координата или комплексное число, применение операторов должно работать в обоих случаях, поскольку это допустимая математическая операция.

Более того, если операторы работают только тогда, когда экземпляр является левым операндом, то во многих случаях мы нарушаем фундаментальный принцип коммутативности. Поэтому, чтобы помочь вам сделать ваши классы математически корректными, Python предоставляет вам обратные специальные методы, такие как __radd__(), __rsub__(), __rmul__() и так далее.

Они обрабатывают такие вызовы, как x + obj, x - obj и x * obj, где x не является экземпляром соответствующего класса. Как и __add__() и другие, эти обратные специальные методы должны возвращать новый экземпляр класса с изменениями операции, а не модифицировать сам вызывающий экземпляр.

Давайте настроим __radd__() в классе Order таким образом, чтобы он добавлял что-то в переднюю часть корзины. Это можно использовать в случаях, когда корзина организована с точки зрения приоритета заказов:

>>> class Order:
...     def __init__(self, cart, customer):
...         self.cart = list(cart)
...         self.customer = customer
...
...     def __add__(self, other):
...         new_cart = self.cart.copy()
...         new_cart.append(other)
...         return Order(new_cart, self.customer)
...
...     def __radd__(self, other):
...         new_cart = self.cart.copy()
...         new_cart.insert(0, other)
...         return Order(new_cart, self.customer)
...
>>> order = Order(['banana', 'apple'], 'Real Python')

>>> order = order + 'orange'
>>> order.cart
['banana', 'apple', 'orange']

>>> order = 'mango' + order
>>> order.cart
['mango', 'banana', 'apple', 'orange']

Полный пример

Чтобы донести все эти моменты до читателя, лучше рассмотреть пример класса, который реализует эти операторы вместе.

Давайте изобретем колесо и реализуем собственный класс для представления комплексных чисел, CustomComplex. Объекты нашего класса будут поддерживать множество встроенных функций и операторов, что сделает их поведение очень похожим на поведение встроенного класса комплексных чисел:

from math import hypot, atan, sin, cos

class CustomComplex:
    def __init__(self, real, imag):
        self.real = real
        self.imag = imag

Конструктор обрабатывает только один вид вызова, CustomComplex(a, b). Он принимает позиционные аргументы, представляющие действительную и мнимую части комплексного числа.

Определим два метода внутри класса, conjugate() и argz(), которые будут давать нам комплексный конъюгат и аргумент комплексного числа соответственно:

def conjugate(self):
    return self.__class__(self.real, -self.imag)

def argz(self):
    return atan(self.imag / self.real)

Примечание: __class__ - это не специальный метод, а атрибут класса, который присутствует по умолчанию. Он содержит ссылку на класс. Используя его здесь, мы получаем его, а затем вызываем конструктор обычным способом. Другими словами, это эквивалентно CustomComplex(real, imag). Это сделано для того, чтобы избежать рефакторинга кода, если имя класса однажды изменится.

Далее мы настраиваем abs() на возврат модуля комплексного числа:

def __abs__(self):
    return hypot(self.real, self.imag)

Мы будем следовать рекомендованному различию между __repr__() и __str__() и использовать первое для разборчивого представления строки, а второе - для "красивого" представления.

Метод __repr__() просто вернет CustomComplex(a, b) в строке, чтобы мы могли вызвать eval() для воссоздания объекта, а метод __str__() вернет комплексное число в скобках, как (a+bj):

def __repr__(self):
    return f"{self.__class__.__name__}({self.real}, {self.imag})"

def __str__(self):
    return f"({self.real}{self.imag:+}j)"

Математически можно сложить любые два комплексных числа или сложить действительное число с комплексным. Давайте настроим оператор + таким образом, чтобы он работал в обоих случаях.

Метод проверяет тип оператора правой части. Если это int или float, он увеличит только вещественную часть (поскольку любое вещественное число a эквивалентно a+0j), а в случае другого комплексного числа он изменит обе части:

def __add__(self, other):
    if isinstance(other, float) or isinstance(other, int):
        real_part = self.real + other
        imag_part = self.imag

    if isinstance(other, CustomComplex):
        real_part = self.real + other.real
        imag_part = self.imag + other.imag

    return self.__class__(real_part, imag_part)

Аналогичным образом мы определяем поведение для - и *:

def __sub__(self, other):
    if isinstance(other, float) or isinstance(other, int):
        real_part = self.real - other
        imag_part = self.imag

    if isinstance(other, CustomComplex):
        real_part = self.real - other.real
        imag_part = self.imag - other.imag

    return self.__class__(real_part, imag_part)

def __mul__(self, other):
    if isinstance(other, int) or isinstance(other, float):
        real_part = self.real * other
        imag_part = self.imag * other

    if isinstance(other, CustomComplex):
        real_part = (self.real * other.real) - (self.imag * other.imag)
        imag_part = (self.real * other.imag) + (self.imag * other.real)

    return self.__class__(real_part, imag_part)

Поскольку сложение и умножение коммутативны, мы можем определить их обратные операторы, вызвав __add__() и __mul__() в __radd__() и __rmul__() соответственно. С другой стороны, необходимо определить поведение __rsub__(), поскольку вычитание не является коммутативным:

def __radd__(self, other):
    return self.__add__(other)

def __rmul__(self, other):
    return self.__mul__(other)

def __rsub__(self, other):
    # x - y != y - x
    if isinstance(other, float) or isinstance(other, int):
        real_part = other - self.real
        imag_part = -self.imag

    return self.__class__(real_part, imag_part)

Примечание: Вы могли заметить, что мы не добавили сюда конструкцию для обработки экземпляра CustomComplex. Это потому, что в таком случае оба операнда являются экземплярами нашего класса, и __rsub__() не будет отвечать за обработку операции. Вместо этого будет вызван __sub__(). Это тонкая, но важная деталь.

Теперь мы позаботимся о двух операторах, == и !=. Для них используются специальные методы __eq__() и __ne__() соответственно. Два комплексных числа считаются равными, если их соответствующие действительная и мнимая части равны. О них говорят, что они неравны, если одно из них неравно:

def __eq__(self, other):
    # Note: generally, floats should not be compared directly
    # due to floating-point precision
    return (self.real == other.real) and (self.imag == other.imag)

def __ne__(self, other):
    return (self.real != other.real) or (self.imag != other.imag)

Примечание:The Floating-Point Guide - это статья, в которой рассказывается о сравнении плавающих чисел и точности плавающей точки. В ней подчеркиваются предостережения, связанные с непосредственным сравнением плавающих точек, что мы и делаем здесь.

Также можно возвести комплексное число в любую степень с помощью простой формулы . Мы настраиваем поведение как встроенного оператора pow(), так и оператора ** с помощью специального метода __pow__():

def __pow__(self, other):
    r_raised = abs(self) ** other
    argz_multiplied = self.argz() * other

    real_part = round(r_raised * cos(argz_multiplied))
    imag_part = round(r_raised * sin(argz_multiplied))

    return self.__class__(real_part, imag_part)

Примечание: Внимательно посмотрите на определение метода. Мы вызываем abs(), чтобы получить модуль комплексного числа. Таким образом, как только вы определили специальный метод для определенной функции или оператора в своем классе, его можно использовать в других методах того же класса.

Создадим два экземпляра этого класса, один с положительной мнимой частью, другой с отрицательной мнимой частью:

>>> a = CustomComplex(1, 2)
>>> b = CustomComplex(3, -4)

Представления строк:

>>> a
CustomComplex(1, 2)
>>> b
CustomComplex(3, -4)
>>> print(a)
(1+2j)
>>> print(b)
(3-4j)

Восстановление объекта с помощью eval() с repr():

>>> b_copy = eval(repr(b))
>>> type(b_copy), b_copy.real, b_copy.imag
(__main__.CustomComplex, 3, -4)

Сложение, вычитание и умножение:

>>> a + b
CustomComplex(4, -2)
>>> a - b
CustomComplex(-2, 6)
>>> a + 5
CustomComplex(6, 2)
>>> 3 - a
CustomComplex(2, -2)
>>> a * 6
CustomComplex(6, 12)
>>> a * (-6)
CustomComplex(-6, -12)

Проверки равенств и неравенств:

>>> a == CustomComplex(1, 2)
True
>>> a ==  b
False
>>> a != b
True
>>> a != CustomComplex(1, 2)
False

Наконец, возведение комплексного числа в некоторую степень:

>>> a ** 2
CustomComplex(-3, 4)
>>> b ** 5
CustomComplex(-237, 3116)

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

 

from math import hypot, atan, sin, cos

class CustomComplex():
    """
    A class to represent a complex number, a+bj.
    Attributes:
        real - int, representing the real part
        imag - int, representing the imaginary part

    Implements the following:

    * Addition with a complex number or a real number using `+`
    * Multiplication with a complex number or a real number using `*`
    * Subtraction of a complex number or a real number using `-`
    * Calculation of absolute value using `abs`
    * Raise complex number to a power using `**`
    * Nice string representation using `__repr__`
    * Nice user-end viewing using `__str__`

    Notes:
        * The constructor has been intentionally kept simple
        * It is configured to support one kind of call:
            CustomComplex(a, b)
        * Error handling was avoided to keep things simple
    """

    def __init__(self, real, imag):
        """
        Initializes a complex number, setting real and imag part
        Arguments:
            real: Number, real part of the complex number
            imag: Number, imaginary part of the complex number
        """
        self.real = real
        self.imag = imag

    def conjugate(self):
        """
        Returns the complex conjugate of a complex number
        Return:
            CustomComplex instance
        """
        return CustomComplex(self.real, -self.imag)

    def argz(self):
        """
        Returns the argument of a complex number
        The argument is given by:
            atan(imag_part/real_part)
        Return:
            float
        """
        return atan(self.imag / self.real)

    def __abs__(self):
        """
        Returns the modulus of a complex number
        Return:
            float
        """
        return hypot(self.real, self.imag)

    def __repr__(self):
        """
        Returns str representation of an instance of the 
        class. Can be used with eval() to get another 
        instance of the class
        Return:
            str
        """
        return f"CustomComplex({self.real}, {self.imag})"


    def __str__(self):
        """
        Returns user-friendly str representation of an instance 
        of the class
        Return:
            str
        """
        return f"({self.real}{self.imag:+}j)"

    def __add__(self, other):
        """
        Returns the addition of a complex number with
        int, float or another complex number
        Return:
            CustomComplex instance
        """
        if isinstance(other, float) or isinstance(other, int):
            real_part = self.real + other
            imag_part = self.imag

        if isinstance(other, CustomComplex):
            real_part = self.real + other.real
            imag_part = self.imag + other.imag

        return CustomComplex(real_part, imag_part)

    def __sub__(self, other):
        """
        Returns the subtration from a complex number of
        int, float or another complex number
        Return:
            CustomComplex instance
        """
        if isinstance(other, float) or isinstance(other, int):
            real_part = self.real - other
            imag_part = self.imag

        if isinstance(other, CustomComplex):
            real_part = self.real - other.real
            imag_part = self.imag - other.imag

        return CustomComplex(real_part, imag_part)

    def __mul__(self, other):
        """
        Returns the multiplication of a complex number with
        int, float or another complex number
        Return:
            CustomComplex instance
        """
        if isinstance(other, int) or isinstance(other, float):
            real_part = self.real * other
            imag_part = self.imag * other

        if isinstance(other, CustomComplex):
            real_part = (self.real * other.real) - (self.imag * other.imag)
            imag_part = (self.real * other.imag) + (self.imag * other.real)

        return CustomComplex(real_part, imag_part)

    def __radd__(self, other):
        """
        Same as __add__; allows 1 + CustomComplex('x+yj')
        x + y == y + x
        """
        pass

    def __rmul__(self, other):
        """
        Same as __mul__; allows 2 * CustomComplex('x+yj')
        x * y == y * x
        """
        pass

    def __rsub__(self, other):
        """
        Returns the subtraction of a complex number from
        int or float
        x - y != y - x
        Subtration of another complex number is not handled by __rsub__
        Instead, __sub__ handles it since both sides are instances of
        this class
        Return:
            CustomComplex instance
        """
        if isinstance(other, float) or isinstance(other, int):
            real_part = other - self.real
            imag_part = -self.imag

        return CustomComplex(real_part, imag_part)

    def __eq__(self, other):
        """
        Checks equality of two complex numbers
        Two complex numbers are equal when:
            * Their real parts are equal AND
            * Their imaginary parts are equal
        Return:
            bool
        """
        # note: comparing floats directly is not a good idea in general
        # due to floating-point precision
        return (self.real == other.real) and (self.imag == other.imag)

    def __ne__(self, other):
        """
        Checks inequality of two complex numbers
        Two complex numbers are unequal when:
            * Their real parts are unequal OR
            * Their imaginary parts are unequal
        Return:
            bool
        """
        return (self.real != other.real) or (self.imag != other.imag)

    def __pow__(self, other):
        """
        Raises a complex number to a power
        Formula:
            z**n = (r**n)*[cos(n*agrz) + sin(n*argz)j], where
            z = complex number
            n = power
            r = absolute value of z
            argz = argument of z
        Return:
            CustomComplex instance
        """
        r_raised = abs(self) ** other
        argz_multiplied = self.argz() * other

        real_part = round(r_raised * cos(argz_multiplied))
        imag_part = round(r_raised * sin(argz_multiplied))

        return CustomComplex(real_part, imag_part)

Выводы

В этом уроке вы узнали о модели данных Python и о том, как модель данных может быть использована для создания классов Pythonic. Вы узнали об изменении поведения встроенных функций, таких как len(), abs(), str(), bool() и так далее. Вы также узнали об изменении поведения встроенных операторов, таких как +, -, *, ** и так далее.

Для получения дополнительной информации о модели данных, а также о перегрузке функций и операторов, посмотрите эти ресурсы:

Вернуться на верх