Is there any way to get source code inside context manager as string?

Question

Source code of function can be received with inspect.getsourcelines(func) function. Is there any way to do same for context manager?

with test():
    print('123')

# How to get "print('123')" as line here?

Asked By: Mikhail Gerasimov

||

Source

Answer 1

What do you think about this solution?

import traceback


class ContextManagerContent(object):

    def __enter__(self):
        return

    def __exit__(self, _type, value, _traceback):
        stack = traceback.extract_stack()
        f, last_line = self._get_origin_info(stack)

        with open(f) as fin:
            lines = list(fin)

        search = 'with {cls_name}'.format(cls_name=self.__class__.__name__)
        for i, x in enumerate(lines[:last_line + 1][::-1]):
            if search in x:
                first_line = len(lines) - i
                break

        selected_lines = lines[first_line:last_line + 1]
        print ''.join(selected_lines)

    def _get_origin_info(self, stack):
        origin = None
        for i, x in enumerate(stack[::-1]):
            if x[2] == '__exit__':
                origin = stack[::-1][i + 1]
                break

        return origin[0], origin[1] - 1


with ContextManagerContent():
    print '123'
    print '456'
    print '789'

If you save this in a .py file and run it you can see being printed the numbers 123, 456 and 789, after that you can see the block of the context manager.

Notice that I did not handle possibile exceptions or formatting of the output, and some parts can be improved, but I think that this is a good starting point.

Answered By: se7entyse7en

Answer 2

Here is another solution inspired from @se7entyse7en’s answer. I think it is a bit cleaner and more efficient this way

from inspect import currentframe, getframeinfo
from contextlib import contextmanager


@contextmanager
def debug(string):
    # before block
    cf = currentframe()
    first_line =  cf.f_back.f_back.f_lineno
    filename = getframeinfo(cf.f_back.f_back).filename

    yield

    # after block
    cf = currentframe()
    last_line =  cf.f_back.f_back.f_lineno
    with open(filename) as f:
        lines = f.readlines()[first_line:last_line]
    print(string + 'n' + ''.join(lines).rstrip())


if __name__ == '__main__':
    with debug("show this code in stdout:"):
        a = 1
        b = 2
        a, b = b, a

Answered By: cwilmot

Is there any way to get source code inside context manager as string?

Question:

Answers: