Интерактивная легенда Ipython в виджете не отвечает при обновлении

Я пытаюсь реализовать виджет в блокноте jupyter для временной шкалы датчика (B), где пользователь должен иметь возможность выбрать, какой фрукт он хотел бы видеть (A). Он также должен иметь возможность выбрать период, который он хотел бы проверить. Основываясь на следующем коде, который я нашел, я попытался включить виджет.

• [https://stackoverflow.com/questions/31410043/hiding-lines-after-showing-a-pyplot-figure][1]

Интерактивный график работает с добавленным кодом при первом построении за весь период данных. Но каждый раз, когда я обновляю сюжет, интерактивная легенда больше не отвечает. Когда я нажимаю легенды, ничего не происходит.

Я не могу найти правильный путь, я пробовал разные способы обновления сюжета или его очистки, но всегда одна и та же проблема.

Интерактивная сюжетная часть Джо Кингстона

import pandas as pd
import seaborn as sns

import ipywidgets as widgets
from ipywidgets import AppLayout, Button, Layout
%matplotlib widget
import matplotlib.pyplot as plt
import matplotlib.dates as mdates
from matplotlib.dates import DateFormatter

# generate dummy data
dti = pd.date_range("2018-01-01", periods=10, freq="H")
d = {"A":["Apple","Banana"] * 5, "B":np.random.randint(0,10,size=(10))}
df = pd.DataFrame(d, index = dti)
df_pump = df

# from Joe Kington, Stackoverflow
def main():
    #prepare labels for plotting
    labels = df_pump["A"].unique()

    # if one of the time pickers is empty, use the whole dataset
    if start_date_picker.value or end_date_picker.value == None:
        df_plot = df_pump
    # else use user entered dates
    else:
        df_plot = df_pump.loc[start_date_picker.value:end_date_picker.value]

    # make a new line for each container of the plant, e.g. HYD1, HYD2 etc.
    fig, ax = plt.subplots(figsize=(9,6))
    for i in range(0, len(labels)):
        sns.scatterplot(x = df_plot.loc[df_plot["A"] == labels[i], "B"].index, 
                   y = df_plot.loc[df_plot["A"] == labels[i], "B"], 
                   label=labels[i],
                   )

    ax.legend(loc='upper left', bbox_to_anchor=(1.05, 1),
              ncol=2, borderaxespad=0)
    fig.subplots_adjust(right=0.55)
    fig.suptitle('pH Verlauf pro Behälter \n' +
                '\n Right-click to hide all\nMiddle-click to show all',
                va='top', size='large')
    plt.ylabel("pH")
    plt.xlabel("")
    ax.xaxis.set_major_formatter(DateFormatter('%d %b %y'),) # maybee use dateformatter

    plt.tight_layout()

    leg = interactive_legend()
    return fig, ax, leg

def interactive_legend(ax=None):
    if ax is None:
        ax = plt.gca()
    if ax.legend_ is None:
        ax.legend()

    return InteractiveLegend(ax.get_legend())

class InteractiveLegend(object):
    def __init__(self, legend):
        self.legend = legend
        self.fig = legend.axes.figure

        self.lookup_artist, self.lookup_handle = self._build_lookups(legend)
        self._setup_connections()

        self.update()

    def _setup_connections(self):
        for artist in self.legend.texts + self.legend.legendHandles:
            artist.set_picker(10) # 10 points tolerance

        self.fig.canvas.mpl_connect('pick_event', self.on_pick)
        self.fig.canvas.mpl_connect('button_press_event', self.on_click)

    def _build_lookups(self, legend):
        labels = [t.get_text() for t in legend.texts]
        handles = legend.legendHandles
        label2handle = dict(zip(labels, handles))
        handle2text = dict(zip(handles, legend.texts))

        lookup_artist = {}
        lookup_handle = {}
        for artist in legend.axes.get_children():
            if artist.get_label() in labels:
                handle = label2handle[artist.get_label()]
                lookup_handle[artist] = handle
                lookup_artist[handle] = artist
                lookup_artist[handle2text[handle]] = artist

        lookup_handle.update(zip(handles, handles))
        lookup_handle.update(zip(legend.texts, handles))

        return lookup_artist, lookup_handle

    def on_pick(self, event):
        handle = event.artist
        if handle in self.lookup_artist:

            artist = self.lookup_artist[handle]
            artist.set_visible(not artist.get_visible())
            self.update()

    def on_click(self, event):
        if event.button == 3:
            visible = False
        elif event.button == 2:
            visible = True
        else:
            return

        for artist in self.lookup_artist.values():
            artist.set_visible(visible)
        self.update()

    def update(self):
        for artist in self.lookup_artist.values():
            handle = self.lookup_handle[artist]
            if artist.get_visible():
                handle.set_visible(True)
            else:
                handle.set_visible(False)
        self.fig.canvas.draw()

    def show(self):
        plt.show()

Часть виджета:

def update_plot(change):
    #clear_output() tried with clearing, no success

    fig, ax, leg = main()
    plt.show()
    
    
### open widget
output = widgets.Output()

with output:
    fig, ax, leg = main()
    

# date pickers
start_date_picker = widgets.DatePicker(
    description='Starting date',
    disabled=False
)
end_date_picker = widgets.DatePicker(
    description='End date',
    disabled=False
)

# plot update and displays
b = widgets.Button(
    value=False,
    description='Time active',

    button_style='', # 'success', 'info', 'warning', 'danger' or ''
    tooltip='Click this button after choosing the desired time period',
    icon='chevron-circle-right' # (FontAwesome names without the `fa-` prefix)
)

# control links
b.on_click(update_plot)

# app layout
controls = widgets.HBox([start_date_picker, 
                         end_date_picker,
                        b])

AppLayout(header=None,
          left_sidebar=None,
          center=output,
          right_sidebar=None,
          footer=controls)