1import sys
2
3
4def typename(t):
5    name = type(t).__name__
6    if sys.version_info < (2,5):
7        if name == 'classobj' and issubclass(t, MyException):
8            name = 'type'
9        elif name == 'instance' and isinstance(t, MyException):
10            name = 'MyException'
11    return "<type '%s'>" % name
12
13
14class MyException(Exception):
15    pass
16
17
18class ContextManager(object):
19    def __init__(self, value, exit_ret = None):
20        self.value = value
21        self.exit_ret = exit_ret
22
23    def __exit__(self, a, b, tb):
24        print("exit %s %s %s" % (typename(a), typename(b), typename(tb)))
25        return self.exit_ret
26
27    def __enter__(self):
28        print("enter")
29        return self.value
30
31
32def multimanager():
33    """
34    >>> multimanager()
35    enter
36    enter
37    enter
38    enter
39    enter
40    enter
41    2
42    value
43    1 2 3 4 5
44    nested
45    exit <type 'NoneType'> <type 'NoneType'> <type 'NoneType'>
46    exit <type 'NoneType'> <type 'NoneType'> <type 'NoneType'>
47    exit <type 'NoneType'> <type 'NoneType'> <type 'NoneType'>
48    exit <type 'NoneType'> <type 'NoneType'> <type 'NoneType'>
49    exit <type 'NoneType'> <type 'NoneType'> <type 'NoneType'>
50    exit <type 'NoneType'> <type 'NoneType'> <type 'NoneType'>
51    """
52    with ContextManager(1), ContextManager(2) as x, ContextManager('value') as y,\
53            ContextManager(3), ContextManager((1, 2, (3, (4, 5)))) as (a, b, (c, (d, e))):
54        with ContextManager('nested') as nested:
55            print(x)
56            print(y)
57            print('%s %s %s %s %s' % (a, b, c, d, e))
58            print(nested)
59
60
61class GetManager(object):
62    def get(self, *args):
63        return ContextManager(*args)
64
65def manager_from_expression():
66    """
67    >>> manager_from_expression()
68    enter
69    1
70    exit <type 'NoneType'> <type 'NoneType'> <type 'NoneType'>
71    enter
72    2
73    exit <type 'NoneType'> <type 'NoneType'> <type 'NoneType'>
74    """
75    with GetManager().get(1) as x:
76        print(x)
77    g = GetManager()
78    with g.get(2) as x:
79        print(x)
80
81
82# Tests borrowed from pyregr test_with.py,
83# modified to follow the constraints of Cython.
84import unittest
85
86class Dummy(object):
87    def __init__(self, value=None, gobble=False):
88        if value is None:
89            value = self
90        self.value = value
91        self.gobble = gobble
92        self.enter_called = False
93        self.exit_called = False
94
95    def __enter__(self):
96        self.enter_called = True
97        return self.value
98
99    def __exit__(self, *exc_info):
100        self.exit_called = True
101        self.exc_info = exc_info
102        if self.gobble:
103            return True
104
105class InitRaises(object):
106    def __init__(self): raise RuntimeError()
107
108class EnterRaises(object):
109    def __enter__(self): raise RuntimeError()
110    def __exit__(self, *exc_info): pass
111
112class ExitRaises(object):
113    def __enter__(self): pass
114    def __exit__(self, *exc_info): raise RuntimeError()
115
116class NestedWith(unittest.TestCase):
117    """
118    >>> NestedWith().runTest()
119    """
120
121    def runTest(self):
122        self.testNoExceptions()
123        self.testExceptionInExprList()
124        self.testExceptionInEnter()
125        self.testExceptionInExit()
126        self.testEnterReturnsTuple()
127
128    def testNoExceptions(self):
129        with Dummy() as a, Dummy() as b:
130            self.assertTrue(a.enter_called)
131            self.assertTrue(b.enter_called)
132        self.assertTrue(a.exit_called)
133        self.assertTrue(b.exit_called)
134
135    def testExceptionInExprList(self):
136        try:
137            with Dummy() as a, InitRaises():
138                pass
139        except:
140            pass
141        self.assertTrue(a.enter_called)
142        self.assertTrue(a.exit_called)
143
144    def testExceptionInEnter(self):
145        try:
146            with Dummy() as a, EnterRaises():
147                self.fail('body of bad with executed')
148        except RuntimeError:
149            pass
150        else:
151            self.fail('RuntimeError not reraised')
152        self.assertTrue(a.enter_called)
153        self.assertTrue(a.exit_called)
154
155    def testExceptionInExit(self):
156        body_executed = False
157        with Dummy(gobble=True) as a, ExitRaises():
158            body_executed = True
159        self.assertTrue(a.enter_called)
160        self.assertTrue(a.exit_called)
161        self.assertTrue(body_executed)
162        self.assertNotEqual(a.exc_info[0], None)
163
164    def testEnterReturnsTuple(self):
165        with Dummy(value=(1,2)) as (a1, a2), \
166             Dummy(value=(10, 20)) as (b1, b2):
167            self.assertEqual(1, a1)
168            self.assertEqual(2, a2)
169            self.assertEqual(10, b1)
170            self.assertEqual(20, b2)
171