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