Context managers can be used for more than just opening and closing files.
If we think about it there are two phases to a context manager:
when the with statement is executing: we enter the context
when the with block is done: we exit the context
We can create our own context manager using a class that implements an __enter__ method which is executed when we enter the context, and an __exit__ method that is executed when we exit the context.
There is a general pattern that context managers can help us deal with:
Open - Close
Lock - Release
Change - Reset
Enter - Exit
Start - Stop
The __enter__ method is quite straightforward. It can (but does not have to) return one or more objects we then use inside the with block.
The __exit__ method however is slightly more complicated.
It needs to return a boolean True/False. This indicates to Python whether to suppress any errors that occurred in the with block. As we saw with files, that was not the case - i.e. it returns a False
If an error does occur in the with block, the error information is passed to the __exit__ method - so it needs three things: the exception type, the exception value and the traceback. If no error occured, then those values will simply be None.
Let's go ahead and create a context manager:
class MyContext:
def __init__(self):
self.obj = None
def __enter__(self):
print('entering context...')
self.obj = 'the Return Object'
return self.obj
def __exit__(self, exc_type, exc_value, exc_traceback):
print('exiting context...')
if exc_type:
print(f'*** Error occurred: {exc_type}, {exc_value}')
return False # do not suppress exceptions
We can even cause an exception inside the with block:
with MyContext() as obj:
raise ValueError
entering context...
exiting context...
*** Error occurred: <class 'ValueError'>,
---------------------------------------------------------------------------
ValueError Traceback (most recent call last)
<ipython-input-13-3f36f003da19> in <module>
1 with MyContext() as obj:
----> 2 raise ValueError
ValueError:
We can change that by returning True from the __exit__ method:
class MyContext:
def __init__(self):
self.obj = None
def __enter__(self):
print('entering context...')
self.obj = 'the Return Object'
return self.obj
def __exit__(self, exc_type, exc_value, exc_traceback):
print('exiting context...')
if exc_type:
print(f'*** Error occurred: {exc_type}, {exc_value}')
return True # suppress exceptions
with MyContext() as obj:
raise ValueError
print('reached here without an exception...')
entering context...
exiting context...
*** Error occurred: <class 'ValueError'>,
reached here without an exception...
Notice that the obj we obtained from the context manager, still exists in our scope after the with statement.The with statement does not have its own local scope - it's not a function! However, the context manager could manipulate the object returned by the context manager:
class Resource:
def __init__(self, name):
self.name = name
self.state = None
class ResourceManager:
def __init__(self, name):
self.name = name
self.resource = None
def __enter__(self):
print('entering context')
self.resource = Resource(self.name)
self.resource.state = 'created'
return self.resource
def __exit__(self, exc_type, exc_value, exc_traceback):
print('exiting context')
self.resource.state = 'destroyed'
if exc_type:
print('error occurred')
return False
with ResourceManager('spam') as res:
print(f'{res.name} = {res.state}')
print(f'{res.name} = {res.state}')
entering context
spam = created
exiting context
spam = destroyed
We still have access to res, but it's internal state was changed by the resource manager's __exit__ method.Although we already have a context manager for files built-in to Python, let's go ahead and write our own anyways - good practice.
class File:
def __init__(self, name, mode):
self.name = name
self.mode = mode
def __enter__(self):
print('opening file...')
self.file = open(self.name, self.mode)
return self.file
def __exit__(self, exc_type, exc_value, exc_traceback):
print('closing file...')
self.file.close()
return False
with File('test.txt', 'w') as f:
f.write('This is a late parrot!')
opening file...
closing file...
Note that the __enter__ method can return anything, including the context manager itself.
If we wanted to, we could re-write our file context manager this way:
class File():
def __init__(self, name, mode):
self.name = name
self.mode = mode
def __enter__(self):
print('opening file...')
self.file = open(self.name, self.mode)
return self
def __exit__(self, exc_type, exc_value, exc_traceback):
print('closing file...')
self.file.close()
return False
#Of course, now we would have to use the context manager object's file property to get a handle to the file:
with File('test.txt', 'r') as file_ctx:
print(next(file_ctx.file))
print(file_ctx.name)
print(file_ctx.mode)
opening file...
This is a late parrot
test.txt
r
closing file...
Working with Iterator
class DataIterator:
def __init__(self, fname):
self._fname = fname
self._f = None
def __iter__(self):
return self
def __next__(self):
row = next(self._f)
return row.strip('\n').split(',')
def __enter__(self):
self._f = open(self._fname)
return self
def __exit__(self, exc_type, exc_value, exc_tb):
if not self._f.closed:
self._f.close()
return False
with DataIterator('nyc_parking_tickets_extract.csv') as data:
for row in data:
print(row)
['Summons Number', 'Plate ID', 'Registration State', 'Plate Type', 'Issue Date', 'Violation Code', 'Vehicle Body Type', 'Vehicle Make', 'Violation Description']
['4006478550', 'VAD7274', 'VA', 'PAS', '10/5/2016', '5', '4D', 'BMW', 'BUS LANE VIOLATION']
['4006462396', '22834JK', 'NY', 'COM', '9/30/2016', '5', 'VAN', 'CHEVR', 'BUS LANE VIOLATION']
['4007117810', '21791MG', 'NY', 'COM', '4/10/2017', '5', 'VAN', 'DODGE', 'BUS LANE VIOLATION']
['4006265037', 'FZX9232', 'NY', 'PAS', '8/23/2016', '5', 'SUBN', 'FORD', 'BUS LANE VIOLATION']
['4006535600', 'N203399C', 'NY', 'OMT', '10/19/2016', '5', 'SUBN', 'FORD', 'BUS LANE VIOLATION']
['4007156700', '92163MG', 'NY', 'COM', '4/13/2017', '5', 'VAN', 'FRUEH', 'BUS LANE VIOLATION']......
Additional Use
common patterns we can implement with context managers:
Of course, this is kind of a pain to have to store the current value, set it to something new, and then remember to set it back to it's original value.
How about writing a context manager to do that seamlessly for us:
In fact, the decimal class already has a context manager, and it's way better than ours, because we can set not only the precision, but anything else we want:
with decimal.localcontext() as ctx:
ctx.prec = 3
print(decimal.Decimal(1) / decimal.Decimal(3))
print(decimal.Decimal(1) / decimal.Decimal(3))
0.333
0.3333333333333333333333333333
Timing a with block
from time import perf_counter, sleep
class Timer:
def __init__(self):
self.elapsed = 0
def __enter__(self):
self.start = perf_counter()
return self
def __exit__(self, exc_type, exc_value, exc_traceback):
self.stop = perf_counter()
self.elapsed = self.stop - self.start
return False
with Timer() as timer:
sleep(1)
print(timer.elapsed)
0.9993739623039163
Decorator
Fortunately the standard library has already though of this - in fact that was one of the critical goals of Python's context managers - the ability to create context managers using generator functions (see PEP 343).
from contextlib import contextmanager
@contextmanager
def open_file(fname, mode='r'):
print('opening file...')
f = open(fname, mode)
try:
yield f
finally:
print('closing file...')
f.close()
with open_file('test.txt') as f:
print(f.readlines())
opening file...
['Sir Spamalot']
closing file...
`Let's implement a timer.
from time import perf_counter, sleep
from contextlib import contextmanager
@contextmanager
def timer():
stats = dict()
start = perf_counter()
stats['start'] = start
try:
yield stats
finally:
end = perf_counter()
stats['end'] = end
stats['elapsed'] = end - start
with timer() as stats:
sleep(1)
print(stats)
