2019-06-20 contextlib with语句与上下文管理器

在Python中,读写文件这样的资源要特别注意,必须在使用完毕后正确关闭它们。正确关闭文件资源的一个方法是使用try...finally:

try:
    f = open('/path/to/file', 'r')
    f.read()
finally:
    if f:
        f.close()

写try...finally非常繁琐。Python的with语句允许我们非常方便地使用资源,而不必担心资源没有关闭,所以上面的代码可以简化为:

with open('/path/to/file', 'r') as f:
    f.read()

并不是只有open()函数返回的fp对象才能使用with语句。实际上,任何对象,只要正确实现了上下文管理,就可以用于with语句。

实现上下文管理是通过enterexit这两个方法实现的。例如,下面的class实现了这两个方法:

class Query(object):

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

    def __enter__(self):
        print('Begin')
        return self
    
    def __exit__(self, exc_type, exc_value, traceback):
        if exc_type:
            print('Error')
        else:
            print('End')
    
    def query(self):
        print('Query info about %s...' % self.name)

这样我们就可以把自己写的资源对象用于with语句:

with Query('Bob') as q:
    q.query()

@contextmanager
编写__enter____exit__仍然很繁琐,因此Python的标准库contextlib提供了更简单的写法,上面的代码可以改写如下:

from contextlib import contextmanager
class Query(object):
    def __init__(self, name):
        self.name = name

    def query(self):
        print('Query info about %s...' % self.name)

@contextmanager
def create_query(name):
    print('Begin')
    q = Query(name)
    yield q
    print('End')

@contextmanager这个decorator接受一个generator,用yield语句把with ... as var把变量输出出去,然后,with语句就可以正常地工作了:

with create_query('Bob') as q:
    q.query()

很多时候,我们希望在某段代码执行前后自动执行特定代码,也可以用@contextmanager实现。例如:

@contextmanager
def tag(name):
    print("<%s>" % name)
    yield
    print("</%s>" % name)

with tag("h1"):
    print("hello")
    print("world")

上述代码执行结果为:

<h1>
hello
world
</h1>

代码的执行顺序是:

with语句首先执行yield之前的语句,因此打印出<h1>;
yield调用会执行with语句内部的所有语句,因此打印出hello和world;
最后执行yield之后的语句,打印出</h1>。
因此,@contextmanager让我们通过编写generator来简化上下文管理。

@closing
如果一个对象没有实现上下文,我们就不能把它用于with语句。这个时候,可以用closing()来把该对象变为上下文对象。例如,用with语句使用urlopen():

from contextlib import closing
from urllib.request import urlopen

with closing(urlopen('https://www.python.org')) as page:
    for line in page:
        print(line)

closing也是一个经过@contextmanager装饰的generator,这个generator编写起来其实非常简单:

@contextmanager
def closing(thing):
    try:
        yield thing
    finally:
        thing.close()

它的作用就是把任意对象变为上下文对象,并支持with语句。

以下是详细版:
在上文中我们提到with语句中的上下文管理器。with语句可以如此简单但强大,主要依赖于上下文管理器。那么什么是上下文管理器?上下文管理器就是实现了上下文协议的类,而上下文协议就是一个类要实现__enter__()__exit__()两个方法。一个类只要实现了__enter__()__exit__(),我们就称之为上下文管理器下面我们具体说下这两个方法。

__enter__():主要执行一些环境准备工作,同时返回一资源对象。如果上下文管理器open("test.txt")的__enter__()函数返回一个文件对象。

__exit__():完整形式为__exit__(type, value, traceback),这三个参数和调用sys.exec_info()函数返回值是一样的,分别为异常类型、异常信息和堆栈。如果执行体语句没有引发异常,则这三个参数均被设为None。否则,它们将包含上下文的异常信息。__exit_()方法返回True或False,分别指示被引发的异常有没有被处理,如果返回False,引发的异常将会被传递出上下文。如果__exit__()函数内部引发了异常,则会覆盖掉执行体的中引发的异常。处理异常时,不需要重新抛出异常,只需要返回False,with语句会检测__exit__()返回False来处理异常。

如果我们要自定义一个上下文管理器,只需要定义一个类并且是实现__enter__()__exit__()即可。下面通过一个简单的例子是演示如果新建自定义的上下文管理器,我们以数据库的连接为例。在使用数据库时,有时要涉及到事务操作。数据库的事务操作当调用commit()执行sql命令时,如果在这个过程中执行失败,则需要执行rollback()回滚数据库,通常实现方式可能如下:

def test_write():
    con = MySQLdb.connection()
    cursor = con.cursor()
    sql = """ #具体的sql语句 """
    try:
        cursor.execute(sql)
        cursor.execute(sql)
        cursor.execute(sql)
        con.commit()      #提交事务 except Exception as ex:
        con.rollback()    #事务执行失败,回滚数据库

