Successor of ContextDecorator doesn't work: 'generator' object has no attribute 'add'

Question

I have a problem with inheriting from the ContextDecorator class. I can't understand why the method session_manager() works:

@contextmanager
def session_manager():
    session = Session()
    yield session
    try:
    session.commit()
except Exception as e:
    session.rollback()
    raise e
finally:
    session.close()

But exactly the same code with ContextDecorator successor class gives an error:

class SessionManager(ContextDecorator):
    def __init__(self):
        self.session = Session()

    def __enter__(self):
        try:
            yield self.session
            self.session.commit()
        except Exception as e:
            self.session.rollback()
            raise e

    def __exit__(self, *exc):
        self.session.close()

---

Exception:

AttributeError: 'generator' object has no attribute 'add'

The documentation and tutorials do not have complex examples (only with 'print' statements) and they works great: https://docs.python.org/3/library/contextlib.html

I don't understand why method session_manager() works, although it returns a generator:

yield session

Here I write some small and simple code: https://gist.github.com/tranebaer/46f94263030dd8f7c1bfcf72d0e37610

Answer 1

The __enter__ method is not supposed to be a generator, unless you want to treat the return value as such in the runtime context. It is called when entering the block governed by the with-statement and its return value is bound to the target(s) specified in the as clause, if any. So the attribute error is the result of calling the method add() on the generator inside the block, when you meant it to be the Session object. Possible cleanup and exception handling should take place in the __exit__ method:

from contextlib import closing, ContextDecorator, ExitStack

class SessionManager(ContextDecorator):

    def __init__(self, session_cls=Session):
        self.session = session_cls()

    def __enter__(self):
        return self.session

    def __exit__(self, type, value, tb):
        with closing(self.session), ExitStack() as stack:
            stack.callback(self.session.rollback)
            if not value:
                self.session.commit()
                # If commit raises an exception, then rollback is left
                # in the exit stack.
                stack.pop_all()

Note that you don't need to inherit from ContextDecorator in order to make a context manager. Just implementing __enter__ and __exit__ is enough. In fact in this case it is a bit pointless, because a function decorated with SessionManager has no access to the Session object.