substituting iadd doesn't work as expected for += operator

I have an object from a library (numpy.ndarray), in which I've substituted the _iadd_ method for a custom one. If I call object._iadd_(x), it works as expected. However, object+=x seems to call the old (unsubstituted) method. I wanted to prevent overflows on numpy from occurring on specific cases, so I created a context manager for that. Here's the (still very crude) code:

class NumpyOverflowPreventer( object ):
    inverse_operator= {'__iadd__':'__sub__', '__isub__':'__add__', '__imul__': '__div__', '__idiv__':'__mul__'}

    def _operate(self, b, forward_operator):
        assert type(b) in (int, float)
        reverse_operator= NumpyOverflowPreventer.inverse_operator[forward_operator]
        uro= getattr(self.upper_range, reverse_operator)
        lro= getattr(self.lower_range, reverse_operator)
        afo= self.originals[ forward_operator ]
        overflows= self.matrix > uro( b )
        underflows= self.matrix < lro( b )
        afo( b )
        self.matrix[overflows]= self.upper_range
        self.matrix[underflows]= self.lower_range
        
    def __init__(self, matrix):
        m= matrix
        assert m.dtype==np.uint8
        self.matrix= m
        self.lower_range= float(0)
        self.upper_range= float(2**8-1)
        
    def __enter__(self):
        import functools
        self.originals={}
        for op in NumpyOverflowPreventer.inverse_operator.keys():
            self.originals[ op ] = getattr( self.matrix, op )
            setattr( self.matrix, op, functools.partial(self._operate, forward_operator=op))
    
    def __exit__(self, type, value, tb):
        for op in NumpyOverflowPreventer.inverse_operator.keys():
            setattr( self.matrix, op, self.originals[ op ] )

running this:

a= np.matrix(255, dtype= np.uint8)
b= np.matrix(255, dtype= np.uint8)
with NumpyOverflowPreventer(a):
    a+=1
with NumpyOverflowPreventer(b):
    b.__iadd__(1)
print a,b

returns this:

[[0]] [[255]]

Solution

The issue you are seeing is that the special built-in methods are not looked up on the instance. They are looked up on the matrix type. So replacing them on the instance will not cause them to be used indirectly.

One way to achieve your goal is to instead make NumpyOverflowPreventer a wrapper for the operations you want to address...

import numpy as np 
import sys

class NumpyOverflowPreventer(object):

    inverse_operator= { 
        '__iadd__': '__sub__', 
        '__isub__': '__add__', 
        '__imul__': '__div__', 
        '__idiv__': '__mul__'
    }

    def __init__(self, matrix):
        m = matrix
        assert m.dtype==np.uint8
        self.matrix = m
        self.lower_range = float(0)
        self.upper_range = float(2**8-1)

    def __iadd__(self, v):
        # dynamic way to get the name "__iadd__"
        self._operate(v, sys._getframe().f_code.co_name)
        return self

    def _operate(self, b, forward_operator):
        assert type(b) in (int, float)
        reverse_operator = self.inverse_operator[forward_operator]
        uro= getattr(self.upper_range, reverse_operator)
        lro= getattr(self.lower_range, reverse_operator)
        afo= getattr(self.matrix, forward_operator)
        overflows= self.matrix > uro( b )
        underflows= self.matrix < lro( b )
        afo( b )
        self.matrix[overflows]= self.upper_range
        self.matrix[underflows]= self.lower_range

I have only defined __iadd__ here, and I am sure you could do all of them dynamically with some metaclass/decorator action...but I am keeping it simple.

Usage:

a = np.matrix(255, dtype= np.uint8)
b = np.matrix(255, dtype= np.uint8)

p = NumpyOverflowPreventer(a)
p+=1

p = NumpyOverflowPreventer(b)
p.__iadd__(1)

print a,b
# [[255]] [[255]]

substituting __iadd__ doesn't work as expected for += operator

substituting iadd doesn't work as expected for += operator