1 #!/usr/bin/env python

     2 

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

     4 

     5 from __future__ import with_statement

     6 

     7 __author__ = "Mike Bland"

     8 __email__ = "mbland at acm dot org"

     9 

    10 import sys

    11 import unittest

    12 from collections import deque

    13 from contextlib import GeneratorContextManager, contextmanager

    14 from test.test_support import run_unittest

    15 

    16 

    17 class MockContextManager(GeneratorContextManager):

    18     def __init__(self, gen):

    19         GeneratorContextManager.__init__(self, gen)

    20         self.enter_called = False

    21         self.exit_called = False

    22         self.exit_args = None

    23 

    24     def __enter__(self):

    25         self.enter_called = True

    26         return GeneratorContextManager.__enter__(self)

    27 

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

    29         self.exit_called = True

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

    31         return GeneratorContextManager.__exit__(self, type,

    32                                                 value, traceback)

    33 

    34 

    35 def mock_contextmanager(func):

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

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

    38     return helper

    39 

    40 

    41 class MockResource(object):

    42     def __init__(self):

    43         self.yielded = False

    44         self.stopped = False

    45 

    46 

    47 @mock_contextmanager

    48 def mock_contextmanager_generator():

    49     mock = MockResource()

    50     try:

    51         mock.yielded = True

    52         yield mock

    53     finally:

    54         mock.stopped = True

    55 

    56 

    57 class Nested(object):

    58 

    59     def __init__(self, *managers):

    60         self.managers = managers

    61         self.entered = None

    62 

    63     def __enter__(self):

    64         if self.entered is not None:

    65             raise RuntimeError("Context is not reentrant")

    66         self.entered = deque()

    67         vars = []

    68         try:

    69             for mgr in self.managers:

    70                 vars.append(mgr.__enter__())

    71                 self.entered.appendleft(mgr)

    72         except:

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

    74                 raise

    75         return vars

    76 

    77     def __exit__(self, *exc_info):

    78         # Behave like nested with statements

    79         # first in, last out

    80         # New exceptions override old ones

    81         ex = exc_info

    82         for mgr in self.entered:

    83             try:

    84                 if mgr.__exit__(*ex):

    85                     ex = (None, None, None)

    86             except:

    87                 ex = sys.exc_info()

    88         self.entered = None

    89         if ex is not exc_info:

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

    91 

    92 

    93 class MockNested(Nested):

    94     def __init__(self, *managers):

    95         Nested.__init__(self, *managers)

    96         self.enter_called = False

    97         self.exit_called = False

    98         self.exit_args = None

    99 

   100     def __enter__(self):

   101         self.enter_called = True

   102         return Nested.__enter__(self)

   103 

   104     def __exit__(self, *exc_info):

   105         self.exit_called = True

   106         self.exit_args = exc_info

   107         return Nested.__exit__(self, *exc_info)

   108 

   109 

   110 class FailureTestCase(unittest.TestCase):

   111     def testNameError(self):

   112         def fooNotDeclared():

   113             with foo: pass

   114         self.assertRaises(NameError, fooNotDeclared)

   115 

   116     def testEnterAttributeError(self):

   117         class LacksEnter(object):

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

   119                 pass

   120 

   121         def fooLacksEnter():

   122             foo = LacksEnter()

   123             with foo: pass

   124         self.assertRaises(AttributeError, fooLacksEnter)

   125 

   126     def testExitAttributeError(self):

   127         class LacksExit(object):

   128             def __enter__(self):

   129                 pass

   130 

   131         def fooLacksExit():

   132             foo = LacksExit()

   133             with foo: pass

   134         self.assertRaises(AttributeError, fooLacksExit)

   135 

   136     def assertRaisesSyntaxError(self, codestr):

   137         def shouldRaiseSyntaxError(s):

   138             compile(s, '', 'single')

   139         self.assertRaises(SyntaxError, shouldRaiseSyntaxError, codestr)

   140 

   141     def testAssignmentToNoneError(self):

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

   143         self.assertRaisesSyntaxError(

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

   145             '  pass')

   146 

   147     def testAssignmentToEmptyTupleError(self):

   148         self.assertRaisesSyntaxError(

   149             'with mock as ():\n'

   150             '  pass')

   151 

   152     def testAssignmentToTupleOnlyContainingNoneError(self):

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

   154         self.assertRaisesSyntaxError(

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

   156             '  pass')

   157 

   158     def testAssignmentToTupleContainingNoneError(self):

   159         self.assertRaisesSyntaxError(

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

   161             '  pass')

   162 

   163     def testEnterThrows(self):

   164         class EnterThrows(object):

   165             def __enter__(self):

   166                 raise RuntimeError("Enter threw")

   167             def __exit__(self, *args):

   168                 pass

   169 

   170         def shouldThrow():

   171             ct = EnterThrows()

   172             self.foo = None

   173             with ct as self.foo:

   174                 pass

   175         self.assertRaises(RuntimeError, shouldThrow)

   176         self.assertEqual(self.foo, None)

   177 

   178     def testExitThrows(self):

   179         class ExitThrows(object):

   180             def __enter__(self):

   181                 return

   182             def __exit__(self, *args):

   183                 raise RuntimeError(42)

   184         def shouldThrow():

   185             with ExitThrows():

   186                 pass

   187         self.assertRaises(RuntimeError, shouldThrow)

   188 

   189 class ContextmanagerAssertionMixin(object):

   190     TEST_EXCEPTION = RuntimeError("test exception")

   191 

   192     def assertInWithManagerInvariants(self, mock_manager):

   193         self.assertTrue(mock_manager.enter_called)

   194         self.assertFalse(mock_manager.exit_called)

   195         self.assertEqual(mock_manager.exit_args, None)

   196 

   197     def assertAfterWithManagerInvariants(self, mock_manager, exit_args):

   198         self.assertTrue(mock_manager.enter_called)

   199         self.assertTrue(mock_manager.exit_called)

   200         self.assertEqual(mock_manager.exit_args, exit_args)

   201 

   202     def assertAfterWithManagerInvariantsNoError(self, mock_manager):

   203         self.assertAfterWithManagerInvariants(mock_manager,

   204             (None, None, None))

   205 

   206     def assertInWithGeneratorInvariants(self, mock_generator):

   207         self.assertTrue(mock_generator.yielded)

   208         self.assertFalse(mock_generator.stopped)

   209 

   210     def assertAfterWithGeneratorInvariantsNoError(self, mock_generator):

   211         self.assertTrue(mock_generator.yielded)

   212         self.assertTrue(mock_generator.stopped)

   213 

   214     def raiseTestException(self):

   215         raise self.TEST_EXCEPTION

   216 

   217     def assertAfterWithManagerInvariantsWithError(self, mock_manager):

   218         self.assertTrue(mock_manager.enter_called)

   219         self.assertTrue(mock_manager.exit_called)

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

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

   222 

   223     def assertAfterWithGeneratorInvariantsWithError(self, mock_generator):

   224         self.assertTrue(mock_generator.yielded)

   225         self.assertTrue(mock_generator.stopped)

   226 

   227 

   228 class NonexceptionalTestCase(unittest.TestCase, ContextmanagerAssertionMixin):

   229     def testInlineGeneratorSyntax(self):

   230         with mock_contextmanager_generator():

   231             pass

   232 

   233     def testUnboundGenerator(self):

   234         mock = mock_contextmanager_generator()

   235         with mock:

   236             pass

   237         self.assertAfterWithManagerInvariantsNoError(mock)

   238 

   239     def testInlineGeneratorBoundSyntax(self):

   240         with mock_contextmanager_generator() as foo:

   241             self.assertInWithGeneratorInvariants(foo)

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

   243         self.assertAfterWithGeneratorInvariantsNoError(foo)

   244 

   245     def testInlineGeneratorBoundToExistingVariable(self):

   246         foo = None

   247         with mock_contextmanager_generator() as foo:

   248             self.assertInWithGeneratorInvariants(foo)

   249         self.assertAfterWithGeneratorInvariantsNoError(foo)

   250 

   251     def testInlineGeneratorBoundToDottedVariable(self):

   252         with mock_contextmanager_generator() as self.foo:

   253             self.assertInWithGeneratorInvariants(self.foo)

   254         self.assertAfterWithGeneratorInvariantsNoError(self.foo)

   255 

   256     def testBoundGenerator(self):

   257         mock = mock_contextmanager_generator()

   258         with mock as foo:

   259             self.assertInWithGeneratorInvariants(foo)

   260             self.assertInWithManagerInvariants(mock)

   261         self.assertAfterWithGeneratorInvariantsNoError(foo)

   262         self.assertAfterWithManagerInvariantsNoError(mock)

   263 

   264     def testNestedSingleStatements(self):

   265         mock_a = mock_contextmanager_generator()

   266         with mock_a as foo:

   267             mock_b = mock_contextmanager_generator()

   268             with mock_b as bar:

   269                 self.assertInWithManagerInvariants(mock_a)

   270                 self.assertInWithManagerInvariants(mock_b)

   271                 self.assertInWithGeneratorInvariants(foo)

   272                 self.assertInWithGeneratorInvariants(bar)

   273             self.assertAfterWithManagerInvariantsNoError(mock_b)

   274             self.assertAfterWithGeneratorInvariantsNoError(bar)

   275             self.assertInWithManagerInvariants(mock_a)

   276             self.assertInWithGeneratorInvariants(foo)

   277         self.assertAfterWithManagerInvariantsNoError(mock_a)

   278         self.assertAfterWithGeneratorInvariantsNoError(foo)

   279 

   280 

   281 class NestedNonexceptionalTestCase(unittest.TestCase,

   282     ContextmanagerAssertionMixin):

   283     def testSingleArgInlineGeneratorSyntax(self):

   284         with Nested(mock_contextmanager_generator()):

   285             pass

   286 

   287     def testSingleArgUnbound(self):

   288         mock_contextmanager = mock_contextmanager_generator()

   289         mock_nested = MockNested(mock_contextmanager)

   290         with mock_nested:

   291             self.assertInWithManagerInvariants(mock_contextmanager)

   292             self.assertInWithManagerInvariants(mock_nested)

   293         self.assertAfterWithManagerInvariantsNoError(mock_contextmanager)

   294         self.assertAfterWithManagerInvariantsNoError(mock_nested)

   295 

   296     def testSingleArgBoundToNonTuple(self):

   297         m = mock_contextmanager_generator()

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

   299         # assigned to foo.

   300         with Nested(m) as foo:

   301             self.assertInWithManagerInvariants(m)

   302         self.assertAfterWithManagerInvariantsNoError(m)

   303 

   304     def testSingleArgBoundToSingleElementParenthesizedList(self):

   305         m = mock_contextmanager_generator()

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

   307         # assigned to foo.

   308         with Nested(m) as (foo):

   309             self.assertInWithManagerInvariants(m)

   310         self.assertAfterWithManagerInvariantsNoError(m)

   311 

   312     def testSingleArgBoundToMultipleElementTupleError(self):

   313         def shouldThrowValueError():

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

   315                 pass

   316         self.assertRaises(ValueError, shouldThrowValueError)

   317 

   318     def testSingleArgUnbound(self):

   319         mock_contextmanager = mock_contextmanager_generator()

   320         mock_nested = MockNested(mock_contextmanager)

   321         with mock_nested:

   322             self.assertInWithManagerInvariants(mock_contextmanager)

   323             self.assertInWithManagerInvariants(mock_nested)

   324         self.assertAfterWithManagerInvariantsNoError(mock_contextmanager)

   325         self.assertAfterWithManagerInvariantsNoError(mock_nested)

   326 

   327     def testMultipleArgUnbound(self):

   328         m = mock_contextmanager_generator()

   329         n = mock_contextmanager_generator()

   330         o = mock_contextmanager_generator()

   331         mock_nested = MockNested(m, n, o)

   332         with mock_nested:

   333             self.assertInWithManagerInvariants(m)

   334             self.assertInWithManagerInvariants(n)

   335             self.assertInWithManagerInvariants(o)

   336             self.assertInWithManagerInvariants(mock_nested)

   337         self.assertAfterWithManagerInvariantsNoError(m)

   338         self.assertAfterWithManagerInvariantsNoError(n)

   339         self.assertAfterWithManagerInvariantsNoError(o)

   340         self.assertAfterWithManagerInvariantsNoError(mock_nested)

   341 

   342     def testMultipleArgBound(self):

   343         mock_nested = MockNested(mock_contextmanager_generator(),

   344             mock_contextmanager_generator(), mock_contextmanager_generator())

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

   346             self.assertInWithGeneratorInvariants(m)

   347             self.assertInWithGeneratorInvariants(n)

   348             self.assertInWithGeneratorInvariants(o)

   349             self.assertInWithManagerInvariants(mock_nested)

   350         self.assertAfterWithGeneratorInvariantsNoError(m)

   351         self.assertAfterWithGeneratorInvariantsNoError(n)

   352         self.assertAfterWithGeneratorInvariantsNoError(o)

   353         self.assertAfterWithManagerInvariantsNoError(mock_nested)

   354 

   355 

   356 class ExceptionalTestCase(unittest.TestCase, ContextmanagerAssertionMixin):

   357     def testSingleResource(self):

   358         cm = mock_contextmanager_generator()

   359         def shouldThrow():

   360             with cm as self.resource:

   361                 self.assertInWithManagerInvariants(cm)

   362                 self.assertInWithGeneratorInvariants(self.resource)

   363                 self.raiseTestException()

   364         self.assertRaises(RuntimeError, shouldThrow)

   365         self.assertAfterWithManagerInvariantsWithError(cm)

   366         self.assertAfterWithGeneratorInvariantsWithError(self.resource)

   367 

   368     def testNestedSingleStatements(self):

   369         mock_a = mock_contextmanager_generator()

   370         mock_b = mock_contextmanager_generator()

   371         def shouldThrow():

   372             with mock_a as self.foo:

   373                 with mock_b as self.bar:

   374                     self.assertInWithManagerInvariants(mock_a)

   375                     self.assertInWithManagerInvariants(mock_b)

   376                     self.assertInWithGeneratorInvariants(self.foo)

   377                     self.assertInWithGeneratorInvariants(self.bar)

   378                     self.raiseTestException()

   379         self.assertRaises(RuntimeError, shouldThrow)

   380         self.assertAfterWithManagerInvariantsWithError(mock_a)

   381         self.assertAfterWithManagerInvariantsWithError(mock_b)

   382         self.assertAfterWithGeneratorInvariantsWithError(self.foo)

   383         self.assertAfterWithGeneratorInvariantsWithError(self.bar)

   384 

   385     def testMultipleResourcesInSingleStatement(self):

   386         cm_a = mock_contextmanager_generator()

   387         cm_b = mock_contextmanager_generator()

   388         mock_nested = MockNested(cm_a, cm_b)

   389         def shouldThrow():

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

   391                 self.assertInWithManagerInvariants(cm_a)

   392                 self.assertInWithManagerInvariants(cm_b)

   393                 self.assertInWithManagerInvariants(mock_nested)

   394                 self.assertInWithGeneratorInvariants(self.resource_a)

   395                 self.assertInWithGeneratorInvariants(self.resource_b)

   396                 self.raiseTestException()

   397         self.assertRaises(RuntimeError, shouldThrow)

   398         self.assertAfterWithManagerInvariantsWithError(cm_a)

   399         self.assertAfterWithManagerInvariantsWithError(cm_b)

   400         self.assertAfterWithManagerInvariantsWithError(mock_nested)

   401         self.assertAfterWithGeneratorInvariantsWithError(self.resource_a)

   402         self.assertAfterWithGeneratorInvariantsWithError(self.resource_b)

   403 

   404     def testNestedExceptionBeforeInnerStatement(self):

   405         mock_a = mock_contextmanager_generator()

   406         mock_b = mock_contextmanager_generator()

   407         self.bar = None

   408         def shouldThrow():

   409             with mock_a as self.foo:

   410                 self.assertInWithManagerInvariants(mock_a)

   411                 self.assertInWithGeneratorInvariants(self.foo)

   412                 self.raiseTestException()

   413                 with mock_b as self.bar:

   414                     pass

   415         self.assertRaises(RuntimeError, shouldThrow)

   416         self.assertAfterWithManagerInvariantsWithError(mock_a)

   417         self.assertAfterWithGeneratorInvariantsWithError(self.foo)

   418 

   419         # The inner statement stuff should never have been touched

   420         self.assertEqual(self.bar, None)

   421         self.assertFalse(mock_b.enter_called)

   422         self.assertFalse(mock_b.exit_called)

   423         self.assertEqual(mock_b.exit_args, None)

   424 

   425     def testNestedExceptionAfterInnerStatement(self):

   426         mock_a = mock_contextmanager_generator()

   427         mock_b = mock_contextmanager_generator()

   428         def shouldThrow():

   429             with mock_a as self.foo:

   430                 with mock_b as self.bar:

   431                     self.assertInWithManagerInvariants(mock_a)

   432                     self.assertInWithManagerInvariants(mock_b)

   433                     self.assertInWithGeneratorInvariants(self.foo)

   434                     self.assertInWithGeneratorInvariants(self.bar)

   435                 self.raiseTestException()

   436         self.assertRaises(RuntimeError, shouldThrow)

   437         self.assertAfterWithManagerInvariantsWithError(mock_a)

   438         self.assertAfterWithManagerInvariantsNoError(mock_b)

   439         self.assertAfterWithGeneratorInvariantsWithError(self.foo)

   440         self.assertAfterWithGeneratorInvariantsNoError(self.bar)

   441 

   442     def testRaisedStopIteration1(self):

   443         # From bug 1462485

   444         @contextmanager

   445         def cm():

   446             yield

   447 

   448         def shouldThrow():

   449             with cm():

   450                 raise StopIteration("from with")

   451 

   452         self.assertRaises(StopIteration, shouldThrow)

   453 

   454     def testRaisedStopIteration2(self):

   455         # From bug 1462485

   456         class cm(object):

   457             def __enter__(self):

   458                 pass

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

   460                 pass

   461 

   462         def shouldThrow():

   463             with cm():

   464                 raise StopIteration("from with")

   465 

   466         self.assertRaises(StopIteration, shouldThrow)

   467 

   468     def testRaisedStopIteration3(self):

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

   470         # From bug 1705170

   471         @contextmanager

   472         def cm():

   473             yield

   474 

   475         def shouldThrow():

   476             with cm():

   477                 raise iter([]).next()

   478 

   479         self.assertRaises(StopIteration, shouldThrow)

   480 

   481     def testRaisedGeneratorExit1(self):

   482         # From bug 1462485

   483         @contextmanager

   484         def cm():

   485             yield

   486 

   487         def shouldThrow():

   488             with cm():

   489                 raise GeneratorExit("from with")

   490 

   491         self.assertRaises(GeneratorExit, shouldThrow)

   492 

   493     def testRaisedGeneratorExit2(self):

   494         # From bug 1462485

   495         class cm (object):

   496             def __enter__(self):

   497                 pass

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

   499                 pass

   500 

   501         def shouldThrow():

   502             with cm():

   503                 raise GeneratorExit("from with")

   504 

   505         self.assertRaises(GeneratorExit, shouldThrow)

   506 

   507 

   508 class NonLocalFlowControlTestCase(unittest.TestCase):

   509 

   510     def testWithBreak(self):

   511         counter = 0

   512         while True:

   513             counter += 1

   514             with mock_contextmanager_generator():

   515                 counter += 10

   516                 break

   517             counter += 100 # Not reached

   518         self.assertEqual(counter, 11)

   519 

   520     def testWithContinue(self):

   521         counter = 0

   522         while True:

   523             counter += 1

   524             if counter > 2:

   525                 break

   526             with mock_contextmanager_generator():

   527                 counter += 10

   528                 continue

   529             counter += 100 # Not reached

   530         self.assertEqual(counter, 12)

   531 

   532     def testWithReturn(self):

   533         def foo():

   534             counter = 0

   535             while True:

   536                 counter += 1

   537                 with mock_contextmanager_generator():

   538                     counter += 10

   539                     return counter

   540                 counter += 100 # Not reached

   541         self.assertEqual(foo(), 11)

   542 

   543     def testWithYield(self):

   544         def gen():

   545             with mock_contextmanager_generator():

   546                 yield 12

   547                 yield 13

   548         x = list(gen())

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

   550 

   551     def testWithRaise(self):

   552         counter = 0

   553         try:

   554             counter += 1

   555             with mock_contextmanager_generator():

   556                 counter += 10

   557                 raise RuntimeError

   558             counter += 100 # Not reached

   559         except RuntimeError:

   560             self.assertEqual(counter, 11)

   561         else:

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

   563 

   564 

   565 class AssignmentTargetTestCase(unittest.TestCase):

   566 

   567     def testSingleComplexTarget(self):

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

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

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

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

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

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

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

   575         with mock_contextmanager_generator() as targets[2]:

   576             keys = targets.keys()

   577             keys.sort()

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

   579         class C: pass

   580         blah = C()

   581         with mock_contextmanager_generator() as blah.foo:

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

   583 

   584     def testMultipleComplexTargets(self):

   585         class C:

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

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

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

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

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

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

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

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

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

   595         class B: pass

   596         blah = B()

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

   598             self.assertEqual(blah.one, 1)

   599             self.assertEqual(blah.two, 2)

   600             self.assertEqual(blah.three, 3)

   601 

   602 

   603 class ExitSwallowsExceptionTestCase(unittest.TestCase):

   604 

   605     def testExitTrueSwallowsException(self):

   606         class AfricanSwallow:

   607             def __enter__(self): pass

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

   609         try:

   610             with AfricanSwallow():

   611                 1/0

   612         except ZeroDivisionError:

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

   614 

   615     def testExitFalseDoesntSwallowException(self):

   616         class EuropeanSwallow:

   617             def __enter__(self): pass

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

   619         try:

   620             with EuropeanSwallow():

   621                 1/0

   622         except ZeroDivisionError:

   623             pass

   624         else:

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

   626 

   627 

   628 def test_main():

   629     run_unittest(FailureTestCase, NonexceptionalTestCase,

   630                  NestedNonexceptionalTestCase, ExceptionalTestCase,

   631                  NonLocalFlowControlTestCase,

   632                  AssignmentTargetTestCase,

   633                  ExitSwallowsExceptionTestCase)

   634 

   635 

   636 if __name__ == '__main__':

   637     test_main()