import io
import pickle
import numpy as np
from copy import copy, deepcopy
from contextlib import contextmanager
import matplotlib.pyplot as plt


@contextmanager
def copy_figure(fig: plt.Figure):
    buf = io.BytesIO()
    pickle.dump(fig, buf)
    buf.seek(0)
    fig_copy = pickle.load(buf)
    # fig_copy = copy(fig)  # you get the same error when using copy instead
    # fig_copy = deepcopy(fig)  # this results in an "NotImplementedError: 
    #    TransformNode instances can not be copied. Consider using frozen() 
    #    instead." which I have no idea what that means either.
    yield fig_copy
    plt.close(fig_copy)


def example():
    image = np.random.uniform(0, 1, (100, 100))
    fig, ax = plt.subplots()
    image_artist = ax.imshow(image)
    fig.colorbar(image_artist)
    with copy_figure(fig) as copy_fig:
        copy_fig.show()

    fig.show()


example()