CODEWITHSUNDEEP
pythonoperator-overloading
pteehan
pteehan

Reputation: 817

python - overloading several operators at once

I have a custom class and I want to overload several artihmetic operators, and wonder if there is a way to avoid having to write out the code for each one individually. I haven't been able to find any examples that don't explicity overload each operator one-by-one.

class Foo(object):
    a=0 

    def __init__(self, a):
        self.a=a            

    def __add__(self, other):
        #common logic here
        return Foo(self.a+other.a)

    def __sub__(self, other):
        #common logic here  
        return Foo(self.a-other.a)

    def __mul__(self, other):
        #common logic here
        return Foo(self.a*other.a)

#etc...

The logic is slightly more complex than this, but the common pattern is that each operator overload method contains some identical code to check that the operation is allowed, and then constructs an operation using the class members. I want to reduce the redundant code. This works: 

class Foo(object):
    a=0 

    def __init__(self, a):
        self.a=a            

    def operate(self, other, operator):
        #common logic here
        a = constructOperation(self.a, other.a, operator)
        return Foo(a)

    def __add__(self, other):
        return self.operate(other, "+")

    def __sub__(self, other):       
        return self.operate(other, "-")     


def constructOperation(operand0, operand1, operator):
    if operator=="+": 
        return operand0 + operand1
    if operator=="-": 
        return operand0 - operand1

But it seems kind of silly to be constructing operations manually like that. Does this approach make sense, or is there a better way here?

Upvotes: 3

Views: 1632

Answers (4)

Bakuriu
Bakuriu

Reputation: 101919

You must define the methods. This is because python does a special look-up when calling special methods, and thus something like this:

import operator

class Foo(object):

    def __init__(self, a):
        self.a = a
    def __getattr__(self, attr):
        try:
            func = getattr(operator, attr)
            return lambda x: Foo(func(self.a, x.a))
        except AttributeError:
            raise AttributeError(attr)

Doesn't work:

>>> f = Foo(1)
>>> g = Foo(3)
>>> f * g     #no __mul__?
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
TypeError: unsupported operand type(s) for *: 'Foo' and 'Foo'
>>> getattr(f, '__mul__')
<function <lambda> at 0x2975140>
>>> f.__mul__(g)
<__main__.Foo object at 0x2922ed0>
>>> f.__mul__(g).a
3

The "best" you can do is to use Antimony's solution, which is the most DRY.

Upvotes: 1

shx2
shx2

Reputation: 64288

I don't know if there's a way to avoid defining all (or at least most) of the operators. It makes sense that there isn't one. After all, there is no one single way to define __sub__ given __add__ (and __mul__). However, one improvement would be to pass callables to constructOperation() instead of the symbolic operators.

E.g.

class Foo(object):
    a=0 

    def __init__(self, a):
        self.a=a            

    def operate(self, other, operator):
        #common logic here
        a = constructOperation(self.a, other.a, operator)
        return Foo(a)

    def __add__(self, other):
        return self.operate(other, sum)

    def __sub__(self, other):       
        return self.operate(other, lambda x, y: x - y)     


def constructOperation(operand0, operand1, operator):
    return operator(operand0, operand1)

Upvotes: 1

NPE
NPE

Reputation: 500227

I would just use the operator module:

import operator

class Foo(object):

    a=0 

    def __init__(self, a):
        self.a=a            

    def operate(self, other, op):
        #common logic here
        return Foo(op(self.a, other.a))

    def __add__(self, other):
        return self.operate(other, operator.add)

    def __sub__(self, other):       
        return self.operate(other, operator.sub)

Upvotes: 5

Antimony
Antimony

Reputation: 39451

You can do it via reflection and higher order functions, though this may not play well with inheritance.

import operator

def apply_a(func):
    def inner(self, other):
        return Foo(func(self.a, other.a))
    return inner

class Foo(object):
    def __init__(self, a=0):
        self.a = a

for name in ['__add__','__mul__','__sub__']:
    setattr(Foo, name, apply_a(getattr(operator, name)))

Upvotes: 7

Related Questions