from .context import (
    GalaxySeleniumContextImpl,
    init as c_init,
)


def init(config=None):
    return c_init(config=config, clazz=JupyterContextImpl)


class JupyterContextImpl(GalaxySeleniumContextImpl):
    def screenshot(self, label):
        path = super().screenshot(label)
        from IPython.display import Image

        return Image(filename=path)
