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?

Asked By: Олексій Холостенко

||

Source

Answer 1

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.)

Answered By: Karl Knechtel

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

Question:

Answers: