decorating class method with arguments

Question

How do I decorate a class method with arguments? Current code is

def establish_con(func):

    con, meta = db.connect(config.user, config.password, config.db)
    meta.reflect(bind=con)

    def inner(self, query, con, *args, **kwargs):
        return func(self, query, con, *args, **kwargs)

    con.close()
    return inner

class DataReader:
    def __init__(self):
        self.data = {}

    @establish_con
    def execQuery(self, query, con):
        # con, meta = db.connect(config.user, config.password, config.db)
        # meta.reflect(bind=con)

        result = pd.read_sql(query, con)

        # con.close()
        return result

test = DataReader()
df = test.execQuery("Select * from backtest limit 10")
print(df)

Currently, the first argument seems to be the class instance. I tried different variations of the code, but always ran into too many/not enough/undefined arguments problems.

I've read other posts like

How do I pass extra arguments to a Python decorator?

Decorators with parameters?

and others, but still can't figure it out.

Edit: Not a duplicate of Python decorators in classes as in this answer no arguments need to be passed to the function.

Answer 1

Adapting answer by @madjardi from this post to use arguments.

import functools

class Example:

    con = "Connection"

    def wrapper(func):
        @functools.wraps(func)
        def wrap(self, *args, **kwargs):
            print("inside wrap")
            return func(self, Example.con, *args, **kwargs) # use *args to pass objects down
        return wrap

    @wrapper
    def method(self, con, arg): 
        print("METHOD {0} {1}".format(con, arg))

    wrapper = staticmethod(wrapper)


e = Example()
e.method(1) # METHOD Connection 1

Answer 2

You should replace:

def inner(self, query, con, *args, **kwargs):
        return func(self, query, con, *args, **kwargs)

With:

def inner(self, query, *args, **kwargs):   # no 'con' here
        return func(self, query, con, *args, **kwargs) 

A minimal working example using your strategy would be (achieving nothing like your goal):

def add_arg(method):

    def decorated_method(self, *args):
        return method(self, 10, *args)

    return decorated_method

class Data:
    @add_arg
    def summation(self, *args):
        return sum(args)


d = Data()
print(d.summation(1, 2))  # prints 13, not 3

Answer 3

Here's how I think you would need to do it. I've split the code up into two separate blocks in an attempt to make it clearer.

The first part of just establishes a minimal scaffold to make it possible to run (and follow the execution of) the code in the second block and keep it as close as possible to what's in your question. It's not terribly important unto itself.

Decorators with arguments effectively become decorator-factories — in the sense they have to create a return a decorator function that that will then be applied to the target function or method.

# Scaffolding
class Pandas:
    @staticmethod
    def read_sql(query, con):
        print(f'in Pandas read_sql({query!r}, {con})')
        return 'pandas_dataframe'


class Connection:
    def close(self):
        print('in Connection.close()')

    def __repr__(self):
        return '<Connection object>'

con = Connection()


class Meta:
    def reflect(self, bind=con):
        print(f'in Meta.reflect(bind={bind}')


class Database:
    def connect(self, user, password, db):
        print(f'in Database.connect({user}, {password}, {db})')
        return Connection(), Meta()

    def __repr__(self):
        return '<Database object>'

class Config:
    def __init__(self, user, password, db):
        self.user = user
        self.password = password
        self.db = db

# Set up a framework for testing.
pd = Pandas()
meta = Meta()
db = Database()
config = Config('username', 'secret', db)

With that environment established, here's how the decorator could be written.

# Decorator
def establish_con(config):

    def wrapper(method):

        def wrapped(*args, **kwargs):
            con, meta = db.connect(config.user, config.password, config.db)
            meta.reflect(bind=con)

            args = args + (con,)
            result = method(*args, **kwargs)

            con.close()

            return result

        return wrapped

    return wrapper


class DataReader:
    def __init__(self):
        self.data = {}

    @establish_con(config)
    def execQuery(self, query, con):
        return pd.read_sql(query, con)


test = DataReader()
df = test.execQuery("Select * from backtest limit 10")
print(df)

Output:

in Database.connect(username, secret, <Database object>)
in Meta.reflect(bind=<Connection object>
in Pandas read_sql('Select * from backtest limit 10', <Connection object>)
in Connection.close()
pandas_dataframe