##// END OF EJS Templates
util: introduce ctxmanager, to avoid nested try/finally blocks...
Bryan O'Sullivan -
r27703:4e27c0a7 default
parent child Browse files
Show More
@@ -0,0 +1,77 b''
1 from __future__ import absolute_import
2
3 import silenttestrunner
4 import unittest
5
6 from mercurial.util import ctxmanager
7
8 class contextmanager(object):
9 def __init__(self, name, trace):
10 self.name = name
11 self.entered = False
12 self.exited = False
13 self.trace = trace
14
15 def __enter__(self):
16 self.entered = True
17 self.trace(('enter', self.name))
18 return self
19
20 def __exit__(self, exc_type, exc_val, exc_tb):
21 self.exited = exc_type, exc_val, exc_tb
22 self.trace(('exit', self.name))
23
24 def __repr__(self):
25 return '<ctx %r>' % self.name
26
27 class ctxerror(Exception):
28 pass
29
30 class raise_on_enter(contextmanager):
31 def __enter__(self):
32 self.trace(('raise', self.name))
33 raise ctxerror(self.name)
34
35 class raise_on_exit(contextmanager):
36 def __exit__(self, exc_type, exc_val, exc_tb):
37 self.trace(('raise', self.name))
38 raise ctxerror(self.name)
39
40 def ctxmgr(name, trace):
41 return lambda: contextmanager(name, trace)
42
43 class test_ctxmanager(unittest.TestCase):
44 def test_basics(self):
45 trace = []
46 addtrace = trace.append
47 with ctxmanager(ctxmgr('a', addtrace), ctxmgr('b', addtrace)) as c:
48 a, b = c()
49 c.atexit(addtrace, ('atexit', 'x'))
50 c.atexit(addtrace, ('atexit', 'y'))
51 self.assertEqual(trace, [('enter', 'a'), ('enter', 'b'),
52 ('atexit', 'y'), ('atexit', 'x'),
53 ('exit', 'b'), ('exit', 'a')])
54
55 def test_raise_on_enter(self):
56 trace = []
57 addtrace = trace.append
58 with self.assertRaises(ctxerror):
59 with ctxmanager(ctxmgr('a', addtrace),
60 lambda: raise_on_enter('b', addtrace)) as c:
61 c()
62 addtrace('unreachable')
63 self.assertEqual(trace, [('enter', 'a'), ('raise', 'b'), ('exit', 'a')])
64
65 def test_raise_on_exit(self):
66 trace = []
67 addtrace = trace.append
68 with self.assertRaises(ctxerror):
69 with ctxmanager(ctxmgr('a', addtrace),
70 lambda: raise_on_exit('b', addtrace)) as c:
71 c()
72 addtrace('running')
73 self.assertEqual(trace, [('enter', 'a'), ('enter', 'b'), 'running',
74 ('raise', 'b'), ('exit', 'a')])
75
76 if __name__ == '__main__':
77 silenttestrunner.main(__name__)
@@ -2641,6 +2641,66 b' def _makedecompressor(decompcls):'
2641 return chunkbuffer(generator(fh))
2641 return chunkbuffer(generator(fh))
2642 return func
2642 return func
2643
2643
2644 class ctxmanager(object):
2645 '''A context manager for use in 'with' blocks to allow multiple
2646 contexts to be entered at once. This is both safer and more
2647 flexible than contextlib.nested.
2648
2649 Once Mercurial supports Python 2.7+, this will become mostly
2650 unnecessary.
2651 '''
2652
2653 def __init__(self, *args):
2654 '''Accepts a list of no-argument functions that return context
2655 managers. These will be invoked at __call__ time.'''
2656 self._pending = args
2657 self._atexit = []
2658
2659 def __enter__(self):
2660 return self
2661
2662 def __call__(self):
2663 '''Create and enter context managers in the order in which they were
2664 passed to the constructor.'''
2665 values = []
2666 for func in self._pending:
2667 obj = func()
2668 values.append(obj.__enter__())
2669 self._atexit.append(obj.__exit__)
2670 del self._pending
2671 return values
2672
2673 def atexit(self, func, *args, **kwargs):
2674 '''Add a function to call when this context manager exits. The
2675 ordering of multiple atexit calls is unspecified, save that
2676 they will happen before any __exit__ functions.'''
2677 def wrapper(exc_type, exc_val, exc_tb):
2678 func(*args, **kwargs)
2679 self._atexit.append(wrapper)
2680 return func
2681
2682 def __exit__(self, exc_type, exc_val, exc_tb):
2683 '''Context managers are exited in the reverse order from which
2684 they were created.'''
2685 received = exc_type is not None
2686 suppressed = False
2687 pending = None
2688 self._atexit.reverse()
2689 for exitfunc in self._atexit:
2690 try:
2691 if exitfunc(exc_type, exc_val, exc_tb):
2692 suppressed = True
2693 exc_type = None
2694 exc_val = None
2695 exc_tb = None
2696 except BaseException as e:
2697 pending = sys.exc_info()
2698 exc_type, exc_val, exc_tb = pending = sys.exc_info()
2699 del self._atexit
2700 if pending:
2701 raise exc_val
2702 return received and suppressed
2703
2644 def _bz2():
2704 def _bz2():
2645 d = bz2.BZ2Decompressor()
2705 d = bz2.BZ2Decompressor()
2646 # Bzip2 stream start with BZ, but we stripped it.
2706 # Bzip2 stream start with BZ, but we stripped it.
General Comments 0
You need to be logged in to leave comments. Login now