[Python] 使用Context manager管理資源

Class使用context manager

__enter__()和__exit__()會分別於class配合with被調用時先後執行

順序:__init()__=>__enter__()=>with內的語句=>__exit__()

import pymysql

class get_data_from_mysql:
    
    def __init__(self, dbName):
        self.dbName = dbName
    
    def __enter__(self):
        try:
            self.mysqldb = pymysql.connect(
                host="127.0.0.1",
                user="root",
                passwd="123456",
                database=self.dbName)
            print 'Connect db successful!'
        except Exception as e:
            raise Exception('Fail to connect db! {}'.format(str(e)))
        
    def __exit__(self, type, value, trace):
        self.mysqldb.close()
        print 'db close.'

if __name__ == '__main__':
    with get_data_from_mysql('school'):
        print 'Execute other code...'

__exit__()內type, value, trace參數可用作處理exception

def __exit__(self, type, value, trace):
    self.mysqldb.close()
    print 'db close.'
    print 'Type: {}'.format(type)
    print 'Value: {}'.format(value)
    print 'Trace: {}'.format(trace)

分別可顯示Exception的類型,值和追蹤

Function使用context manager

以decorator的方式對function產生context manager

from contextlib import contextmanager
import pymysql

@contextmanager
def run(dbName):
    try:
        mysqldb = pymysql.connect(
            host="127.0.0.1",
            user="root",
            passwd="123456",
            database=dbName)
        print 'Connect db successful!'
    except Exception as e:
        raise Exception('Fail to connect db! {}'.format(str(e)))
    yield
    mysqldb.close()
    print 'db close.'
        
if __name__ == '__main__':
    with run('school'):
        print 'Using decorator...'