Context Manager

Conception

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:

  1. when the with statement is executing: we enter the context

  2. 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.

  1. 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

  2. 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:

  • Open - Close

  • Change - Reset

  • Start - Stop

Decimal Contexts

import decimal
decimal.getcontext()

Context(prec=28, rounding=ROUND_HALF_EVEN, Emin=-999999, Emax=999999, capitals=1, clamp=0, flags=[], traps=[InvalidOperation, DivisionByZero, Overflow])

If we create a decimal object, then it will use those settings.We can certainly change the properties of that global context:

decimal.getcontext().prec=14
decimal.getcontext()

Context(prec=14, rounding=ROUND_HALF_EVEN, Emin=-999999, Emax=999999, capitals=1, clamp=0, flags=[], traps=[InvalidOperation, DivisionByZero, Overflow])

Suppose now that we just want to temporarily change something in the context - we would have to do something like this:

old_prec = decimal.getcontext().prec
decimal.getcontext().prec = 4
print(decimal.Decimal(1) / decimal.Decimal(3))
decimal.getcontext().prec = old_prec
print(decimal.Decimal(1) / decimal.Decimal(3))

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:

class precision:
    def __init__(self, prec):
        self.prec = prec
        self.current_prec = decimal.getcontext().prec
        
    def __enter__(self):
        decimal.getcontext().prec = self.prec
        
    def __exit__(self, exc_type, exc_value, exc_traceback):
        decimal.getcontext().prec = self.current_prec
        return False   
        
        
with precision(3):
    print(decimal.Decimal(1) / decimal.Decimal(3))
print(decimal.Decimal(1) / decimal.Decimal(3))  

 
0.333
0.3333333333333333333333333333
  

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)

Last updated