CODEWITHSUNDEEP
pythonpython-3.xooppython-decoratorscontextmanager
Robert Axe
Robert Axe

Reputation: 414

Modifying class attribute by calling class method from inside decorator

I have quite a few methods that requires opening some file (each method has different), doing some stuff, closing it. Althou it is possible to open and close a file within each method, I am wondering whether it is possible to do it through decorators so that a decorator could take method as a parameter which would load data into obj attribute that the decorated class is acting on.

class Foo:
    def __init__(self, fname):
        self.fname = fname

    def load_file(self):
        with open(self.fname, 'rb') as f:
            self.file_ = pickle.load(f)
    
    def do_some_work(self):
        self.load_file()
        ... # some calculation and so on
        delattr(self, 'file_')

    @loader(self.load_file)
    def do_some_work_decorated(self):
         ... # only some calculation and so on, file loading is done by the defined method in decorator

My question is whether this is even possible and are there better ways to approach it?

Upvotes: 1

Views: 94

Answers (1)

martineau
martineau

Reputation: 123491

Well, you can do something similar to that, but you'd need to pass the name the method to the decorator because there's no self for it to refer to when it's used to decorate the method(s) — which happens when the class is defined not when the code in it is executed afterwards.

Here's what I mean:

import pickle

def loader(do_stuff):
    def decorator(method):
        def decorated(self, *args, **kwargs):
            getattr(self, do_stuff)()  # Call specified class method.
            return method(self, *args, **kwargs)  # Then call decorated method.
        return decorated
    return decorator


class Foo:
    def __init__(self, fname):
        self.fname = fname

    def load_file(self):
        with open(self.fname, 'rb') as f:
            self.file_ = pickle.load(f)

    def do_some_work(self):
        self.load_file()
        ... # some calculation and so on
        delattr(self, 'file_')

    @loader('load_file')
    def do_some_work_decorated(self):
        ... # only some calculation and so on, file loading is done by the defined method in decorator
        print(f'{self.file_}')

if __name__ == '__main__':

    # Create a test file.
    with open('foo.pkl', 'wb') as outp:
        pickle.dump(42, outp)

    # See if decorator worked.
    foo = Foo('foo.pkl')
    foo.do_some_work_decorated()  # -> 42

Upvotes: 1

Related Questions