decorate method after declaration

Question

I have a class from a library, let say MyClass with a method my_method.

from MyLibrary import MyClass
obj = MyClass()
result = obj.my_method('arg')

I can't thouch the original implementation, but I want to decorate the my_method, in particular I want to check if the returning value is not None and in case to raise an exception, as:

result = obj.my_method('arg')
if result is None: raise ValueError()

this is quite boring since I need to do it many times. In particular I don't want to change the code when I am calling obj.my_method('arg'). I am trying to:

def safe_method(self, *args):
  result = self.my_method(*args)
  if result is None: raise ValueError()
  return result

MyClass.my_method = safe_method

this is not working:

RuntimeError: maximum recursion depth exceeded while calling a Python object

I understand why, but I can't solve the problem

Answer 1

You can decorate it like this:

def safe_decorator(func):
    def wrapper(self, *args):
        result = func(self, *args)
        if result is None:
            raise ValueError()
        return result
    return wrapper

MyClass.my_method = safe_decorator(MyClass.my_method)

You can then reuse this decorator for whatever method you want from whatever class definition. This is what the @ notation does more or less :)

Hope this helps!

Answer 2

def safe_method(self, *args):
  result = self.my_method(*args)
  if result is None: raise ValueError()
  return result

MyClass.my_method = safe_method

What is happening above is you are replacing MyClass.my_method with another method that in turn calls my_method. Hence, you are getting a recursion error.

Try this:

MyClass._my_method = MyClass.my_method

def safe_method(self, *args):
    result = self._my_method(*args)
    if result is None:
        raise ValueError()
    return result

MyClass.my_method = safe_method