This st.echo() context manager really looks like black magic to me!
I understand what it does and how to use it. But I struggle to understand how such a thing can be achieved in python : a context manager that somewhat captures the raw string of python code it delimits (while it’s being executed)… O_O
Does it call a function that analyses the whole script syntacticaly to get the relevant code piece and pass it back to the context manager so that it knows what to display ?
I’m really curious to learn how this can be implemented.
Understanding it might also shed some light on why this widget is resisting the StreamPy interactive implementation I’m working on.
Found my way to the source code of st.echo(). Pretty clever !
I now understand why it doesn’t work in StreamPy (the source code is not a file anymore)
Will have to modify it to fit the StreamPy interactive environment.
Great question! What’s cool about Streamlit is that it’s open-source. To view the source code for st.echo, click the [source] button to the right of its function signature in the docs.
Does it call a function that analyses the whole script syntacticaly to get the relevant code piece and pass it back to the context manager so that it knows what to display ?
Your intuition is correct, and yeah, it does feel like black magic! Here’s a high-level overview of how it works, followed by a walkthrough of the source code:
TL;DR
The st.echo command utilizes Python’s introspection abilities and the Abstract Syntax Tree (AST) to capture the source code of the with block where st.echo is called. It uses the traceback module to capture the current stack frame to identify the context of the call to st.echo, and reads and parses the entire Python file where st.echo was called.
After parsing the Python file into an AST, it creates a mapping from line numbers to AST nodes, determines the starting and ending lines of the with st.echo(): block, and extracts the corresponding lines of source code. After executing the with block, it then displays the extracted source code in the app.
Source code walkthrough for st.echo
As you’ve identified, st.echo is a context manager, which means it has special behavior when used in a with statement.
When the st.echo context manager is entered (but before the block of code within it is executed), st.echo captures the stack frame of where st.echo was invoked in your script by using traceback.extract_stack()[-3]. Here’s how the traceback indexing works:
traceback.extract_stack()[-1] refers to the current location inside the st.echo function
traceback.extract_stack()[-2] is the location of the with st.echo(): call in the user’s code
traceback.extract_stack()[-3] points to the location in your code where st.echo was invoked
It then reads the Python source file where st.echo was called, using the filename from the stack frame. This gives us a list of all lines of source code in that file.
The function then uses the astmodule to parse the entire Python file into a tree of nodes, each representing a Python construct (like a function definition, a loop, a conditional, an expression, etc).
And then constructs a mapping from line numbers to nodes (line_to_node_map) by traversing the AST and collecting all nodes that have a body attribute (which includes function definitions, class definitions, loops, and other constructs that contain blocks of code).
The starting line of the block to echo is determined from the body[0].lineno of the node corresponding to the line where st.echo was called, and the ending line is the end_lineno of the same node. We then extract the lines of source code for the block to echo from the list of all source lines, dedent the block to remove leading whitespace, and store the resulting string.
The yield statement then allows the code inside the with block to be executed. After the with block is executed, st.echo finally displays the captured source code in the app.
So you’re right, st.echo is implemented by capturing the Python source code, parsing it into an AST, and then analyzing the AST to find the block of code to display.
Thanks for the thorough answer. It does help a lot !
This approach will turn very useful elsewhere (not necessarily streamlit related) to capture raw python code at runtime.
I will have to tweak it for Streampy, because if the source code containing the with statement is not directly in the .py file (but rather a string of code sent to some kind of exec command, such as code inputted in an interactive console’s input cell) the current implementation will fail.
For instance, this following simple script throws an exception:
import streamlit as st
code=r"""
with st.echo():
st.write("Something")
"""
exec(code,globals())
will return a filename equal to <string> in case the code being executed is not a file… And apparently, there’s no possibility to retrieve the raw code from within the executed script.
In my case, I found a way to work around this limitation, as I know the code to analyze since it comes from the input cell, I just pass this code to the ast module instead of reading it from a file.
It works, you can test it in the StreamPy app with the same syntax as usual.
Here’s a high-level overview of how it works, followed by a walkthrough of the source code:
