1 #!/usr/bin/env python

     2 

     3 """Unit tests for the with statement specified in PEP 343."""

     4 

     5 

     6 __author__ = "Mike Bland"

     7 __email__ = "mbland at acm dot org"

     8 

     9 import sys

    10 import unittest

    11 from collections import deque

    12 from contextlib import GeneratorContextManager, contextmanager

    13 from test.test_support import run_unittest

    14 

    15 

    16 class MockContextManager(GeneratorContextManager):

    17     def __init__(self, gen):

    18         GeneratorContextManager.__init__(self, gen)

    19         self.enter_called = False

    20         self.exit_called = False

    21         self.exit_args = None

    22 

    23     def __enter__(self):

    24         self.enter_called = True

    25         return GeneratorContextManager.__enter__(self)

    26 

    27     def __exit__(self, type, value, traceback):

    28         self.exit_called = True

    29         self.exit_args = (type, value, traceback)

    30         return GeneratorContextManager.__exit__(self, type,

    31                                                 value, traceback)

    32 

    33 

    34 def mock_contextmanager(func):

    35     def helper(*args, **kwds):

    36         return MockContextManager(func(*args, **kwds))

    37     return helper

    38 

    39 

    40 class MockResource(object):

    41     def __init__(self):

    42         self.yielded = False

    43         self.stopped = False

    44 

    45 

    46 @mock_contextmanager

    47 def mock_contextmanager_generator():

    48     mock = MockResource()

    49     try:

    50         mock.yielded = True

    51         yield mock

    52     finally:

    53         mock.stopped = True

    54 

    55 

    56 class Nested(object):

    57 

    58     def __init__(self, *managers):

    59         self.managers = managers

    60         self.entered = None

    61 

    62     def __enter__(self):

    63         if self.entered is not None:

    64             raise RuntimeError("Context is not reentrant")

    65         self.entered = deque()

    66         vars = []

    67         try:

    68             for mgr in self.managers:

    69                 vars.append(mgr.__enter__())

    70                 self.entered.appendleft(mgr)

    71         except:

    72             if not self.__exit__(*sys.exc_info()):

    73                 raise

    74         return vars

    75 

    76     def __exit__(self, *exc_info):

    77         # Behave like nested with statements

    78         # first in, last out

    79         # New exceptions override old ones

    80         ex = exc_info

    81         for mgr in self.entered:

    82             try:

    83                 if mgr.__exit__(*ex):

    84                     ex = (None, None, None)

    85             except:

    86                 ex = sys.exc_info()

    87         self.entered = None

    88         if ex is not exc_info:

    89             raise ex[0], ex[1], ex[2]

    90 

    91 

    92 class MockNested(Nested):

    93     def __init__(self, *managers):

    94         Nested.__init__(self, *managers)

    95         self.enter_called = False

    96         self.exit_called = False

    97         self.exit_args = None

    98 

    99     def __enter__(self):

   100         self.enter_called = True

   101         return Nested.__enter__(self)

   102 

   103     def __exit__(self, *exc_info):

   104         self.exit_called = True

   105         self.exit_args = exc_info

   106         return Nested.__exit__(self, *exc_info)

   107 

   108 

   109 class FailureTestCase(unittest.TestCase):

   110     def testNameError(self):

   111         def fooNotDeclared():

   112             with foo: pass

   113         self.assertRaises(NameError, fooNotDeclared)

   114 

   115     def testEnterAttributeError(self):

   116         class LacksEnter(object):

   117             def __exit__(self, type, value, traceback):

   118                 pass

   119 

   120         def fooLacksEnter():

   121             foo = LacksEnter()

   122             with foo: pass

   123         self.assertRaises(AttributeError, fooLacksEnter)

   124 

   125     def testExitAttributeError(self):

   126         class LacksExit(object):

   127             def __enter__(self):

   128                 pass

   129 

   130         def fooLacksExit():

   131             foo = LacksExit()

   132             with foo: pass

   133         self.assertRaises(AttributeError, fooLacksExit)

   134 

   135     def assertRaisesSyntaxError(self, codestr):

   136         def shouldRaiseSyntaxError(s):

   137             compile(s, '', 'single')

   138         self.assertRaises(SyntaxError, shouldRaiseSyntaxError, codestr)

   139 

   140     def testAssignmentToNoneError(self):

   141         self.assertRaisesSyntaxError('with mock as None:\n  pass')

   142         self.assertRaisesSyntaxError(

   143             'with mock as (None):\n'

   144             '  pass')

   145 

   146     def testAssignmentToEmptyTupleError(self):

   147         self.assertRaisesSyntaxError(

   148             'with mock as ():\n'

   149             '  pass')

   150 

   151     def testAssignmentToTupleOnlyContainingNoneError(self):

   152         self.assertRaisesSyntaxError('with mock as None,:\n  pass')

   153         self.assertRaisesSyntaxError(

   154             'with mock as (None,):\n'

   155             '  pass')

   156 

   157     def testAssignmentToTupleContainingNoneError(self):

   158         self.assertRaisesSyntaxError(

   159             'with mock as (foo, None, bar):\n'

   160             '  pass')

   161 

   162     def testEnterThrows(self):

   163         class EnterThrows(object):

   164             def __enter__(self):

   165                 raise RuntimeError("Enter threw")

   166             def __exit__(self, *args):

   167                 pass

   168 

   169         def shouldThrow():

   170             ct = EnterThrows()

   171             self.foo = None

   172             with ct as self.foo:

   173                 pass

   174         self.assertRaises(RuntimeError, shouldThrow)

   175         self.assertEqual(self.foo, None)

   176 

   177     def testExitThrows(self):

   178         class ExitThrows(object):

   179             def __enter__(self):

   180                 return

   181             def __exit__(self, *args):

   182                 raise RuntimeError(42)

   183         def shouldThrow():

   184             with ExitThrows():

   185                 pass

   186         self.assertRaises(RuntimeError, shouldThrow)

   187 

   188 class ContextmanagerAssertionMixin(object):

   189     TEST_EXCEPTION = RuntimeError("test exception")

   190 

   191     def assertInWithManagerInvariants(self, mock_manager):

   192         self.assertTrue(mock_manager.enter_called)

   193         self.assertFalse(mock_manager.exit_called)

   194         self.assertEqual(mock_manager.exit_args, None)

   195 

   196     def assertAfterWithManagerInvariants(self, mock_manager, exit_args):

   197         self.assertTrue(mock_manager.enter_called)

   198         self.assertTrue(mock_manager.exit_called)

   199         self.assertEqual(mock_manager.exit_args, exit_args)

   200 

   201     def assertAfterWithManagerInvariantsNoError(self, mock_manager):

   202         self.assertAfterWithManagerInvariants(mock_manager,

   203             (None, None, None))

   204 

   205     def assertInWithGeneratorInvariants(self, mock_generator):

   206         self.assertTrue(mock_generator.yielded)

   207         self.assertFalse(mock_generator.stopped)

   208 

   209     def assertAfterWithGeneratorInvariantsNoError(self, mock_generator):

   210         self.assertTrue(mock_generator.yielded)

   211         self.assertTrue(mock_generator.stopped)

   212 

   213     def raiseTestException(self):

   214         raise self.TEST_EXCEPTION

   215 

   216     def assertAfterWithManagerInvariantsWithError(self, mock_manager):

   217         self.assertTrue(mock_manager.enter_called)

   218         self.assertTrue(mock_manager.exit_called)

   219         self.assertEqual(mock_manager.exit_args[0], RuntimeError)

   220         self.assertEqual(mock_manager.exit_args[1], self.TEST_EXCEPTION)

   221 

   222     def assertAfterWithGeneratorInvariantsWithError(self, mock_generator):

   223         self.assertTrue(mock_generator.yielded)

   224         self.assertTrue(mock_generator.stopped)

   225 

   226 

   227 class NonexceptionalTestCase(unittest.TestCase, ContextmanagerAssertionMixin):

   228     def testInlineGeneratorSyntax(self):

   229         with mock_contextmanager_generator():

   230             pass

   231 

   232     def testUnboundGenerator(self):

   233         mock = mock_contextmanager_generator()

   234         with mock:

   235             pass

   236         self.assertAfterWithManagerInvariantsNoError(mock)

   237 

   238     def testInlineGeneratorBoundSyntax(self):

   239         with mock_contextmanager_generator() as foo:

   240             self.assertInWithGeneratorInvariants(foo)

   241         # FIXME: In the future, we'll try to keep the bound names from leaking

   242         self.assertAfterWithGeneratorInvariantsNoError(foo)

   243 

   244     def testInlineGeneratorBoundToExistingVariable(self):

   245         foo = None

   246         with mock_contextmanager_generator() as foo:

   247             self.assertInWithGeneratorInvariants(foo)

   248         self.assertAfterWithGeneratorInvariantsNoError(foo)

   249 

   250     def testInlineGeneratorBoundToDottedVariable(self):

   251         with mock_contextmanager_generator() as self.foo:

   252             self.assertInWithGeneratorInvariants(self.foo)

   253         self.assertAfterWithGeneratorInvariantsNoError(self.foo)

   254 

   255     def testBoundGenerator(self):

   256         mock = mock_contextmanager_generator()

   257         with mock as foo:

   258             self.assertInWithGeneratorInvariants(foo)

   259             self.assertInWithManagerInvariants(mock)

   260         self.assertAfterWithGeneratorInvariantsNoError(foo)

   261         self.assertAfterWithManagerInvariantsNoError(mock)

   262 

   263     def testNestedSingleStatements(self):

   264         mock_a = mock_contextmanager_generator()

   265         with mock_a as foo:

   266             mock_b = mock_contextmanager_generator()

   267             with mock_b as bar:

   268                 self.assertInWithManagerInvariants(mock_a)

   269                 self.assertInWithManagerInvariants(mock_b)

   270                 self.assertInWithGeneratorInvariants(foo)

   271                 self.assertInWithGeneratorInvariants(bar)

   272             self.assertAfterWithManagerInvariantsNoError(mock_b)

   273             self.assertAfterWithGeneratorInvariantsNoError(bar)

   274             self.assertInWithManagerInvariants(mock_a)

   275             self.assertInWithGeneratorInvariants(foo)

   276         self.assertAfterWithManagerInvariantsNoError(mock_a)

   277         self.assertAfterWithGeneratorInvariantsNoError(foo)

   278 

   279 

   280 class NestedNonexceptionalTestCase(unittest.TestCase,

   281     ContextmanagerAssertionMixin):

   282     def testSingleArgInlineGeneratorSyntax(self):

   283         with Nested(mock_contextmanager_generator()):

   284             pass

   285 

   286     def testSingleArgUnbound(self):

   287         mock_contextmanager = mock_contextmanager_generator()

   288         mock_nested = MockNested(mock_contextmanager)

   289         with mock_nested:

   290             self.assertInWithManagerInvariants(mock_contextmanager)

   291             self.assertInWithManagerInvariants(mock_nested)

   292         self.assertAfterWithManagerInvariantsNoError(mock_contextmanager)

   293         self.assertAfterWithManagerInvariantsNoError(mock_nested)

   294 

   295     def testSingleArgBoundToNonTuple(self):

   296         m = mock_contextmanager_generator()

   297         # This will bind all the arguments to nested() into a single list

   298         # assigned to foo.

   299         with Nested(m) as foo:

   300             self.assertInWithManagerInvariants(m)

   301         self.assertAfterWithManagerInvariantsNoError(m)

   302 

   303     def testSingleArgBoundToSingleElementParenthesizedList(self):

   304         m = mock_contextmanager_generator()

   305         # This will bind all the arguments to nested() into a single list

   306         # assigned to foo.

   307         with Nested(m) as (foo):

   308             self.assertInWithManagerInvariants(m)

   309         self.assertAfterWithManagerInvariantsNoError(m)

   310 

   311     def testSingleArgBoundToMultipleElementTupleError(self):

   312         def shouldThrowValueError():

   313             with Nested(mock_contextmanager_generator()) as (foo, bar):

   314                 pass

   315         self.assertRaises(ValueError, shouldThrowValueError)

   316 

   317     def testSingleArgUnbound(self):

   318         mock_contextmanager = mock_contextmanager_generator()

   319         mock_nested = MockNested(mock_contextmanager)

   320         with mock_nested:

   321             self.assertInWithManagerInvariants(mock_contextmanager)

   322             self.assertInWithManagerInvariants(mock_nested)

   323         self.assertAfterWithManagerInvariantsNoError(mock_contextmanager)

   324         self.assertAfterWithManagerInvariantsNoError(mock_nested)

   325 

   326     def testMultipleArgUnbound(self):

   327         m = mock_contextmanager_generator()

   328         n = mock_contextmanager_generator()

   329         o = mock_contextmanager_generator()

   330         mock_nested = MockNested(m, n, o)

   331         with mock_nested:

   332             self.assertInWithManagerInvariants(m)

   333             self.assertInWithManagerInvariants(n)

   334             self.assertInWithManagerInvariants(o)

   335             self.assertInWithManagerInvariants(mock_nested)

   336         self.assertAfterWithManagerInvariantsNoError(m)

   337         self.assertAfterWithManagerInvariantsNoError(n)

   338         self.assertAfterWithManagerInvariantsNoError(o)

   339         self.assertAfterWithManagerInvariantsNoError(mock_nested)

   340 

   341     def testMultipleArgBound(self):

   342         mock_nested = MockNested(mock_contextmanager_generator(),

   343             mock_contextmanager_generator(), mock_contextmanager_generator())

   344         with mock_nested as (m, n, o):

   345             self.assertInWithGeneratorInvariants(m)

   346             self.assertInWithGeneratorInvariants(n)

   347             self.assertInWithGeneratorInvariants(o)

   348             self.assertInWithManagerInvariants(mock_nested)

   349         self.assertAfterWithGeneratorInvariantsNoError(m)

   350         self.assertAfterWithGeneratorInvariantsNoError(n)

   351         self.assertAfterWithGeneratorInvariantsNoError(o)

   352         self.assertAfterWithManagerInvariantsNoError(mock_nested)

   353 

   354 

   355 class ExceptionalTestCase(unittest.TestCase, ContextmanagerAssertionMixin):

   356     def testSingleResource(self):

   357         cm = mock_contextmanager_generator()

   358         def shouldThrow():

   359             with cm as self.resource:

   360                 self.assertInWithManagerInvariants(cm)

   361                 self.assertInWithGeneratorInvariants(self.resource)

   362                 self.raiseTestException()

   363         self.assertRaises(RuntimeError, shouldThrow)

   364         self.assertAfterWithManagerInvariantsWithError(cm)

   365         self.assertAfterWithGeneratorInvariantsWithError(self.resource)

   366 

   367     def testNestedSingleStatements(self):

   368         mock_a = mock_contextmanager_generator()

   369         mock_b = mock_contextmanager_generator()

   370         def shouldThrow():

   371             with mock_a as self.foo:

   372                 with mock_b as self.bar:

   373                     self.assertInWithManagerInvariants(mock_a)

   374                     self.assertInWithManagerInvariants(mock_b)

   375                     self.assertInWithGeneratorInvariants(self.foo)

   376                     self.assertInWithGeneratorInvariants(self.bar)

   377                     self.raiseTestException()

   378         self.assertRaises(RuntimeError, shouldThrow)

   379         self.assertAfterWithManagerInvariantsWithError(mock_a)

   380         self.assertAfterWithManagerInvariantsWithError(mock_b)

   381         self.assertAfterWithGeneratorInvariantsWithError(self.foo)

   382         self.assertAfterWithGeneratorInvariantsWithError(self.bar)

   383 

   384     def testMultipleResourcesInSingleStatement(self):

   385         cm_a = mock_contextmanager_generator()

   386         cm_b = mock_contextmanager_generator()

   387         mock_nested = MockNested(cm_a, cm_b)

   388         def shouldThrow():

   389             with mock_nested as (self.resource_a, self.resource_b):

   390                 self.assertInWithManagerInvariants(cm_a)

   391                 self.assertInWithManagerInvariants(cm_b)

   392                 self.assertInWithManagerInvariants(mock_nested)

   393                 self.assertInWithGeneratorInvariants(self.resource_a)

   394                 self.assertInWithGeneratorInvariants(self.resource_b)

   395                 self.raiseTestException()

   396         self.assertRaises(RuntimeError, shouldThrow)

   397         self.assertAfterWithManagerInvariantsWithError(cm_a)

   398         self.assertAfterWithManagerInvariantsWithError(cm_b)

   399         self.assertAfterWithManagerInvariantsWithError(mock_nested)

   400         self.assertAfterWithGeneratorInvariantsWithError(self.resource_a)

   401         self.assertAfterWithGeneratorInvariantsWithError(self.resource_b)

   402 

   403     def testNestedExceptionBeforeInnerStatement(self):

   404         mock_a = mock_contextmanager_generator()

   405         mock_b = mock_contextmanager_generator()

   406         self.bar = None

   407         def shouldThrow():

   408             with mock_a as self.foo:

   409                 self.assertInWithManagerInvariants(mock_a)

   410                 self.assertInWithGeneratorInvariants(self.foo)

   411                 self.raiseTestException()

   412                 with mock_b as self.bar:

   413                     pass

   414         self.assertRaises(RuntimeError, shouldThrow)

   415         self.assertAfterWithManagerInvariantsWithError(mock_a)

   416         self.assertAfterWithGeneratorInvariantsWithError(self.foo)

   417 

   418         # The inner statement stuff should never have been touched

   419         self.assertEqual(self.bar, None)

   420         self.assertFalse(mock_b.enter_called)

   421         self.assertFalse(mock_b.exit_called)

   422         self.assertEqual(mock_b.exit_args, None)

   423 

   424     def testNestedExceptionAfterInnerStatement(self):

   425         mock_a = mock_contextmanager_generator()

   426         mock_b = mock_contextmanager_generator()

   427         def shouldThrow():

   428             with mock_a as self.foo:

   429                 with mock_b as self.bar:

   430                     self.assertInWithManagerInvariants(mock_a)

   431                     self.assertInWithManagerInvariants(mock_b)

   432                     self.assertInWithGeneratorInvariants(self.foo)

   433                     self.assertInWithGeneratorInvariants(self.bar)

   434                 self.raiseTestException()

   435         self.assertRaises(RuntimeError, shouldThrow)

   436         self.assertAfterWithManagerInvariantsWithError(mock_a)

   437         self.assertAfterWithManagerInvariantsNoError(mock_b)

   438         self.assertAfterWithGeneratorInvariantsWithError(self.foo)

   439         self.assertAfterWithGeneratorInvariantsNoError(self.bar)

   440 

   441     def testRaisedStopIteration1(self):

   442         # From bug 1462485

   443         @contextmanager

   444         def cm():

   445             yield

   446 

   447         def shouldThrow():

   448             with cm():

   449                 raise StopIteration("from with")

   450 

   451         self.assertRaises(StopIteration, shouldThrow)

   452 

   453     def testRaisedStopIteration2(self):

   454         # From bug 1462485

   455         class cm(object):

   456             def __enter__(self):

   457                 pass

   458             def __exit__(self, type, value, traceback):

   459                 pass

   460 

   461         def shouldThrow():

   462             with cm():

   463                 raise StopIteration("from with")

   464 

   465         self.assertRaises(StopIteration, shouldThrow)

   466 

   467     def testRaisedStopIteration3(self):

   468         # Another variant where the exception hasn't been instantiated

   469         # From bug 1705170

   470         @contextmanager

   471         def cm():

   472             yield

   473 

   474         def shouldThrow():

   475             with cm():

   476                 raise iter([]).next()

   477 

   478         self.assertRaises(StopIteration, shouldThrow)

   479 

   480     def testRaisedGeneratorExit1(self):

   481         # From bug 1462485

   482         @contextmanager

   483         def cm():

   484             yield

   485 

   486         def shouldThrow():

   487             with cm():

   488                 raise GeneratorExit("from with")

   489 

   490         self.assertRaises(GeneratorExit, shouldThrow)

   491 

   492     def testRaisedGeneratorExit2(self):

   493         # From bug 1462485

   494         class cm (object):

   495             def __enter__(self):

   496                 pass

   497             def __exit__(self, type, value, traceback):

   498                 pass

   499 

   500         def shouldThrow():

   501             with cm():

   502                 raise GeneratorExit("from with")

   503 

   504         self.assertRaises(GeneratorExit, shouldThrow)

   505 

   506 

   507 class NonLocalFlowControlTestCase(unittest.TestCase):

   508 

   509     def testWithBreak(self):

   510         counter = 0

   511         while True:

   512             counter += 1

   513             with mock_contextmanager_generator():

   514                 counter += 10

   515                 break

   516             counter += 100 # Not reached

   517         self.assertEqual(counter, 11)

   518 

   519     def testWithContinue(self):

   520         counter = 0

   521         while True:

   522             counter += 1

   523             if counter > 2:

   524                 break

   525             with mock_contextmanager_generator():

   526                 counter += 10

   527                 continue

   528             counter += 100 # Not reached

   529         self.assertEqual(counter, 12)

   530 

   531     def testWithReturn(self):

   532         def foo():

   533             counter = 0

   534             while True:

   535                 counter += 1

   536                 with mock_contextmanager_generator():

   537                     counter += 10

   538                     return counter

   539                 counter += 100 # Not reached

   540         self.assertEqual(foo(), 11)

   541 

   542     def testWithYield(self):

   543         def gen():

   544             with mock_contextmanager_generator():

   545                 yield 12

   546                 yield 13

   547         x = list(gen())

   548         self.assertEqual(x, [12, 13])

   549 

   550     def testWithRaise(self):

   551         counter = 0

   552         try:

   553             counter += 1

   554             with mock_contextmanager_generator():

   555                 counter += 10

   556                 raise RuntimeError

   557             counter += 100 # Not reached

   558         except RuntimeError:

   559             self.assertEqual(counter, 11)

   560         else:

   561             self.fail("Didn't raise RuntimeError")

   562 

   563 

   564 class AssignmentTargetTestCase(unittest.TestCase):

   565 

   566     def testSingleComplexTarget(self):

   567         targets = {1: [0, 1, 2]}

   568         with mock_contextmanager_generator() as targets[1][0]:

   569             self.assertEqual(targets.keys(), [1])

   570             self.assertEqual(targets[1][0].__class__, MockResource)

   571         with mock_contextmanager_generator() as targets.values()[0][1]:

   572             self.assertEqual(targets.keys(), [1])

   573             self.assertEqual(targets[1][1].__class__, MockResource)

   574         with mock_contextmanager_generator() as targets[2]:

   575             keys = targets.keys()

   576             keys.sort()

   577             self.assertEqual(keys, [1, 2])

   578         class C: pass

   579         blah = C()

   580         with mock_contextmanager_generator() as blah.foo:

   581             self.assertEqual(hasattr(blah, "foo"), True)

   582 

   583     def testMultipleComplexTargets(self):

   584         class C:

   585             def __enter__(self): return 1, 2, 3

   586             def __exit__(self, t, v, tb): pass

   587         targets = {1: [0, 1, 2]}

   588         with C() as (targets[1][0], targets[1][1], targets[1][2]):

   589             self.assertEqual(targets, {1: [1, 2, 3]})

   590         with C() as (targets.values()[0][2], targets.values()[0][1], targets.values()[0][0]):

   591             self.assertEqual(targets, {1: [3, 2, 1]})

   592         with C() as (targets[1], targets[2], targets[3]):

   593             self.assertEqual(targets, {1: 1, 2: 2, 3: 3})

   594         class B: pass

   595         blah = B()

   596         with C() as (blah.one, blah.two, blah.three):

   597             self.assertEqual(blah.one, 1)

   598             self.assertEqual(blah.two, 2)

   599             self.assertEqual(blah.three, 3)

   600 

   601 

   602 class ExitSwallowsExceptionTestCase(unittest.TestCase):

   603 

   604     def testExitTrueSwallowsException(self):

   605         class AfricanSwallow:

   606             def __enter__(self): pass

   607             def __exit__(self, t, v, tb): return True

   608         try:

   609             with AfricanSwallow():

   610                 1/0

   611         except ZeroDivisionError:

   612             self.fail("ZeroDivisionError should have been swallowed")

   613 

   614     def testExitFalseDoesntSwallowException(self):

   615         class EuropeanSwallow:

   616             def __enter__(self): pass

   617             def __exit__(self, t, v, tb): return False

   618         try:

   619             with EuropeanSwallow():

   620                 1/0

   621         except ZeroDivisionError:

   622             pass

   623         else:

   624             self.fail("ZeroDivisionError should have been raised")

   625 

   626 

   627 def test_main():

   628     run_unittest(FailureTestCase, NonexceptionalTestCase,

   629                  NestedNonexceptionalTestCase, ExceptionalTestCase,

   630                  NonLocalFlowControlTestCase,

   631                  AssignmentTargetTestCase,

   632                  ExitSwallowsExceptionTestCase)

   633 

   634 

   635 if __name__ == '__main__':

   636     test_main()