如果想通过with语句来实现数据库执行失败的回滚操作,则我们需要自定义一个数据库连接的上下文管理器,假设为DBConnection,则我们将上面例子用with语句来实现的话,应该是这样子的,如下:

def test_write():
    sql = """ #具体的sql语句 """ 
    con = DBConnection()
    with con as cursor:   
        cursor.execute(sql)
        cursor.execute(sql)
        cursor.execute(sql)

要实现上面with语句的功能,则我们的DBConnection数据库上下文管理器则需要提供一下功能:__enter__()要返回一个连接的cursor; 当没有异常发生是,__exit__()函数commit所有的数据库操作。如果有异常发生则_exit__()会回滚数据库,调用rollback()。所以我们可以实现DBConnection如下:

def DBConnection(object): 
    def __init__(self):
         pass 
    def cursor(self): 
         #返回一个游标并且启动一个事务
          pass
    def commit(self):          #提交当前事务
         pass
    def rollback(self):         #回滚当前事务
         pass
    def __enter__(self):          #返回一个cursor
         cursor = self.cursor() 
         return cursor
    def __exit__(self, type, value, tb): 
         if tb is None: 
             #没有异常则提交事务
             self.commit()         
         else:              #有异常则回滚数据库
             self.rollback()
contextlib模块
  • contextmanage对象

上文提到如果我们要实现一个自定义的上下文管理器,需要定义一个实现了__enter____exit__两个方法的类, 这显示不是很方便。Python的contextlib模块给我们提供了更方便的方式来实现一个自定义的上下文管理器。contextlib模块包含一个装饰器contextmanager和一些辅助函数,装饰器contextmanager只需要写一个生成器函数就可以代替自定义的上下文管理器,典型用法如下:
  需要使用yield先定义一个生成器函数.

 @contextmanager 
def some_generator(<arguments>):
            <setup>
            try: 
                 yield <value>
            finally: 
                 <cleanup>

然后便可以用with语句调用contextmanage生成的上下文管理器了,with语句用法如下:

with some_generator(<arguments>) as <variable>: 
    <body>

生成器函数some_generator就和我们普通的函数一样,它的原理如下:

  1. some_generator函数在在yield之前的代码等同于上下文管理器中的__enter__函数。
  2. yield的返回值等同于__enter__函数的返回值,即如果with语句声明了as <variable>,则yield的值会赋给variable
  3. 然后执行<cleanup>代码块,等同于上下文管理器的__exit__函数。此时发生的任何异常都会再次通过yield函数返回。

下面举几个简单的例子,

例子1:锁资源自动获取和释放的例子

@contextmanager
 def locked(lock): 
     lock.acquire()
     try:
         yield
     finally:
         lock.release()
 
 with locked(myLock): 
     #代码执行到这里时,myLock已经自动上锁
     pass
     #执行完后会,会自动释放锁

例子2:文件打开后自动管理的实现

 @contextmanager
 def myopen(filename, mode="r"):
     f = open(filename,mode) 4     try:
         yield f
     finally:
         f.close()
 
 with myopen("test.txt") as f:
     for line in f: 
         print(line)

例子3:数据库事务的处理

@contextmanager def transaction(db):
    db.begin() try: yield 
    except:
        db.rollback() raise
    else:
        db.commit()

with transaction(mydb):
    mydb.cursor.execute(sql)
    mydb.cursor.execute(sql)
    mydb.cursor.execute(sql)
    mydb.cursor.execute(sql)
  • nested函数

contextlib模块还提供了一个函数给我们:nested(mgr1,mgr2...mgrn)函数,用来嵌套多个上下文管理器,等同于下面的形式:

with mgr1:
    with mgr2:
        ...
        with mgrn: pass</pre>

但是with语句本身已经支持了多个下文管理器的使用,所以nested的意义不是很大。我们可以写一个例子来看下nested函数的使用,以及与直接使用with来嵌套多个上下文管理器的区别,如下所示:

from contextlib import contextmanager 
from contextlib import nested 
from contextlib import closing 
@contextmanager
def my_context(name):
     print("enter")
     try:
         yield name 
     finally:
         print("exit") 
 #使用nested函数来调用多个管理器
print("---------使用nested函数调用多个管理器-----------") 
with nested(my_context("管理器一"), my_context("管理器二"),my_context("管理器三")) as (m1,m2,m3): 
     print(m1)
     print(m2) 
     print(m3) 
 #直接使用with来调用调用多个管理器
     print("---------使用with调用多个管理器-----------")
with my_context("管理器一") as m1, my_context("管理器二") as m2, my_context("管理器三") as m3: 
     print(m1) 
     print(m2) 
     print(m3)

输出结果为:

  • closing对象

