Effective way to “wrap” subclass methods

Question

The following code has the desired functionality:

class MainClass:
    def __init__(self, worker):
        self.worker = worker

    def foo(self):
        self.worker.run()

class Worker:
    def run(self):
        print('about to run...')
        self._run()
        print('just ran...')

    def _run(self):
        print('running...')

class SpecialWorker(Worker)
    def _run(self):
        print('running very fast...')

The idea is that I have a generic worker class which implements an important _run method. However, external processes access _run through the run method of the Worker class which basically just wraps _run . The reason I work this way is that subclasses to Worker will have the appropriate wrapping no matter how _run is defined, but also people developing new subclasses do not need to take care to explicitly include the wrapping code.

My question is if there is a better way to do this that doesn't involve splitting the run method out into two methods, one of which calls the other. What I basically want is a way that the Worker run method will be wrapped with certain functionality in all subclasses of Worker .

Any suggestions other that what I've presented?

Answer 1

Whenever one function needs to be wrapped by another in Python, decorators come to mind. Outside of classes they are more straightforward, but they could be used in this situation as well. I'm not sure the approach is more elegant than your original solution, but here is what a decorator based approach might look like:

class MainClass:
    def __init__(self, worker):
        self.worker = worker

    def foo(self):
        self.worker.run()

class Worker:    
    @staticmethod
    def orchestrate_run(func):
        def wrapped_run(self):
            print('about to run...')
            func(self)
            print('just ran...')
        return wrapped_run
    
class SpecialWorker(Worker):
    @Worker.orchestrate_run
    def run(self):
        print('running very fast...')

worker = SpecialWorker()
mc = MainClass(worker)
mc.foo()

about to run...
running very fast...
just ran...