Python: use the decorator to pass an extra parameter to the function

Question

I have a function than depends on db connection. This function has a lot of return statements, of this kind:

def do_something_with_data(db: Session, data: DataClass):
    db = Session()
    
    if condition1:
         db.do_something1()
         db.close()
         return
    
    if condition2:
         db.do_something2()
         db.close()
         return
    
    if condition3:
         db.do_something3()
         db.close()
         return
    ...

After executing the function, I need to run db.close(), but because of the structure of the function this entry will have to be duplicated many times for each return as shown above.

So I made a decorator that passes the created session to the function and closes the session at the end of the execution of the function instead.


def db_depends(function_that_depends_on_db):
    def inner(*args, **kwargs):
        db = Session()
        result = function_that_depends_on_db(db, *args, **kwargs)
        db.close()
        return result
    return inner

@db_depends
def do_something_with_data(db: Session, data: DataClass):    
    if condition1:
         db.do_something1()
         return
    
    if condition2:
         db.do_something2()
         return
    
    if condition3:
         db.do_something3()
         return
    ...

All works great, but the fact, that user see two required arguments in definition, however there is only one (data) seems kinda dirty.

Is it possible to do the same thing without misleading people who will read the code or IDE hints?

If your function is full of such if statements, you can use elif to avoid multiple returns. — Mechanic Pig, Sep 02 '22 at 06:54

score 2 · Accepted Answer · answered Sep 02 '22 at 06:51

Just have the function accept the Session parameter normally:

def do_something_with_data(db: Session, data: DataClass):    
    if condition1:
         db.do_something1()
         return
    
    if condition2:
         db.do_something2()
         return
    
    if condition3:
         db.do_something3()
         return

This allows the user to specify a Session explicitly, for example to reuse the same Session to do multiple things.

Yes, that doesn't close the Session. Because that is the responsibility of the calling code, since that's where the Session came from in the first place. After all, if the calling code wants to reuse a Session, then it shouldn't be closed.

If you want a convenience method to open a new, temporary Session for the call, you can easily do that using the existing decorator code:

do_something_in_new_session = db_depends(do_something_with_data)

But if we don't need to apply this logic to multiple functions, then "simple is better than complex" - just write an ordinary wrapper:

def do_something_in_new_session(data: DataClass):
    db = Session()
    result = do_something_with_data(db, data)
    db.close()
    return result

Either way, it would be better to write the closing logic using a with block, assuming your library supports it:

def do_something_in_new_session(data: DataClass):
    with Session() as db:
        return do_something_with_data(db, data)

Among other things, this ensures that .close is called even if an exception is raised in do_something_with_data.

If your DB library doesn't support that (i.e., the Session class isn't defined as a context manager - although that should only be true for very old libraries now), that's easy to fix using contextlib.closing from the standard library:

from contextlib import closing

def do_something_in_new_session(data: DataClass):
    with closing(Session()) as db:
        return do_something_with_data(db, data)

(And of course, if you don't feel the need to make a wrapper like that, you can easily use such a with block directly at the call site.)

score 0 · Answer 2 · answered May 16 '23 at 11:54

@KarlKnechtel answer did a really good job of describing what I need to do specifically in my case, however, if you really need to change the signature (type) of a function if it's used with a decorator that passes one of the argument for you, well, it turns out the documentation for the typing module has the answer (from doc):

typing.Concatenate Used with Callable and ParamSpec to type annotate a higher order callable which adds, removes, or transforms parameters of another callable.

Usage (also from doc) essentially one to one match with my case:

from collections.abc import Callable
from threading import Lock
from typing import Concatenate, ParamSpec, TypeVar

P = ParamSpec('P')
R = TypeVar('R')

# Use this lock to ensure that only one thread is executing a function
# at any time.
my_lock = Lock()

def with_lock(f: Callable[Concatenate[Lock, P], R]) -> Callable[P, R]:
    '''A type-safe decorator which provides a lock.'''
    def inner(*args: P.args, **kwargs: P.kwargs) -> R:
        # Provide the lock as the first argument.
        return f(my_lock, *args, **kwargs)
    return inner

@with_lock
def sum_threadsafe(lock: Lock, numbers: list[float]) -> float:
    '''Add a list of numbers together in a thread-safe manner.'''
    with lock:
        return sum(numbers)

# We don't need to pass in the lock ourselves thanks to the decorator.
sum_threadsafe([1.1, 2.2, 3.3])

Python: use the decorator to pass an extra parameter to the function

2 Answers2