08. Classes and Objects/03. Making Objects Support the Context-Management Protocol.ipynb

{
 "metadata": {
  "name": ""
 },
 "nbformat": 3,
 "nbformat_minor": 0,
 "worksheets": [
  {
   "cells": [
    {
     "cell_type": "heading",
     "level": 1,
     "metadata": {},
     "source": [
      "8.3. Making Objects Support the Context-Management\n",
      "Protocol"
     ]
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "# Problem\n",
      "You want to make your objects support the context-management protocol (the `with`\n",
      "statement).\n",
      "\n",
      "# Solution\n",
      "In order to make an object compatible with the with statement, you need to implement\n",
      "`__enter__()` and `__exit__()` methods. For example, consider the following class,\n",
      "which provides a network connection:"
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "from socket import socket, AF_INET, SOCK_STREAM\n",
      "\n",
      "class LazyConnection:\n",
      "    def __init__(self, address, family=AF_INET, type=SOCK_STREAM):\n",
      "        self.address = address\n",
      "        self.family = AF_INET\n",
      "        self.type = SOCK_STREAM\n",
      "        self.sock = None\n",
      "\n",
      "    def __enter__(self):\n",
      "        if self.sock is not None:\n",
      "            raise RuntimeError('Already connected')\n",
      "        self.sock = socket(self.family, self.type)\n",
      "        self.sock.connect(self.address)\n",
      "        return self.sock\n",
      "\n",
      "    def __exit__(self, exc_ty, exc_val, tb):\n",
      "        self.sock.close()\n",
      "        self.sock = None"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [],
     "prompt_number": 1
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "The key feature of this class is that it represents a network connection, but it doesn\u2019t\n",
      "actually do anything initially (e.g., it doesn\u2019t establish a connection). Instead, the \n",
      "connection is established and closed using the `with` statement (essentially on demand). For\n",
      "example:"
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "from functools import partial\n",
      "\n",
      "conn = LazyConnection(('www.python.org', 80))\n",
      "# Connection closed\n",
      "with conn as s:\n",
      "    # conn.__enter__() executes: connection open\n",
      "    s.send(b'GET /index.html HTTP/1.0\\r\\n')\n",
      "    s.send(b'Host: www.python.org\\r\\n')\n",
      "    s.send(b'\\r\\n')\n",
      "    resp = b''.join(iter(partial(s.recv, 8192), b''))\n",
      "    # conn.__exit__() executes: connection closed"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [],
     "prompt_number": 2
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "# Discussion\n",
      "The main principle behind writing a context manager is that you\u2019re writing code that\u2019s\n",
      "meant to surround a block of statements as defined by the use of the with statement.\n",
      "When the `with` statement is first encountered, the `__enter__()` method is triggered.\n",
      "The return value of `__enter__()` (if any) is placed into the variable indicated with the\n",
      "`as` qualifier. Afterward, the statements in the body of the `with` statement execute. Finally,\n",
      "the `__exit__()` method is triggered to clean up.\n",
      "\n",
      "This control flow happens regardless of what happens in the body of the `with` statement,\n",
      "including if there are exceptions. In fact, the three arguments to the `__exit__()` method\n",
      "contain the exception type, value, and traceback for pending exceptions (if any). The\n",
      "`__exit__()` method can choose to use the exception information in some way or to\n",
      "ignore it by doing nothing and returning `None` as a result. If `__exit__()` returns `True`,\n",
      "the exception is cleared as if nothing happened and the program continues executing\n",
      "statements immediately after the `with` block.\n",
      "\n",
      "One subtle aspect of this recipe is whether or not the `LazyConnection` class allows nested\n",
      "use of the connection with multiple `with` statements. As shown, only a single socket\n",
      "connection at a time is allowed, and an exception is raised if a repeated `with` statement\n",
      "is attempted when a socket is already in use. You can work around this limitation with\n",
      "a slightly different implementation, as shown here:"
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "from socket import socket, AF_INET, SOCK_STREAM\n",
      "\n",
      "class LazyConnection:\n",
      "    def __init__(self, address, family=AF_INET, type=SOCK_STREAM):\n",
      "        self.address = address\n",
      "        self.family = AF_INET\n",
      "        self.type = SOCK_STREAM\n",
      "        self.connections = []\n",
      "\n",
      "    def __enter__(self):\n",
      "        sock = socket(self.family, self.type)\n",
      "        sock.connect(self.address)\n",
      "        self.connections.append(sock)\n",
      "        return sock\n",
      "    \n",
      "    def __exit__(self, exc_ty, exc_val, tb):\n",
      "        self.connections.pop().close()"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [],
     "prompt_number": 4
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "```python\n",
      "# Example use\n",
      "from functools import partial\n",
      "\n",
      "conn = LazyConnection(('www.python.org', 80))\n",
      "with conn as s1:\n",
      "    ...\n",
      "    with conn as s2:\n",
      "        ...\n",
      "        # s1 and s2 are independent sockets\n",
      "```"
     ]
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "In this second version, the `LazyConnection` class serves as a kind of factory for \n",
      "connections. Internally, a list is used to keep a stack. Whenever `__enter__()` executes, it\n",
      "makes a new connection and adds it to the stack. The `__exit__()` method simply pops\n",
      "the last connection off the stack and closes it. It\u2019s subtle, but this allows multiple \n",
      "connections to be created at once `with` nested with statements, as shown.\n",
      "\n",
      "Context managers are most commonly used in programs that need to manage resources\n",
      "such as files, network connections, and locks. A key part of such resources is they have\n",
      "to be explicitly closed or released to operate correctly. For instance, if you acquire a lock,\n",
      "then you have to make sure you release it, or else you risk deadlock. By implementing\n",
      "`__enter__()`, `__exit__()`, and using the `with` statement, it is much easier to avoid such\n",
      "problems, since the cleanup code in the `__exit__()` method is guaranteed to run no\n",
      "matter what.\n",
      "\n",
      "An alternative formulation of context managers is found in the `contextmanager` \n",
      "module. See Recipe 9.22. A thread-safe version of this recipe can be found in Recipe 12.6."
     ]
    }
   ],
   "metadata": {}
  }
 ]
}