Inspect python with block

user_99024 0
Viewed 174

Is it possible to get the source code of the with block?

E.g.:

with print_source():
  someFunction()
  someOtherFunction()

# "someFunction()\nsomeOtherFunction()\n" should be printed

If not, is it something that is available in Python C API?

I've tried using sys.settrace, but it only catches function calls, which with-block isn't.

python
1 Answers

It is not super straightforward to do this because a with block isn't an object in its own right like a module or function. However, while I don't have a full implementation, I have some suggestions from which you should be able to put something together:

  • In the __enter__() function of the print_source context manager, use inspect.stack() to get the current frames on the stack. The top frame is the __enter__ method, and the next one down is the function where the with block lives. I will call this the "function frame".
  • The function frame info includes the line number of the with print_source line. You will need this shortly.
  • Pass the frame from the function frame info to inspect.getsourcelines() to get all the source lines from the function (like it says on the tin). This function also identifies the line number of the def foo(): line.
  • Using the line number of the with ... line, offset by the line number of the def ... line, you can identify the indentation level of the with block. Iterate through the source lines of the function, starting with the one after the with block, until you hit another line with the same (or smaller) indentation as the with ... line. You have now identified all the source within the with block.

Things that I have not allowed for (but should be possible) include:

  • comment lines within the with block that have smaller indentation than the rest of the block (should be easy to just check if a line is a comment and ignore it for the end-of-block check). You probably also want to handle the case where a with block is the last code in a file, but is followed by comments that have smaller indentation. Do something like tracking the line number of the last code that had the correct indentation.
  • multiple context managers being combined, like with cm1(), cm2(), cm3():. I haven't checked how this appears in the stack
  • having the with statement spread across multiple lines. I haven't checked which line the frame reports the statement as being on, but I have a strong suspicion that it is the line with the :

Hope that gives enough to get something working!

user_11983210
Related