treetrace.ext.pyplot

 1from ..tracing.data.blob import DataWithMime
 2
 3
 4def capture_figure(figure=None, file_format: str = "png", **kwargs):
 5    import io
 6
 7    import matplotlib.pyplot as plt
 8
 9    buf = io.BytesIO()
10    plt.savefig(buf, format=file_format, **kwargs)
11    buf.seek(0)
12    if figure:
13        plt.close(figure)
14    if file_format == "png":
15        mime = "image/png"
16    else:
17        mime = None
18    return DataWithMime(buf.read(), mime_type=mime)