Skip to content

🌊  Streaming Write

Submitted by Lukas Masuch

Summary

Drop-in replacement for st.write with streaming support.

Functions

write

Drop-in replacement for st.write with streaming support.

This function is a drop-in replacement for st.write that adds additional capabilities: * Supports streaming data via generator functions. * Executes callable objects (e.g. functions) and writes the return value.

Source code in src/streamlit_extras/streaming_write/__init__.py 
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
@extra
def write(*args: Any, unsafe_allow_html: bool = False, **kwargs) -> List[Any]:
    """Drop-in replacement for `st.write` with streaming support.

    This function is a drop-in replacement for `st.write` that adds additional capabilities:
    * Supports streaming data via generator functions.
    * Executes callable objects (e.g. functions) and writes the return value.
    """
    if not args:
        return []

    written_content: List[Any] = []
    string_buffer: List[str] = []

    def flush_buffer():
        if string_buffer:
            text_content = " ".join(string_buffer)
            text_container = st.empty()
            text_container.markdown(text_content)
            written_content.append(text_content)
            string_buffer[:] = []

    for arg in args:
        # Order matters!
        if isinstance(arg, str):
            string_buffer.append(arg)
        elif callable(arg) or inspect.isgenerator(arg):
            flush_buffer()
            if inspect.isgeneratorfunction(arg) or inspect.isgenerator(arg):
                # This causes greyed out effect since this element is missing on rerun:
                stream_container = None
                streamed_response = ""

                def flush_stream_response():
                    nonlocal streamed_response
                    nonlocal stream_container
                    if streamed_response and stream_container:
                        stream_container.write(
                            streamed_response,
                            unsafe_allow_html=unsafe_allow_html,
                            **kwargs,
                        )
                        written_content.append(streamed_response)
                        stream_container = None
                        streamed_response = ""

                generator = arg() if inspect.isgeneratorfunction(arg) else arg
                for chunk in generator:
                    if isinstance(chunk, str):
                        first_text = False
                        if not stream_container:
                            stream_container = st.empty()
                            first_text = True
                        streamed_response += chunk
                        # Only add the streaming symbol on the second text chunk
                        stream_container.write(
                            streamed_response + ("" if first_text else " β–Œ"),
                            unsafe_allow_html=unsafe_allow_html,
                            **kwargs,
                        )
                    elif callable(chunk):
                        flush_stream_response()
                        chunk()
                        written_content.append(chunk)
                    else:
                        flush_stream_response()
                        st.write(chunk, unsafe_allow_html=unsafe_allow_html, **kwargs)
                        written_content.append(chunk)
                flush_stream_response()

            else:
                return_value = arg()
                written_content.append(arg)
                if return_value is not None:
                    flush_buffer()
                    st.write(
                        return_value, unsafe_allow_html=unsafe_allow_html, **kwargs
                    )
        else:
            flush_buffer()
            st.write(arg, unsafe_allow_html=unsafe_allow_html, **kwargs)
            written_content.append(arg)
    flush_buffer()
    return written_content

Import:

from streamlit_extras.streaming_write import write # (1)!
  1. You should add this to the top of your .py file πŸ› 

Examples

example

def example():
    def stream_example():
        for word in _LOREM_IPSUM.split():
            yield word + " "
            time.sleep(0.1)

        # Also supports any other object supported by `st.write`
        yield pd.DataFrame(
            np.random.randn(5, 10),
            columns=["a", "b", "c", "d", "e", "f", "g", "h", "i", "j"],
        )

        for word in _LOREM_IPSUM.split():
            yield word + " "
            time.sleep(0.05)

    if st.button("Stream data"):
        write(stream_example)