A contextmanager that force the fixed width image output in jupyter notebook

gistfile1.txt

from contextlib import contextmanager
import matplotlib.pyplot as plt
import base64
from io import BytesIO
from IPython.display import HTML, display

@contextmanager
def dp(width=600, dpi=300):
    def display_figure():
        # Save the current figure as a PNG image in memory
        png_image = BytesIO()
        plt.gcf().savefig(png_image, format='png', dpi=dpi)

        # Rewind to the start of the file
        png_image.seek(0)

        # Create a base64 encoding of the image
        png_image_b64_string = "data:image/png;base64,"
        png_image_b64_string += base64.b64encode(png_image.read()).decode('utf8')

        # Close the figure
        plt.close(plt.gcf())

        # Display the image in an HTML img tag with a specific width
        display(HTML(f'<img src="{png_image_b64_string}" width="{width}"/>'))

    try:
        yield display_figure
    finally:
        pass
    
# with fixed_size_plot(800) as display_plot:
#     plt.figure(figsize=(10,6))
#     plt.plot([0, 1, 2, 3, 4], [0, 1, 4, 9, 16])
#     display_plot()