Making Objects Support the Context-Management Protocol

Contributed by: David Beazley and Brian K. Jones

Difficulty level: Easy

Visit the Python Cookbook: Classes and Objects Lab to test out this Shortcut in a live practice environment.

Problem

You want to make your objects support the context-management protocol (the with statement).

Solution

In order to make an object compatible with the with statement, you need to implement __enter__() and __exit__() methods. For example, consider the following class, which provides a network connection:

from socket import socket, AF_INET, SOCK_STREAM

class LazyConnection(object):
    def __init__(self, address, family=AF_INET, type=SOCK_STREAM):
        self.address = address
        self.family = family
        self.type = type
        self.sock = None

    def __enter__(self):
        if self.sock is not None:
            raise RuntimeError('Already connected')
        self.sock = socket(self.family, self.type)
        self.sock.connect(self.address)
        return self.sock

    def __exit__(self, exc_ty, exc_val, tb):
        self.sock.close()
        self.sock = None

The key feature of this class is that it represents a network connection, but it doesn’t actually do anything initially (e.g., it doesn’t establish a connection). Instead, the connection is established and closed using the with statement (essentially on demand). For example:

from functools import partial

conn = LazyConnection(('www.python.org', 80))
# Connection closed
with conn as s:
    # conn.__enter__() executes: connection open
    s.send(b'GET /index.html HTTP/1.0\r\n')
    s