Python decorate methods with variable number of positional args and optional arg

Question

I am writing my first Python (3.4) application using SQLalchemy. I have several methods which all have a very similar pattern. They take an optional argument session which defaults to None. If session is passed, the function uses that session, otherwise it opens and uses a new session. For example, consider the following method:

def _stocks(self, session=None):
    """Return a list of all stocks in database."""
    newsession = False
    if not session:
        newsession = True
        session = self.db.Session()
    stocks = [stock.ticker for stock in session.query(Stock).all()]
    if newsession:
        session.close()
    return stocks

So, being new to Python and eager to learn all of its power, I thought this smelt like the perfect time to learn a little something about Python decorators. So after a lot of reading, like this this series of blog posts and this fantastic SO answer, I wrote the following decorator:

from functools import wraps

def session_manager(func):
    """
    Manage creation of session for given function.

    If a session is passed to the decorated function, it is simply
    passed through, otherwise a new session is created.  Finally after
    execution of decorated function, the new session (if created) is
    closed/
    """
    @wraps(func)
    def inner(that, session=None, *args, **kwargs):
        newsession = False
        if not session:
            newsession = True
            session = that.db.Session()
        func(that, session, *args, **kwargs)
        if newsession:
            session.close()
        return func(that, session, *args, **kwargs)
    return inner

And it seems to work great. The original method is now reduced to:

@session_manager
def _stocks(self, session=None):
    """Return a list of all stocks in database."""
    return [stock.ticker for stock in session.query(Stock).all()]

HOWEVER, when I apply the decorator to a function that takes some positional arguments in addition to the optional session, I get an error. So trying to write:

@session_manager
def stock_exists(self, ticker, session=None):
    """
    Check for existence of stock in database.

    Args:
        ticker (str): Ticker symbol for a given company's stock.
        session (obj, optional): Database session to use.  If not
            provided, opens, uses and closes a new session.

    Returns:
        bool: True if stock is in database, False otherwise.
    """
    return bool(session.query(Stock)
                .filter_by(ticker=ticker)
                .count()
                )

and running like print(client.manager.stock_exists('AAPL')) gives an AttributeError with the following traceback:

Traceback (most recent call last):
  File "C:\Code\development\Pynance\pynance.py", line 33, in <module>
    print(client.manager.stock_exists('GPX'))
  File "C:\Code\development\Pynance\pynance\decorators.py", line 24, in inner
    func(that, session, *args, **kwargs)
  File "C:\Code\development\Pynance\pynance\database\database.py", line 186, in stock_exists
    .count()
AttributeError: 'NoneType' object has no attribute 'query'
[Finished in 0.7s]

So I am guessing by the traceback, that I am messing up the order of the arguments, but I can't figure out how to order them properly. I have functions that I want to decorate that can take 0-3 arguments in addition to the session. Can someone please point out the error in my methodology?

Answer 1

Change

def inner(that, session=None, *args, **kwargs):

to

def inner(that, *args, session=None, **kwargs):

and

return func(that, session, *args, **kwargs)

to

return func(that, *args, session=session, **kwargs)

It works:

def session_manager(func):

    def inner(that, *args, session=None, **kwargs):
        if not session:
            session = object()
        return func(that, *args, session=session, **kwargs)

    return inner


class A():

    @session_manager
    def _stocks(self, session=None):
        print(session)
        return True

    @session_manager
    def stock_exists(self, ticker, session=None):
        print(ticker, session)
        return True

a = A()
a._stocks()
a.stock_exists('ticker')

Output:

$ python3 test.py
<object object at 0x7f4197810070>
ticker <object object at 0x7f4197810070>

When you use def inner(that, session=None, *args, **kwargs) any second positional argument (counting self) is treated as session argument. So when you call manager.stock_exists('AAPL') session gets value AAPL.

Answer 2

First thing I noticed was that You are calling decorated function twice

@wraps(func)
    def inner(that, session=None, *args, **kwargs):
        newsession = False
        if not session:
            newsession = True
            session = that.db.Session()
        #calling first time
        func(that, session, *args, **kwargs)
        if newsession:
            session.close()
        #calling second time
        return func(that, session, *args, **kwargs)
    return inner

During second call session would be already closed. Also, You don't need to explicitly accept that and session parameters in decorator function, they are already in args and kwargs. Take a look at this solution:

@wraps(func)
def inner(*args, **kwargs):
    session = None
    if not 'session' in kwargs:
        session = that.db.Session()
        kwargs['session'] = session
    result = func(*args, **kwargs)
    if session:
        session.close()
    return result
return inner

You may also want to put session closing code in finally block, then You will be sure that it is closed even if decorated function throws an exception