'with' statement

Python have a great number of cool language features and one of them is the with statement. The with statement can be used in numerous ways but I think that the most common usage is to track and release some kind of resource and at the same time encapsulate the code block with a try/finally clause.

For example:

with open('my_file.txt') as f:
    for line in f:
        print line

which could be translated to something like this:

f = open('my_file.txt')
try:
    for line in f:
        print line
finally:
    f.close()

and in a general form:

with <EXPRESSION> [as <VAR>]:
    <BLOCK>

The with statement guarantees that the file will be closed no matter what happens, for example an exception is thrown or a return is executed inside the for statement etc. The result from the expression that is evaluated after the with keyword will be bound to the variable specified after the as keyword. The as part is optional and I'll show you a use case later.

But wait a minute, how does the with keyword know how to close the file? Well, the magic is called 'Context Manager'.

A context manager is a class which defines two methods, __enter__() and __exit__(). The expression in the with statement must evaluate to a context manager else you'll get an AttributeError. The __enter__ method is called when the with statement is evaluated and should return the value which should be assigned to the variable specified after the as keyword. After the block has finished executing the __exit__ method is called. The __exit__ method will be passed any exception occurred while executing the block and can decide if the exception should be muted by returning True or be propagated to the caller by returning False.

Let's look at the following example (not a good design, just an example):

class TransactionContextManager(object):
    def __init__(self, obj):
        self.obj = obj

    def __enter__(self):
        self.obj.begin()
        return self.obj

    def __exit__(self, exc_type, exc_val, exc_tb):
        if exc_type:
            self.obj.rollback()
        else:
            self.obj.commit()
        return False

class Transaction(object):
    def begin(self):
        print "Start transaction"

    def query(self, sql):
        print "SQL:", sql

    def rollback(self):
        print "Rollback transaction"

    def commit(self):
        print "Commit transaction"

print "Before 'with'"

with TransactionContextManager(Transaction()) as tx:
    tx.query("SELECT * FROM users")

print "After 'with'"

The code produces the following output:

Before 'with'
Start transaction
SQL: SELECT * FROM users
Commit transaction
After 'with'

It's easy to trace the execution but I'll explain it anyway.

1. The code starts by printing "Before 'with'".
2. The with statement is executed and instances of TransactionContextManager and Transaction are created. The context manager's __enter__ method is invoked which will invoke the Transaction instance's begin method. The begin method prints "Start transaction" and finally the __enter__ method returns the Transaction instance.
3. The return value from the context manager's __enter__ method is bound to the variable tx.
4. The block is executed and the query method is called.
5. The query method prints "SQL: SELECT * FROM users".
6. The block is done executing and the context manager's __exit__ method is invoked. Since no exception occurred the commit method is invoked on the Transaction instance. The __exit__ method returns False which means that any exceptions should be propagated to the caller, none in this case.
7. The code ends executing by printing "After 'with'".
   

Let's change the Transaction class so that it throws an exception from the query method. This should execute the rollback method instead of the commit method. The exception should also be propagated to the caller since we return False from the __exit__ method.

def query(self, sql):
    raise Exception("SQL Error occurred")

And the output looks like following:

Before 'with'
Start transaction
Rollback transaction
Traceback (most recent call last):
  File "transaction.py", line 32, in 
    tx.query("SELECT * FROM users")
  File "transaction.py", line 21, in query
    raise Exception("SQL Error occurred")
Exception: SQL Error occurred

Here we see that the rollback method is called instead of commit. The rollback method is called because the code inside the with block raised an exception and a valid exc_type argument was passed to the __exit__ method. We can also see that the exception is propagated to the caller because we return False from the __exit__ method.

If we change the __exit__ method to return True instead of False we get the following output:

Before 'with'
Start transaction
Rollback transaction
After 'with'

As expected, we don't see the exception anymore. Note that the __exit__ method is not called if an exception occurs before the with statement starts executing the block.

Alright, I hope you agree with me that the with statement is very useful. Now I'll show you a final example where we don't need to bind the return value from the context manager to a variable. I've re-factored the previous example and hopefully made a better design.

class Connection(object):
    def __init__(self):
        # Default behaviour is to do auto commits
        self._auto_commit = True

    def connect(self):
        pass

    def auto_commit(self):
        return self._auto_commit

    def set_auto_commit(self, mode):
        print "Auto commit:", mode
        self._auto_commit = mode

    def rollback(self):
        print "Rollback transaction"

    def commit(self):
        print "Commit transaction"

    def executeQuery(self, sql):
        # Handle auto commit here if it's enabled
        print "SQL:", sql

class Transaction(object):
    def __init__(self, conn):
        self.conn = conn
        self.auto_commit = False

    def __enter__(self):
        self.auto_commit = self.conn.auto_commit()
        self.conn.set_auto_commit(False)
        return self.conn

    def __exit__(self, exc_type, exc_val, exc_tb):
        if exc_type:
            self.conn.rollback()
        else:
            self.conn.commit()

        self.conn.set_auto_commit(self.auto_commit)
        return False

conn = Connection()
conn.connect()

with Transaction(conn):
    conn.executeQuery("SELECT a user FROM users")
    conn.executeQuery("INSERT some data INTO users")

The output:

Auto commit: False
SQL: SELECT a user FROM users
SQL: INSERT some data INTO users
Commit transaction
Auto commit: True

As the example shows, we don't need to use the return value from the context manager to still have use of them. We only want to be sure that the queries inside the with block is executed in the same transaction and rollbacked if an exception occurs or commited if successful. Without the with statement the code would be implemented something like the following snippet:

conn = Connection()
conn.connect()
auto_comm = conn.auto_commit()

try:
    conn.set_auto_commit(False)
    conn.executeQuery("SELECT a user FROM users")
    conn.executeQuery("INSERT some data INTO users")
    conn.commit()
except:
    conn.rollback()
    raise
finally:
    conn.set_auto_commit(auto_comm)

Well, for me the with statement seems more clean and convenient to use than the last example.