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.
- The code starts by printing "Before 'with'".
- 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.
- The return value from the context manager's __enter__ method is bound to the variable tx.
- The block is executed and the query method is called.
- The query method prints "SQL: SELECT * FROM users".
- 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.
- 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.
