Search code examples
pythoncontextmanager

Is it possible to access the context object (code block) inside the __exit__() method of a context manager?


I would like to invoke the code object again in the exit() method if it raises an exception (maybe several times, maybe with delay). I know it is very easy to do with a decorator, but my motivation is that sometimes I want to repeat just some fragment of code that I don't want to extract to a separate function and decorate it. I'm looking for something along these lines:

class again(object):
    def __enter__(self):
        pass

    def __exit__(self, exc_type, exc_val, exc_tb):
        if exc_type is not None:
            ????        # Invoke the code object again
            return True # eat exception

It would used like so:

x = 0
with again():
    print x
    x += 1
    if x == 1:
         raise Exception('I hate 1')

and the expected output would be:

0
1  

I could find a way to get hold of the code object. None of the context manager attributes seem to reference it (I guess it is not really needed, because it's job is just to do stuff before and after).

Is it possible to do it?


Solution

  • The with block doesn't exist as a separate code object, so no. See this similar question. In that case, the questioner was trying to do the reverse (access the context manager from inside the code block), but as this answer explains, the with block is not a separate scope, so it doesn't really have any separate status.

    You can see this with an example:

    import contextlib
    import dis
    
    @contextlib.contextmanager
    def silly():
        yield
    
    def foo():
        print "Hello"
        with silly():
            print "Inside"
        print "Goodbye"
    

    and then

    >>> dis.dis(foo.__code__)
      2           0 LOAD_CONST               1 (u'Hello')
                  3 PRINT_ITEM          
                  4 PRINT_NEWLINE       
    
      3           5 LOAD_GLOBAL              0 (silly)
                  8 CALL_FUNCTION            0
                 11 SETUP_WITH              10 (to 24)
                 14 POP_TOP             
    
      4          15 LOAD_CONST               2 (u'Inside')
                 18 PRINT_ITEM          
                 19 PRINT_NEWLINE       
                 20 POP_BLOCK           
                 21 LOAD_CONST               0 (None)
            >>   24 WITH_CLEANUP        
                 25 END_FINALLY         
    
      5          26 LOAD_CONST               3 (u'Goodbye')
                 29 PRINT_ITEM          
                 30 PRINT_NEWLINE       
                 31 LOAD_CONST               0 (None)
                 34 RETURN_VALUE  
    

    You can see that the with block's code is just inside the function's code object along with everything else. It doesn't exist as a separate code object and isn't distinguished from the rest of the function's code. You can't get it out in any sane way (by which I mean, without hacking the bytecode).