---
tags: Python
---
# Context Manager

[TOC]

以開檔來說，以往以 `try-except` 的寫法，在使用檔案過程中發生例外時，在 `finally` 處確保正確關檔。

這種寫法雖然不會有問題，但是缺點就是必須手動加入關閉檔案的程式碼，不是很方便，也很容易忘記，Python2.5 後提供 `with`，可以**自動地開始和結束一些事情**。

> This PEP adds a new statement "with" to the Python language to make it possible to factor out standard uses of try/finally statements.
> In this PEP, context managers provide `__enter__()` and `__exit__()` methods that are invoked on entry to and exit from the body of the with statement.
>
> [name=[PEP 343 -- The "with" Statement](https://www.python.org/dev/peps/pep-0343/)]

context manager 用到了 `__enter__` 和 `__exit__` 。用 `with` 會呼叫 `__enter__`，並將回傳值指派給 `as` 的變數。

## 純 Class

`__exit__` 像是 `finally`，就算 `with` 的 scope 內有出錯，一樣會執行。

class DBHandler:

    def __enter__(self):
        print('start')

    def __exit__(self, exception_type, exception_value, traceback):
        print('end')

with DBHandler():
    print('do something')

class DBHandler:

    def __init__(self, name):
        self.name = name

    def __enter__(self):
        print(f'start with {self.name}')
        return f'db_{self.name}'

    def __exit__(self, exception_type, exception_value, traceback):
        print('end')

with DBHandler('PET') as db:
    print(f'do something with {db}')

`__exit__` 內可以在判斷 `with` 是否為正常結束：

def __exit__(self, exception_type, exception_value, traceback):
    if exception_type is None:
        print('success')
    else:
        print('error')

## [`contextlib.ContextDecorator`](https://docs.python.org/3/library/contextlib.html#contextlib.ContextDecorator)

可以以 decorator 或 with 的形式給予。以 decorator 的方式的缺點為：不能使用 `as` 來取得回傳物件。

class DBHandler_decorator(contextlib.ContextDecorator):
    def __enter__(self):
        print('start')

    def __exit__(self, exception_type, exception_value, traceback):
        print('end')

@DBHandler_decorator()
def function():
    print('do something')

@DBHandler_decorator()
def function2(name):
    print('do something')
    return f'db_{name}'

with DBHandler_decorator():
    print('do something')

`ContextDecorator` 在 [`__call__`](https://github.com/python/cpython/blob/master/Lib/contextlib.py#L56) 包裝成 decorator。

## [`contextlib.contextmanager`](https://docs.python.org/3/library/contextlib.html#contextlib.contextmanager)

> This function is a decorator that can be used to define a factory function for with statement context managers, without needing to create a class or separate `__enter__()` and `__exit__()` methods.
>
> [name=[`contextlib.contextmanager` / contextlib](https://docs.python.org/3/library/contextlib.html#contextlib.contextmanager)]

import contextlib

@contextlib.contextmanager
def db_handler():
    print('start')
    yield
    print('end')

with db_handler():
    print('do something')

import contextlib

@contextlib.contextmanager
def db_handler(name):
    print(f'start with {name}')
    yield f'db_{name}'
    print('end')

with db_handler('PET') as db:
    print(f'do something with {db}')

注意上方的用法比較像是 `else` 而非 `finally`，當 `with` 的 scope 內有錯時，並不會執行 `yield` 後的程式，可以加上 `try-except` 修改。

import contextlib

@contextlib.contextmanager
def db_handler(name):
    print(f'start with {name}')
    try:
        yield f'db_{name}'
    finally:
        print('end')

with db_handler('PET') as db:
    print(f'do something with {db}')
    raise Exception()

## 應用

* openpyxl

   @contextlib.contextmanager
   def open_excel(filename, *args, **kwargs):
       wb = openpyxl.load_workbook(filename, *args, **kwargs)
       try:
           yield wb
       finally:
           wb.close()

   with open_excel('test.xlsx') as wb:
       print(wb)

* session

   @contextlib.contextmanager
   def session_scope():
       session = Session()
       try:
           yield session
           session.commit()
       except:
           session.rollback()
           raise
       finally:
           session.close()

* timer

   import time

   class Timer:    
       def __enter__(self):
           self.start = time.clock()
           return self

       def __exit__(self, *args):
           self.end = time.clock()
           self.interval = self.end - self.start
           print(self.interval)

   with Timer() as t:
       print('do something')

## 參考資料

* [PEP 343 -- The "with" Statement](https://www.python.org/dev/peps/pep-0343/)
* [`contextlib.contextmanager` / contextlib](https://docs.python.org/3/library/contextlib.html#contextlib.contextmanager)
* [簡潔的Python｜重構你的舊程式](https://www.books.com.tw/products/E050051199)