# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/02_display.ipynb.

# %% auto 0
__all__ = ['PassthroughCapturer']

# %% ../nbs/02_display.ipynb 2
from IPython.core.displaypub import DisplayPublisher
from IPython.core.display_functions import _new_id
from IPython.display import display, clear_output, Javascript
from IPython import get_ipython
from typing import List
from PIL import Image
from collections import defaultdict
import warnings
import sys

# %% ../nbs/02_display.ipynb 3
class PassthroughCapturer(DisplayPublisher):
    """A DisplayPublisher that stores AND displays!"""

    raw_outputs: List

    def __init__(self):
        self._publisher = None
        self.raw_outputs = []

    def publish(
        self, data, metadata=None, source=None, *, transient=None, update=False
    ):
        self.raw_outputs.append(
            {
                "data": data,
                "metadata": metadata,
                "transient": transient,
                "update": update,
            }
        )
        self._publisher.publish(data, metadata=metadata, transient=transient, update=update)

    def clear_output(self, wait=False):
        self.raw_outputs = []
        self._publisher.clear_output(wait)

    def start(self):
        ip = get_ipython()
        self._publisher = ip.display_pub
        ip.display_pub = self

        # print("starting passthrough: ", id(self._publisher), id(self))

    # Note: stop must be called from the cell that started the
    # passthrough, or from the associated post-cell callback
    def stop(self):
        ip = get_ipython()
        assert self._publisher
        ip.display_pub = self._publisher

    def get_outputs(self, with_js=False):
        coalesced_outputs = []
        display_id_map = defaultdict(list)

        outputs = self.raw_outputs
        if not with_js:
            outputs = [o for o in outputs if 'application/javascript' not in o['data']]

        # dicts are ordered in Python 3.7+
        for output in outputs:
            display_id = output['transient']['display_id'] if output['transient'] else _new_id()
            if not output['update']:
                display_id_map[display_id] += [output["data"]]
            else:
                display_id_map[display_id][-1] = output["data"]

        for display_id, output_list in display_id_map.items():
            coalesced_outputs += output_list

        return coalesced_outputs
