Download Plotly figures as PDF with download_button

Hi!

I’m trying to figure out the [best] way to download my plotly figures saved in session state with the download button.

I do a similar thing with saved data frames for which I have the following code:

@st.cache
def save_df(df: pandas.DataFrame):
    return df.to_csv(sep='\t', index=False).encode('utf-8')

# some app code here

if not st.session_state.processed_data and not st.session_state.final_feature_table:
    st.warning('Process data!')

else:
    files = {
        "Feature dataframe": st.session_state.final_feature_table[0],
        "Processed dataframe": st.session_state.processed_data[0]
    }

    available_files = [key for (key, value) in files.items()]
    file_name = st.text_input('', value='your_filename')
    selected_file = st.selectbox('Select a file for download:', available_files, index=0)
    file = save_df(files[selected_file])

    st.download_button(
        label="Download table",
        data=file,
        file_name=f'{file_name}.tab',
        mime='text/csv'
    )

This works like a charm.

Now, I generate figures using a form, resetting the state after each re-submission. I check and it works outside of the form using st.plotly_chart.

The problem occurs when I try to save a static figure to a PDF file using the following code:

@st.cache
def save_figure(figure: go.Figure, file):
    return figure.write_image(file, format='pdf')


# this is outside of the form

    if not st.session_state.figures:
        st.warning('Generate figures!')
    else:
        st.markdown('## Download Figures')
        available_figures = [key for (key, value) in st.session_state.figures.items()]
        file_name = st.text_input('', value='your_filename')
        selected_figure = st.selectbox('Select a figure for download:', available_figures, index=0)
        figure_object = st.session_state.figures[selected_figure]  # I can see this figure with 
        save_figure = save_figure(figure_object, file_name)
        btn = st.download_button(
            label="Download image",
            data=save_figure,
            file_name=f"{file_name}.pdf",
            mime="figure/pdf")

I get the this error:

RuntimeError: Invalid binary data format: <class 'NoneType'>

I’m guessing it’s something trivial that I’m missing but I just cannot figure out what it is!

Thanks,
Simon

Hi @simonb :wave:

The docs for figure.write_image states that the function returns None. To get this working, let’s save the generated pdf of the figure in an in-memory buffer using io.BytesIO(), and download the buffer contents with st.download_button. Note: the write_image function also requires the installation of the kaleido package.

Here’s a working reproducible example:

Solution

import streamlit as st
import pandas as pd
import plotly.graph_objects as go
import io

# Load the data
@st.experimental_memo
def load_data():
    return pd.DataFrame(
        {
            "Fruit": ["Apples", "Oranges", "Bananas", "Apples", "Oranges", "Bananas"],
            "Contestant": ["Alex", "Alex", "Alex", "Jordan", "Jordan", "Jordan"],
            "Number Eaten": [2, 1, 3, 1, 3, 2],
        }
    )

# Create and cache a Plotly figure
@st.experimental_memo
def create_figure(df):
    fig = go.Figure()
    for contestant, group in df.groupby("Contestant"):
        fig.add_trace(
            go.Bar(
                x=group["Fruit"],
                y=group["Number Eaten"],
                name=contestant,
                hovertemplate="Contestant=%s<br>Fruit=%%{x}<br>Number Eaten=%%{y}<extra></extra>"
                % contestant,
            )
        )
    fig.update_layout(legend_title_text="Contestant")
    fig.update_xaxes(title_text="Fruit")
    fig.update_yaxes(title_text="Number Eaten")
    return fig

df = load_data()
fig = create_figure(df)

# Create an in-memory buffer
buffer = io.BytesIO()

# Save the figure as a pdf to the buffer
fig.write_image(file=buffer, format="pdf")

# Download the pdf from the buffer
st.download_button(
    label="Download PDF",
    data=buffer,
    file_name="figure.pdf",
    mime="application/pdf",
)

st.plotly_chart(fig)

Output

plotly-download-pdf

Happy Streamlit-ing! :balloon:
Snehan

1 Like

Thank you @snehankekre, app’s working beautifully now! :vulcan_salute:t2:

Simon

1 Like