Function wrappers inside class using decorators

Question

I have a class which interacts with a database and so there are repetitive actions (establish session, commit, close session) before and after each member method of the class.

As follows:

class UserDatabaseManager(object):

    DEFAULT_DB_PATH = 'test.db'

    def __init__(self, dbpath=DEFAULT_DB_PATH):
        dbpath = 'sqlite:///' + dbpath
        self.engine = create_engine(dbpath, echo=True)

    def add_user(self, username, password):
        Session = sessionmaker(bind=self.engine)
        session = Session()
        # <============================== To be wrapped
        user = User(username, password)
        session.add(user)
        # ==============================>
        session.commit()
        session.close()

    def delete_user(self, user):
        Session = sessionmaker(bind=self.engine)
        session = Session()
        # <============================== To be wrapped
        # Delete user here
        # ==============================>
        session.commit()
        session.close()

What is an idiomatic way to abstract out the repeated session calls with a function wrapper?

I would prefer to do this with decorators by declaring a private _Decorators class inside UserDatabaseManager and implementing the wrapper function inside there, but then such class won't be able to access the self.engine instance attribute of the outer class.

@Luke I actually upvoted both your answers as they work, but it's hard to tell which approach is better at the moment and both are quite similar. I've been meaning to see how either approach scales with my project for a few more days before casting my verdict. But thanks so much for replying to my question! — Katie, May 12 '18 at 00:50

Ajax1234 · Answer 1 · 2018-05-05T23:07:30.817

You can create a simple function outside the class to wrap each method:

def create_session(**kwargs):
   def outer(f):
     def wrapper(cls, *args):
       Session = sessionmaker(bind=getattr(cls, 'engine'))
       session = Session()
       getattr(session, kwargs.get('action', 'add'))(f(cls, *args))
       session.commit()
       session.close()
     return wrapper
   return outer

class UserDatabaseManager(object):
  DEFAULT_DB_PATH = 'test.db'
  def __init__(self, dbpath=DEFAULT_DB_PATH):
    dbpath = 'sqlite:///' + dbpath
    self.engine = create_engine(dbpath, echo=True)
  @create_session(action = 'add')
  def add_user(self, username, password):
    return User(username, password)

  @create_session(action = 'delete')
  def delete_user(self, user):
     return User(username, password)

Generally, setup and tear-down operations like the above are best put in a contextmanager:

class UserDatabaseManager(object):
  DEFAULT_DB_PATH = 'test.db'
  def __init__(self, dbpath=DEFAULT_DB_PATH):
     dbpath = 'sqlite:///' + dbpath
     self.engine = create_engine(dbpath, echo=True)

class UserAction(UserDatabaseManager):
  def __init__(self, path):
    UserDatabaseManager.__init__(self, path)
  def __enter__(self):
    self.session = sessionmaker(bind=self.engine)()
    return self.session
  def __exit__(self, *args):
     self.session.commit()
     self.session.close()

with UserAction('/the/path') as action:
   action.add(User(username, password))

with UserAction('/the/path') as action:
   action.remove(User(username, password))

Wouldn't it be a bad thing to commit a session where some statements were valid but then an exception was raised? — Rainy, May 05 '18 at 23:39
@Rainy can you elaborate? with `__enter__` and `__exit__`, any exceptions raised in the setup or the body of the `with` statement will crash the current program. — Ajax1234, May 05 '18 at 23:47
Sorry for taking long to respond; my point is that an exception raised in body of `with` will cause `__exit__` to be run and will commit session. (tested on Python3.x). The user most likely doesn't want that, he would either want to commit explicitly at some point or to avoid committing at all if there was an error. — Rainy, May 10 '18 at 17:52

user2856 · Accepted Answer · 2018-05-07T06:50:32.377

A simple (and in my opinion, the most idiomatic) way of doing this is to wrap the setup/teardown boilerplate code in a context manager using contextlib.contextmanager. You then simply use a with statement in the functions that do the work (rather than trying to wrap that function itself).

For example:

from contextlib import contextmanager

class UserDatabaseManager(object):

    DEFAULT_DB_PATH = 'test.db'

    def __init__(self, dbpath=DEFAULT_DB_PATH):
        dbpath = 'sqlite:///' + dbpath
        self.engine = create_engine(dbpath, echo=True)

    @contextmanager
    def session(self):
        try:
            Session = sessionmaker(bind=self.engine)
            session = Session()
            yield session
            session.commit()
        except:
            session.rollback()
        finally:
            session.close()

    def add_user(self, username, password):
        with self.session() as session:
            user = User(username, password)
            session.add(user)

    def delete_user(self, user):
        with self.session() as session:
            session.delete(user)

I accepted this answer as it's cleaner, but I have to give it to Ajax1234's answer for actually accomplishing the same purpose with decorators and with more gory details around the dunder methods, this was a close call. — Katie, May 18 '18 at 16:12

Function wrappers inside class using decorators

2 Answers2