CODEWITHSUNDEEP
pythonpython-3.x
eto3542
eto3542

Reputation: 91

How can I use a method as a decorator in python?

I have a class that utilizes a connection to a database to interact with the data like this:

class DataBase:
    def __init__(self, user, password):
        # open connection in init to have only one active connection per class instance
        self.connection = Connection(user, password)

    def get(self):
        with self.connection:
            # get data magic
            pass
    
    def post(self):
        with self.connection:
            # post data magic
            pass

To avoid writing with self.connection for all methods, I would like to add a decorator like this:

class DataBase:
    def __init__(self, user, password):
        # open connection in init to have only one active connection per class instance
        self.connection = Connection(user, password)
    
    def _with_connection(self, func):
        def wrapper(*args, **kwargs):
            with self.connection:
                result = func(*args, **kwargs)
            return result
        return wrapper
    
    @_with_connection
    def get(self):
        # get data magic
        pass
        
    @_with_connection
    def post(self):
        # post data magic
        pass

However, this obviously raises the following problem: TypeError: _with_connection() missing 1 required positional argument: 'func'

So how could this be implemented?

Upvotes: 0

Views: 46

Answers (1)

chepner
chepner

Reputation: 530960

The problem is that you aren't using _with_connection as an instance method (never mind that you can't, because no instance of Database exists yet), so it shouldn't be defined like one. Get rid of the self parameter. It is used in wrapper, so you can define that separately from *args.

class DataBase:
    def __init__(self, user, password):
        # open connection in init to have only one active connection per class instance
        self.connection = Connection(user, password)
    
    def _with_connection(func):
        def wrapper(self, *args, **kwargs):
            with self.connection:
                result = func(self, *args, **kwargs)
            return result
        return wrapper
    
    @_with_connection
    def get(self):
        # get data magic
        pass
        
    @_with_connection
    def post(self):
        # post data magic
        pass

I don't see this as any kind of improvement over your original code, though. You aren't saving any lines of code (replacing the with statement with a decoration), and you make it less obvious how get and post actually use the connection.

Upvotes: 2

Related Questions