# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/04_utils.cell_cbs.ipynb.

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

# %% ../../nbs/04_utils.cell_cbs.ipynb 2
import sys
from IPython import get_ipython
from . import PassthroughCapturer, TeeIO, SaveTraceback

# %% ../../nbs/04_utils.cell_cbs.ipynb 4
class ExecCBs:
    stdout = None
    stderr = None
    result = None # The thing at the end of the cell
    display =None
    saved_streams = None

    def __init__(self, after_run=None, once=True):
        self.do_after_run = after_run
        self.once = once

    def pre_callback(self, info=None):
        # print("# pre")
        ip = get_ipython()
        # This will capture diplay() and friends
        self.display = PassthroughCapturer()
        self.display.start()

        # This will capture stdout/stderr
        sys.stdout = self.stdout = TeeIO(sys.stdout)
        sys.stderr = self.stderr = TeeIO(sys.stderr)


        self.tracebacks = SaveTraceback(ip._showtraceback)
        self.tracebacks.register()

        # We register the callback here, so it does not run after the cell
        # that registered the pre

        ip.events.register("post_run_cell", self.post_callback)
        # We get the result in the post callback info, no need to register anything else.

    def post_callback(self, result):
        # Post will be called in the cell that registers the
        # callbacks, but it wuld not have ran pre.
        # print("# post")
        self.result = result

        if self.stdout:
            sys.stdout = self.stdout._stream
            sys.stderr = self.stderr._stream

        if self.display:
            self.display.stop()

        if self.tracebacks:
            self.tracebacks.unregister()

        if self.once:
            self.unregister()

        if self.do_after_run:
            self.do_after_run(
                self.stdout.getvalue(),
                self.stderr.getvalue(),
                self.display.get_outputs(),
                result,
                self.tracebacks.tracebacks
            )

    def register(self):
        ip = get_ipython()
        ip.events.register("pre_run_cell", self.pre_callback)


    def unregister(self):
        try:
            ip = get_ipython()
            ip.events.unregister("pre_run_cell", self.pre_callback)
            ip.events.unregister("post_run_cell", self.post_callback)
        except: pass # In case unregister was aclled before a cell ran, and we never regisstered pre.

