How to implement @dataclass to define arithmetic operations in Python?

1.1k Views Asked by At

I'm learning Python on my own and I found a task that requires using a decorator @dataclass to create a class with basic arithmetic operations.

from dataclasses import dataclass
from numbers import Number

@dataclass
class MyClass:
    x: float
    y: float
    def __add__(self, other):
         match other:
            case Number():
                return MyClass(float(other) + self.x, self.y)    
            case MyClass(ot_x, ot_y):
                return MyClass(self.x + ot_x, self.y + ot_y)
              
    __radd__ = __add__        

I have implemented the addition operation. But I also need to do the operations of subtraction __sub__, multiplication __mul__, division __truediv__, negation __neg__, also __mod__ and __pow__. But I couldn't realize these operations. The main thing for me is to use the construction match/case. Maybe there are simpler ways to create it. I will be glad of your help.

1

There are 1 best solutions below

0
ShadowRanger On

If you're trying to make a complete numeric type, I strongly suggest checking out the implementation of the fractions.Fraction type in the fractions source code. The class was intentionally designed as a model for how you'd overload all the pairs of operators needed to implement a numeric type at the Python layer (it's explicitly pointed out in the numbers module's guide to type implementers).

The critical parts for minimizing boilerplate begin with the definition of the _operator_fallbacks utility function within the class (which is used to take a single implementation of the operation and the paired operator module function representing it, and generate the associated __op__ and __rop__ operators, being type strict for the former and relaxed for the latter, matching the intended behavior of each operator based on whether it's the first chance or last chance to implement the method).

It's far too much code to include here, but to show how you'd implement addition using it, I'll adapt your code to call it (you'd likely use a slightly different implementation of _operator_fallbacks, but the idea is the same):

import operator

# Optional, but if you want to act like built-in numeric types, you
# should be immutable, and using slots (if you can rely on Python 3.10+)
# dramatically reduce per-instance memory overhead
# Pre-3.10, since x and y don't have defaults, you could define __slots__ manually
@dataclass(frozen=True, slots=True)
class MyClass:
    x: float
    y: float

    # _operator_fallbacks defined here
    # When it received a non-MyClass, it would construct a MyClass from it, e.g.
    # to match your original code it would construct it as MyClass(float(val), 0)
    # and then invoke the monomorphic_operator, e.g. the _add passed to it below
    # or, if the type did not make sense to convert to MyClass, but it made sense
    # to convert the MyClass instance to the other type, it would do so, then use the
    # provided fallback operator to perform the computation
    # For completely incompatible types, it just returns NotImplemented

    def _add(a, b):
        """a + b"""
        return MyClass(a.x + b.x, a.y + b.y)  # _operator_fallback has already coerced the types appropriately

    __add__, __radd__ = _operator_fallbacks(_add, operator.add)

By putting the ugliness of type-checking and coercion in common code found in _operator_fallbacks, and putting only the real work of addition in _add, it avoids a lot of every-operator-overload boilerplate (as you can see here; _operator_fallbacks will be a page of code to make the forward and reverse functions and return them, but each new operator is only a few lines, defining the monomorphic operator and calling _operator_fallbacks to generate the __op__/__rop__ pair.