contextlib中还包含一个closing对象,这个对象就是一个上下文管理器,它的exit函数仅仅调用传入参数的close函数,closing对象的源码如下:

class closing(object):
     def __init__(self, thing):
         self.thing = thing 
     def __enter__(self):
          return self.thing 
     def __exit__(self, *exc_info): 
         self.thing.close()

所以closeing上下文管理器仅使用于具有close()方法的资源对象。例如,如果我们通过urllib.urlopen打开一个网页,urlopen返回的request有close方法,所以我们就可以使用closing上下文管理器,如下:

import urllib, sys from contextlib import closing
with closing(urllib.urlopen('http://www.yahoo.com')) as f: for line in f:
        sys.stdout.write(line)
最后编辑于
©著作权归作者所有,转载或内容合作请联系作者
  • 序言:七十年代末,一起剥皮案震惊了整个滨河市,随后出现的几起案子,更是在滨河造成了极大的恐慌,老刑警刘岩,带你破解...
    沈念sama阅读 204,445评论 6 478
  • 序言:滨河连续发生了三起死亡事件,死亡现场离奇诡异,居然都是意外死亡,警方通过查阅死者的电脑和手机,发现死者居然都...
    沈念sama阅读 85,889评论 2 381
  • 文/潘晓璐 我一进店门,熙熙楼的掌柜王于贵愁眉苦脸地迎上来,“玉大人,你说我怎么就摊上这事。” “怎么了?”我有些...
    开封第一讲书人阅读 151,047评论 0 337
  • 文/不坏的土叔 我叫张陵,是天一观的道长。 经常有香客问我,道长,这世上最难降的妖魔是什么? 我笑而不...
    开封第一讲书人阅读 54,760评论 1 276
  • 正文 为了忘掉前任,我火速办了婚礼,结果婚礼上,老公的妹妹穿的比我还像新娘。我一直安慰自己,他们只是感情好,可当我...
    茶点故事阅读 63,745评论 5 367
  • 文/花漫 我一把揭开白布。 她就那样静静地躺着,像睡着了一般。 火红的嫁衣衬着肌肤如雪。 梳的纹丝不乱的头发上,一...
    开封第一讲书人阅读 48,638评论 1 281
  • 那天,我揣着相机与录音,去河边找鬼。 笑死,一个胖子当着我的面吹牛,可吹牛的内容都是我干的。 我是一名探鬼主播,决...
    沈念sama阅读 38,011评论 3 398
  • 文/苍兰香墨 我猛地睁开眼,长吁一口气:“原来是场噩梦啊……” “哼!你这毒妇竟也来了?” 一声冷哼从身侧响起,我...
    开封第一讲书人阅读 36,669评论 0 258
  • 序言:老挝万荣一对情侣失踪,失踪者是张志新(化名)和其女友刘颖,没想到半个月后,有当地人在树林里发现了一具尸体,经...
    沈念sama阅读 40,923评论 1 299
  • 正文 独居荒郊野岭守林人离奇死亡,尸身上长有42处带血的脓包…… 初始之章·张勋 以下内容为张勋视角 年9月15日...
    茶点故事阅读 35,655评论 2 321
  • 正文 我和宋清朗相恋三年,在试婚纱的时候发现自己被绿了。 大学时的朋友给我发了我未婚夫和他白月光在一起吃饭的照片。...
    茶点故事阅读 37,740评论 1 330
  • 序言:一个原本活蹦乱跳的男人离奇死亡,死状恐怖,灵堂内的尸体忽然破棺而出,到底是诈尸还是另有隐情,我是刑警宁泽,带...
    沈念sama阅读 33,406评论 4 320
  • 正文 年R本政府宣布,位于F岛的核电站,受9级特大地震影响,放射性物质发生泄漏。R本人自食恶果不足惜,却给世界环境...
    茶点故事阅读 38,995评论 3 307
  • 文/蒙蒙 一、第九天 我趴在偏房一处隐蔽的房顶上张望。 院中可真热闹,春花似锦、人声如沸。这庄子的主人今日做“春日...
    开封第一讲书人阅读 29,961评论 0 19
  • 文/苍兰香墨 我抬头看了看天上的太阳。三九已至,却和暖如春,着一层夹袄步出监牢的瞬间,已是汗流浃背。 一阵脚步声响...
    开封第一讲书人阅读 31,197评论 1 260
  • 我被黑心中介骗来泰国打工, 没想到刚下飞机就差点儿被人妖公主榨干…… 1. 我叫王不留,地道东北人。 一个月前我还...
    沈念sama阅读 45,023评论 2 350
  • 正文 我出身青楼,却偏偏与公主长得像,于是被迫代替她去往敌国和亲。 传闻我的和亲对象是个残疾皇子,可洞房花烛夜当晚...
    茶点故事阅读 42,483评论 2 342