Plotly Resampler Compatibility with fig.layout.uirevision

[benlw22]

I am wondering if there is a way to preserve the figure state (zoom window) when changing the figure data attributes. To be specific, I have the following working example that uses the uirevision attribute in the figure layout so that when the user changes the color of a trace, the figure does not completely reset the window. This works as expected even with a FigureResampler figure, but when the callback for dynamically updating the figure based on the user zooming in or out is included, the uirevision functionality fails.

from dash import html, dcc, no_update, Input, Output, State #, Patch
import pandas as pd
import plotly.graph_objects as go
from plotly_resampler import FigureResampler
from dash_extensions.enrich import DashProxy, Serverside, ServersideOutputTransform

df = pd.read_csv('https://raw.githubusercontent.com/plotly/datasets/master/stockdata.csv')

df['reference'] = df[df.columns[0]]

app = DashProxy(__name__, transforms=[ServersideOutputTransform()])
app.scripts.config.serve_locally = True
app.css.config.serve_locally = True

app.layout = html.Div([
    html.Label('Color'),
    dcc.Dropdown(
        id='color',
        options=[
            {'label': 'Navy', 'value': '#001f3f'},
            {'label': 'Blue', 'value': '#0074D9'},
            {'label': 'Aqua', 'value': '#7FDBFF'},
            {'label': 'TEAL', 'value': '#39CCCC'},
            {'label': 'OLIVE', 'value': '#3D9970'},
            {'label': 'GREEN', 'value': '#2ECC40'},
            {'label': 'LIME', 'value': '#01FF70'},
            {'label': 'YELLOW', 'value': '#FFDC00'},
            {'label': 'ORANGE', 'value': '#FF851B'},
            {'label': 'RED', 'value': '#FF4136'},
            {'label': 'MAROON', 'value': '#85144b'},
            {'label': 'FUCHSIA', 'value': '#F012BE'},
            {'label': 'PURPLE', 'value': '#B10DC9'},
        ],
        value='#001f3f'
    ),

    html.Label('Reference'),
    dcc.RadioItems(
        id='reference',
        options=[{'label': i, 'value': i} for i in ['Display', 'Hide']],
        value='Display'
    ),

    html.Label('Dataset'),
    dcc.Dropdown(
        id='dataset',
        options=[{'label': i, 'value': i} for i in df.columns],
        value=df.columns[0]
    ),

    dcc.Graph(id='graph'),
    dcc.Store(id="store"),
])

@app.callback(
    Output('graph', 'figure'),
    Output("store", "data"),
    [Input('color', 'value'),
     Input('reference', 'value'),
     Input('dataset', 'value')])
def display_graph(color, reference, dataset):
    # initialize a figure resampler figure
    fig: FigureResampler = FigureResampler(
            go.Figure()
        )
    fig.add_trace(go.Scattergl(mode= 'lines', name= dataset, marker = {'color': color}), hf_x=df.index, hf_y=df[dataset])
    if reference == 'Display':
        fig.add_trace(go.Scattergl(mode= 'lines', name= 'Reference'), hf_x=df.index, hf_y=df['reference'])
    fig.update_layout(
                uirevision = dataset,
                legend = {'x': 0, 'y': 1}
                )
    return fig, Serverside(fig)

# The plotly-resampler callback to update the graph after a relayout event (= zoom/pan)
# As we use the figure again as output, we need to set: allow_duplicate=True
# including this callback breaks the uirevision functionality
# comment out this callback and the uirevision functionality works as expected
@app.callback(
    Output("graph", "figure", allow_duplicate=True),
    Input("graph", "relayoutData"),
    State("store", "data"),  # The server side cached FigureResampler per session
    State('dataset', 'value'),
    prevent_initial_call=True,
    memoize=True,
)
def update_fig(relayoutdata: dict, fig: FigureResampler, dataset):
    if fig is None:
        return no_update
    print(fig.layout.uirevision)
    return fig.construct_update_data_patch(relayoutdata)

if __name__ == '__main__':
    app.run_server(debug=True)

I use plotly-resampler all the time and really appreciate your tremendous work! The functionality described above would really enhance my user experience and I am sure there must be a way to accomplish this that I am not considering. Thanks for your help.