14adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao"""Utilities for with-statement contexts. See PEP 343.""" 24adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao 34adfde8bc82dd39f59e0445588c3e599ada477dJosh Gaoimport sys 44adfde8bc82dd39f59e0445588c3e599ada477dJosh Gaofrom functools import wraps 54adfde8bc82dd39f59e0445588c3e599ada477dJosh Gaofrom warnings import warn 64adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao 74adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao__all__ = ["contextmanager", "nested", "closing"] 84adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao 94adfde8bc82dd39f59e0445588c3e599ada477dJosh Gaoclass GeneratorContextManager(object): 104adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao """Helper for @contextmanager decorator.""" 114adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao 124adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao def __init__(self, gen): 134adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao self.gen = gen 144adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao 154adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao def __enter__(self): 164adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao try: 174adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao return self.gen.next() 184adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao except StopIteration: 194adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao raise RuntimeError("generator didn't yield") 204adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao 214adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao def __exit__(self, type, value, traceback): 224adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao if type is None: 234adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao try: 244adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao self.gen.next() 254adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao except StopIteration: 264adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao return 274adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao else: 284adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao raise RuntimeError("generator didn't stop") 294adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao else: 304adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao if value is None: 314adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao # Need to force instantiation so we can reliably 324adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao # tell if we get the same exception back 334adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao value = type() 344adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao try: 354adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao self.gen.throw(type, value, traceback) 364adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao raise RuntimeError("generator didn't stop after throw()") 374adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao except StopIteration, exc: 384adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao # Suppress the exception *unless* it's the same exception that 394adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao # was passed to throw(). This prevents a StopIteration 404adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao # raised inside the "with" statement from being suppressed 414adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao return exc is not value 424adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao except: 434adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao # only re-raise if it's *not* the exception that was 444adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao # passed to throw(), because __exit__() must not raise 454adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao # an exception unless __exit__() itself failed. But throw() 464adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao # has to raise the exception to signal propagation, so this 474adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao # fixes the impedance mismatch between the throw() protocol 484adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao # and the __exit__() protocol. 494adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao # 504adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao if sys.exc_info()[1] is not value: 514adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao raise 524adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao 534adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao 544adfde8bc82dd39f59e0445588c3e599ada477dJosh Gaodef contextmanager(func): 554adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao """@contextmanager decorator. 564adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao 574adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao Typical usage: 584adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao 594adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao @contextmanager 604adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao def some_generator(<arguments>): 614adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao <setup> 624adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao try: 634adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao yield <value> 644adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao finally: 654adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao <cleanup> 664adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao 674adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao This makes this: 684adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao 694adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao with some_generator(<arguments>) as <variable>: 704adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao <body> 714adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao 724adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao equivalent to this: 734adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao 744adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao <setup> 754adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao try: 764adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao <variable> = <value> 774adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao <body> 784adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao finally: 794adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao <cleanup> 804adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao 814adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao """ 824adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao @wraps(func) 834adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao def helper(*args, **kwds): 844adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao return GeneratorContextManager(func(*args, **kwds)) 854adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao return helper 864adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao 874adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao 884adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao@contextmanager 894adfde8bc82dd39f59e0445588c3e599ada477dJosh Gaodef nested(*managers): 904adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao """Combine multiple context managers into a single nested context manager. 914adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao 924adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao This function has been deprecated in favour of the multiple manager form 934adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao of the with statement. 944adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao 954adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao The one advantage of this function over the multiple manager form of the 964adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao with statement is that argument unpacking allows it to be 974adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao used with a variable number of context managers as follows: 984adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao 994adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao with nested(*managers): 1004adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao do_something() 1014adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao 1024adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao """ 1034adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao warn("With-statements now directly support multiple context managers", 1044adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao DeprecationWarning, 3) 1054adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao exits = [] 1064adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao vars = [] 1074adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao exc = (None, None, None) 1084adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao try: 1094adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao for mgr in managers: 1104adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao exit = mgr.__exit__ 1114adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao enter = mgr.__enter__ 1124adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao vars.append(enter()) 1134adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao exits.append(exit) 1144adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao yield vars 1154adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao except: 1164adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao exc = sys.exc_info() 1174adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao finally: 1184adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao while exits: 1194adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao exit = exits.pop() 1204adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao try: 1214adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao if exit(*exc): 1224adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao exc = (None, None, None) 1234adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao except: 1244adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao exc = sys.exc_info() 1254adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao if exc != (None, None, None): 1264adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao # Don't rely on sys.exc_info() still containing 1274adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao # the right information. Another exception may 1284adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao # have been raised and caught by an exit method 1294adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao raise exc[0], exc[1], exc[2] 1304adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao 1314adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao 1324adfde8bc82dd39f59e0445588c3e599ada477dJosh Gaoclass closing(object): 1334adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao """Context to automatically close something at the end of a block. 1344adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao 1354adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao Code like this: 1364adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao 1374adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao with closing(<module>.open(<arguments>)) as f: 1384adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao <block> 1394adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao 1404adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao is equivalent to this: 1414adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao 1424adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao f = <module>.open(<arguments>) 1434adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao try: 1444adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao <block> 1454adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao finally: 1464adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao f.close() 1474adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao 1484adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao """ 1494adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao def __init__(self, thing): 1504adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao self.thing = thing 1514adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao def __enter__(self): 1524adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao return self.thing 1534adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao def __exit__(self, *exc_info): 1544adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao self.thing.close() 155