Wrapping only the 'most recent' implementation of a method in a class?

Question

I've implemented automatic wrapping of a method as described in how-to-wrap-every-method-of-a-class in the following manner:

from functools import wraps
from somewhere import State

def wrapper(method):
    @wraps(method)
    def wrapped(_instance, *args, **kwargs):
        _instance.state = State.Running
        method(_instance, *args, **kwargs)
        _instance.state = State.Finished
    return wrapped

class MetaClass(type):
    def __new__(meta, classname, bases, classDict):
        newClassDict = {}
        for attributeName, attribute in classDict.items():
            if isinstance(attribute, FunctionType) and attributeName == 'run':
                # replace it with a wrapped version
                attribute = wrapper(attribute)
            newClassDict[attributeName] = attribute
        return type.__new__(meta, classname, bases, newClassDict)

class AmazingClass(metaclass=MetaClass):
    def __init__(self):
        self.state = State.Idle

    def run(self, *args, **kwargs):
        print(self.state)  # State.Running
        do_something()
        # After returning: self.state == State.Finished

An issue arises when deriving from AmazingClass and overriding run()

class EvenMoreAmazingClass(AmazingClass):
    def run(self, *args, **kwargs):
        print(self.state)  # State.Running
        super().run(*args, **kwargs)
        print(self.state)  # State.Finished set by base-implementation of run (undesired) !!!!
        do_something_else()
        # self.state should only reach State.Finished after returning from here

The problem is that the base-implementation of run() already sets self.state to State.Finished.
I want to call the base-implementation of run() and also not have to do any trickery in any classes derived from AmazingClass to keep self.state == State.Running.

So the question arises, is there a way to detect and wrap only the outermost implementation of run()? (In this case EvenMoreAmazingClass.run())

Answer 1

If you'll never have anything other than the run method wrapped with wrapper you can implement a StateHandler class to handle the inner calls.

from functools import wraps
from types import FunctionType

class State:
    Running = 'State.Running'
    Finished = 'State.Finished'
    Idle = 'State.Idle'

class StateHandler:
    state = 0

    @staticmethod
    def incr():
        StateHandler.state += 1
        # always return State.running when calling StateHandler.incr()
        return State.Running 

    @staticmethod
    def decr():
        StateHandler.state -= 1
        # only return State.Finished when the outer 
        # method has completed. 
        if StateHandler.state:
            return State.Running
        return State.Finished

def wrapper(method):
    @wraps(method)
    def wrapped(self, *args, **kwargs):
        self.state = StateHandler.incr()
        method(self, *args, **kwargs)
        self.state = StateHandler.decr()
    return wrapped

class MetaClass(type):
    def __new__(meta, class_name, bases, dct):
        new_dct = {}
        for name, attribute in dct.items():
            if isinstance(attribute, FunctionType) and name == 'run':
                attribute = wrapper(attribute)
            new_dct[name] = attribute
        return type.__new__(meta, class_name, bases, new_dct)


class AmazingClass(metaclass=MetaClass):
    def __init__(self):
        self.state = State.Idle

    def run(self, *args, **kwargs):
        print(self.state)  # State.Running


class EvenMoreAmazingClass(AmazingClass):
    def run(self, *args, **kwargs):
        print(self.state)  # State.Running
        super().run(*args, **kwargs)
        print(self.state)


foo = EvenMoreAmazingClass()
print(foo.state)
foo.run()
print(foo.state)

I added a toy State implementation to make the example runnable.

This works by incrementing StateHandler.state every time a run method is called and decrementing when it ends. So StateHandler.state is 1 at the outer call, 2 at the inner call, 1 at the inner return and 0 at the outer return.

This effectively causes State.Finished to be set ONLY when the initial run method that was called has returned. So it works for both your base class and derived class implementations.

Answer 2

In MetaClass.__new__, you should keep a copy of the non-wrapped method, name that copy _run.

Then when you call super(), you can use that copy instead of the wrapped run.

Full code:

from functools import wraps


class Wild:
    def __getattr__(self, attr):
        return attr

State = Wild()

def do_something(*args):
    print('Doing something', *args)


def wrapper(method):
    @wraps(method)
    def wrapped(_instance, *args, **kwargs):
        _instance.state = State.Running
        method(_instance, *args, **kwargs)
        _instance.state = State.Finished
    return wrapped


class MetaClass(type):
    def __new__(meta, classname, bases, classDict):
        newClassDict = {}
        for attributeName, attribute in classDict.items():
            if callable(attribute) and attributeName == '_run':
                newClassDict['run'] = wrapper(attribute)
            newClassDict[attributeName] = attribute
        print(classDict, newClassDict)
        return type.__new__(meta, classname, bases, newClassDict)


class AmazingClass(metaclass=MetaClass):
    def __init__(self):
        self.state = State.Idle

    def _run(self, *args, **kwargs):
        print(self.state)  # State.Running
        do_something()
        # After returning: self.state == State.Finished


class EvenMoreAmazingClass(AmazingClass):
    def _run(self, *args, **kwargs):
        print(self.state)  # State.Running
        super()._run(*args, **kwargs)
        print(self.state)  # State.Running -> this is what you want
        do_something('else')


ddd = EvenMoreAmazingClass()
ddd.run()
print(ddd.state)  # Finished