05-advanced-python/05.11-context-managers-and-the-with-statement.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# with 语句和上下文管理器"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```python\n",
    "# create/aquire some resource\n",
    "...\n",
    "try:\n",
    "    # do something with the resource\n",
    "    ...\n",
    "finally:\n",
    "    # destroy/release the resource\n",
    "    ...\n",
    "```\n",
    "\n",
    "处理文件，线程，数据库，网络编程等等资源的时候，我们经常需要使用上面这样的代码形式，以确保资源的正常使用和释放。\n",
    "\n",
    "好在`Python` 提供了 `with` 语句帮我们自动进行这样的处理，例如之前在打开文件时我们使用： "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "with open('my_file', 'w') as fp:\n",
    "    # do stuff with fp\n",
    "    data = fp.write(\"Hello world\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "这等效于下面的代码，但是要更简便："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "fp = open('my_file', 'w')\n",
    "try:\n",
    "    # do stuff with f\n",
    "    data = fp.write(\"Hello world\")\n",
    "finally:\n",
    "    fp.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 上下文管理器"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "其基本用法如下：\n",
    "```\n",
    "with <expression>:\n",
    "    <block>\n",
    "```\n",
    "\n",
    "`<expression>` 执行的结果应当返回一个实现了上下文管理器的对象，即实现这样两个方法，`__enter__` 和 `__exit__`："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<built-in method __enter__ of file object at 0x0000000003C1D540>\n",
      "<built-in method __exit__ of file object at 0x0000000003C1D540>\n"
     ]
    }
   ],
   "source": [
    "print fp.__enter__\n",
    "print fp.__exit__"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`__enter__` 方法在 `<block>` 执行前执行，而 `__exit__` 在 `<block>` 执行结束后执行：\n",
    "\n",
    "比如可以这样定义一个简单的上下文管理器："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "class ContextManager(object):\n",
    "    \n",
    "    def __enter__(self):\n",
    "        print \"Entering\"\n",
    "    \n",
    "    def __exit__(self, exc_type, exc_value, traceback):\n",
    "        print \"Exiting\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "使用 `with` 语句执行："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Entering\n",
      "  Inside the with statement\n",
      "Exiting\n"
     ]
    }
   ],
   "source": [
    "with ContextManager():\n",
    "    print \"  Inside the with statement\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "即使 `<block>` 中执行的内容出错，`__exit__` 也会被执行："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Entering\n",
      "Exiting\n"
     ]
    },
    {
     "ename": "ZeroDivisionError",
     "evalue": "integer division or modulo by zero",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mZeroDivisionError\u001b[0m                         Traceback (most recent call last)",
      "\u001b[1;32m<ipython-input-6-b509c97cb388>\u001b[0m in \u001b[0;36m<module>\u001b[1;34m()\u001b[0m\n\u001b[0;32m      1\u001b[0m \u001b[1;32mwith\u001b[0m \u001b[0mContextManager\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m----> 2\u001b[1;33m     \u001b[1;32mprint\u001b[0m \u001b[1;36m1\u001b[0m\u001b[1;33m/\u001b[0m\u001b[1;36m0\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[1;31mZeroDivisionError\u001b[0m: integer division or modulo by zero"
     ]
    }
   ],
   "source": [
    "with ContextManager():\n",
    "    print 1/0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## `__`enter`__` 的返回值"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "如果在 `__enter__` 方法下添加了返回值，那么我们可以使用 `as` 把这个返回值传给某个参数："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "class ContextManager(object):\n",
    "    \n",
    "    def __enter__(self):\n",
    "        print \"Entering\"\n",
    "        return \"my value\"\n",
    "    \n",
    "    def __exit__(self, exc_type, exc_value, traceback):\n",
    "        print \"Exiting\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "将 `__enter__` 返回的值传给 `value` 变量："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Entering\n",
      "my value\n",
      "Exiting\n"
     ]
    }
   ],
   "source": [
    "with ContextManager() as value:\n",
    "    print value"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "一个通常的做法是将 `__enter__` 的返回值设为这个上下文管理器对象本身，文件对象就是这样做的："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<open file 'my_file', mode 'r' at 0x0000000003B63030>\n"
     ]
    }
   ],
   "source": [
    "fp = open('my_file', 'r')\n",
    "print fp.__enter__()\n",
    "fp.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import os\n",
    "os.remove('my_file')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "实现方法非常简单："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "class ContextManager(object):\n",
    "    \n",
    "    def __enter__(self):\n",
    "        print \"Entering\"\n",
    "        return self\n",
    "    \n",
    "    def __exit__(self, exc_type, exc_value, traceback):\n",
    "        print \"Exiting\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Entering\n",
      "<__main__.ContextManager object at 0x0000000003D48828>\n",
      "Exiting\n"
     ]
    }
   ],
   "source": [
    "with ContextManager() as value:\n",
    "    print value"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 错误处理"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "上下文管理器对象将错误处理交给 `__exit__` 进行，可以将错误类型，错误值和 `traceback` 等内容作为参数传递给 `__exit__` 函数："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "class ContextManager(object):\n",
    "    \n",
    "    def __enter__(self):\n",
    "        print \"Entering\"\n",
    "    \n",
    "    def __exit__(self, exc_type, exc_value, traceback):\n",
    "        print \"Exiting\"\n",
    "        if exc_type is not None:\n",
    "            print \"  Exception:\", exc_value"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "如果没有错误，这些值都将是 `None`, 当有错误发生的时候："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Entering\n",
      "Exiting\n",
      "  Exception: integer division or modulo by zero\n"
     ]
    },
    {
     "ename": "ZeroDivisionError",
     "evalue": "integer division or modulo by zero",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mZeroDivisionError\u001b[0m                         Traceback (most recent call last)",
      "\u001b[1;32m<ipython-input-14-b509c97cb388>\u001b[0m in \u001b[0;36m<module>\u001b[1;34m()\u001b[0m\n\u001b[0;32m      1\u001b[0m \u001b[1;32mwith\u001b[0m \u001b[0mContextManager\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m----> 2\u001b[1;33m     \u001b[1;32mprint\u001b[0m \u001b[1;36m1\u001b[0m\u001b[1;33m/\u001b[0m\u001b[1;36m0\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[1;31mZeroDivisionError\u001b[0m: integer division or modulo by zero"
     ]
    }
   ],
   "source": [
    "with ContextManager():\n",
    "    print 1/0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "在这个例子中，我们只是简单的显示了错误的值，并没有对错误进行处理，所以错误被向上抛出了，如果不想让错误抛出，只需要将 `__exit__` 的返回值设为 `True`： "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "class ContextManager(object):\n",
    "    \n",
    "    def __enter__(self):\n",
    "        print \"Entering\"\n",
    "    \n",
    "    def __exit__(self, exc_type, exc_value, traceback):\n",
    "        print \"Exiting\"\n",
    "        if exc_type is not None:\n",
    "            print \" Exception suppresed:\", exc_value\n",
    "            return True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Entering\n",
      "Exiting\n",
      " Exception suppresed: integer division or modulo by zero\n"
     ]
    }
   ],
   "source": [
    "with ContextManager():\n",
    "    print 1/0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "在这种情况下，错误就不会被向上抛出。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 数据库的例子"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "对于数据库的 transaction 来说，如果没有错误，我们就将其 `commit` 进行保存，如果有错误，那么我们将其回滚到上一次成功的状态。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "class Transaction(object):\n",
    "    \n",
    "    def __init__(self, connection):\n",
    "        self.connection = connection\n",
    "    \n",
    "    def __enter__(self):\n",
    "        return self.connection.cursor()\n",
    "    \n",
    "    def __exit__(self, exc_type, exc_value, traceback):\n",
    "        if exc_value is None:\n",
    "            # transaction was OK, so commit\n",
    "            self.connection.commit()\n",
    "        else:\n",
    "            # transaction had a problem, so rollback\n",
    "            self.connection.rollback()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "建立一个数据库，保存一个地址表："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import sqlite3 as db\n",
    "connection = db.connect(\":memory:\")\n",
    "\n",
    "with Transaction(connection) as cursor:\n",
    "    cursor.execute(\"\"\"CREATE TABLE IF NOT EXISTS addresses (\n",
    "        address_id INTEGER PRIMARY KEY,\n",
    "        street_address TEXT,\n",
    "        city TEXT,\n",
    "        state TEXT,\n",
    "        country TEXT,\n",
    "        postal_code TEXT\n",
    "    )\"\"\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "插入数据："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "with Transaction(connection) as cursor:\n",
    "    cursor.executemany(\"\"\"INSERT OR REPLACE INTO addresses VALUES (?, ?, ?, ?, ?, ?)\"\"\", [\n",
    "        (0, '515 Congress Ave', 'Austin', 'Texas', 'USA', '78701'),\n",
    "        (1, '245 Park Avenue', 'New York', 'New York', 'USA', '10167'),\n",
    "        (2, '21 J.J. Thompson Ave.', 'Cambridge', None, 'UK', 'CB3 0FA'),\n",
    "        (3, 'Supreme Business Park', 'Hiranandani Gardens, Powai, Mumbai', 'Maharashtra', 'India', '400076'),\n",
    "    ])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "假设插入数据之后出现了问题："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "ename": "Exception",
     "evalue": "out of addresses",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mException\u001b[0m                                 Traceback (most recent call last)",
      "\u001b[1;32m<ipython-input-20-ed8abdd56558>\u001b[0m in \u001b[0;36m<module>\u001b[1;34m()\u001b[0m\n\u001b[0;32m      3\u001b[0m         \u001b[1;33m(\u001b[0m\u001b[1;36m4\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;34m'2100 Pennsylvania Ave'\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;34m'Washington'\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;34m'DC'\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;34m'USA'\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;34m'78701'\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m,\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m      4\u001b[0m     )\n\u001b[1;32m----> 5\u001b[1;33m     \u001b[1;32mraise\u001b[0m \u001b[0mException\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;34m\"out of addresses\"\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[1;31mException\u001b[0m: out of addresses"
     ]
    }
   ],
   "source": [
    "with Transaction(connection) as cursor:\n",
    "    cursor.execute(\"\"\"INSERT OR REPLACE INTO addresses VALUES (?, ?, ?, ?, ?, ?)\"\"\",\n",
    "        (4, '2100 Pennsylvania Ave', 'Washington', 'DC', 'USA', '78701'),\n",
    "    )\n",
    "    raise Exception(\"out of addresses\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "那么最新的一次插入将不会被保存，而是返回上一次 `commit` 成功的状态："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(0, u'515 Congress Ave', u'Austin', u'Texas', u'USA', u'78701')\n",
      "(1, u'245 Park Avenue', u'New York', u'New York', u'USA', u'10167')\n",
      "(2, u'21 J.J. Thompson Ave.', u'Cambridge', None, u'UK', u'CB3 0FA')\n",
      "(3, u'Supreme Business Park', u'Hiranandani Gardens, Powai, Mumbai', u'Maharashtra', u'India', u'400076')\n"
     ]
    }
   ],
   "source": [
    "cursor.execute(\"SELECT * FROM addresses\")\n",
    "for row in cursor:\n",
    "    print row"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## contextlib 模块"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "很多的上下文管理器有很多相似的地方，为了防止写入很多重复的模式，可以使用 `contextlib` 模块来进行处理。\n",
    "\n",
    "最简单的处理方式是使用 `closing` 函数确保对象的 `close()` 方法始终被调用："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<!DOCTYPE html><!--STATUS OK--><html><head><meta http-equiv=\"content-type\" content=\"text/html;charse\n"
     ]
    }
   ],
   "source": [
    "from contextlib import closing\n",
    "import urllib\n",
    "\n",
    "with closing(urllib.urlopen('http://www.baidu.com')) as url:\n",
    "    html = url.read()\n",
    "\n",
    "print html[:100]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "另一个有用的方法是使用修饰符 `@contextlib`："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Enter\n",
      "  Inside the with statement\n",
      "Exit\n"
     ]
    }
   ],
   "source": [
    "from contextlib import contextmanager\n",
    "\n",
    "@contextmanager\n",
    "def my_contextmanager():\n",
    "    print \"Enter\"\n",
    "    yield\n",
    "    print \"Exit\"\n",
    "\n",
    "with my_contextmanager():\n",
    "    print \"  Inside the with statement\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`yield` 之前的部分可以看成是 `__enter__` 的部分，`yield` 的值可以看成是 `__enter__` 返回的值，`yield` 之后的部分可以看成是 `__exit__` 的部分。\n",
    "\n",
    "使用 `yield` 的值："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Enter\n",
      "my value\n",
      "Exit\n"
     ]
    }
   ],
   "source": [
    "@contextmanager\n",
    "def my_contextmanager():\n",
    "    print \"Enter\"\n",
    "    yield \"my value\"\n",
    "    print \"Exit\"\n",
    "    \n",
    "with my_contextmanager() as value:\n",
    "    print value"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "错误处理可以用 `try` 块来完成："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "@contextmanager\n",
    "def my_contextmanager():\n",
    "    print \"Enter\"\n",
    "    try:\n",
    "        yield\n",
    "    except Exception as exc:\n",
    "        print \"   Error:\", exc\n",
    "    finally:\n",
    "        print \"Exit\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Enter\n",
      "   Error: integer division or modulo by zero\n",
      "Exit\n"
     ]
    }
   ],
   "source": [
    "with my_contextmanager():\n",
    "    print 1/0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "对于之前的数据库 `transaction` 我们可以这样定义："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "@contextmanager\n",
    "def transaction(connection):\n",
    "    cursor = connection.cursor()\n",
    "    try:\n",
    "        yield cursor\n",
    "    except:\n",
    "        connection.rollback()\n",
    "        raise\n",
    "    else:\n",
    "        connection.commit()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}