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

Question

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?

Answer 1

Although the with block doesn't exist as a separate scope, you can obtain the with block by tokenizing the caller's source starting from the line number of the caller's frame, and keeping track of the depth of the indentation, where the body of the block always starts from the first INDENT token, and after all the nested INDENT and DEDENT tokens are accounted for, ends with a DEDENT token that decreases the indentation depth back to 0.

With the body of the with statement identified, it's then a simple matter of enclosing the body in a dummy block (such as if 1:), compiling it and executing it with the caller frame's global and local variables for as many times as specified.

Note that with this approach changes made to local variables in the block will not be reflected after the block, since the exec function cannot modify the frame's locals dictionary per the documentation.

import sys
from linecache import getline
from tokenize import tokenize, INDENT, DEDENT

class again:
    def __init__(self, times=1):
        self.times = times

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        if exc_type:
            def readline():
                lineno = caller.f_lineno
                while line := getline(filename, lineno):
                    if lineno == caller.f_lineno: # dedent the with statement
                        line = line.lstrip()      # so it can be parsed alone
                    yield line.encode()
                    lineno += 1
                yield b''
            caller = sys._getframe(1)
            filename = caller.f_code.co_filename
            first = end = depth = 0
            try:
                for token, _, (start, _), (end, _), _ in tokenize(readline().__next__):
                    if token == INDENT:
                        depth += 1
                        if not first:
                            first = start
                    elif token == DEDENT:
                        if depth == 1:
                            break
                        depth -= 1
            except IndentationError:
                end += 1
            body = ''.join(
                getline(filename, caller.f_lineno + lineno - 1)
                for lineno in range(first, end)
            )
            code = compile('if 1:\n' + body, '\n' + body, 'exec')
            namespace = caller.f_locals.copy()
            for _ in range(self.times):
                try:
                    exec(code, caller.f_globals, namespace)
                    return True
                except Exception as e:
                    pass

so that:

def f():
    x = 0
    with again(2):
        print(x)
        x += 1
        if x < 3:
             raise Exception(f'I hate {x}')
f()

outputs:

0
1
2

Demo here

Answer 2

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).