<commit_before>import os
import time as ttime
from .core import CallbackBase
from ..utils import ensure_uid
import numpy as np
import doct
import matplotlib.pyplot as plt


class LiveImage(CallbackBase):
    """
    Stream 2D images in a cross-section viewer.

    Parameters
    ----------
    field : string
        name of data field in an Event
    """

    def __init__(self, field, *, fs=None):
        from xray_vision.backend.mpl.cross_section_2d import CrossSection
        import matplotlib.pyplot as plt
        super().__init__()
        self.field = field
        fig = plt.figure()
        self.cs = CrossSection(fig)
        self.cs._fig.show()
        if fs is None:
            import filestore.api as fs
        self.fs = fs

    def event(self, doc):
        uid = doc['data'][self.field]
        if 'filled' not in doc.keys() or \
                        doc['filled'].get(self.field, False) is False:
            data = self.fs.retrieve(uid)
        self.update(data)
        super().event(doc)

    def update(self, data):
        self.cs.update_image(data)
        self.cs._fig.canvas.draw_idle()


def post_run(callback, db=None, fill=False):
    """Trigger a callback to process all the Documents from a run at the end.

    This function does not receive the Document stream during collection.
    It retrieves the complete set of Documents from the DataBroker after
    collection is complete.

    Parameters
    ----------
    callback : callable
        Expected signature ::

            def func(doc_name, doc):
                pass
<commit_msg>FIX: pull the data from the doc if not a uid<commit_after>import os
import time as ttime
from .core import CallbackBase
from ..utils import ensure_uid
import numpy as np
import doct
import matplotlib.pyplot as plt


class LiveImage(CallbackBase):
    """
    Stream 2D images in a cross-section viewer.

    Parameters
    ----------
    field : string
        name of data field in an Event
    """

    def __init__(self, field, *, fs=None):
        from xray_vision.backend.mpl.cross_section_2d import CrossSection
        import matplotlib.pyplot as plt
        super().__init__()
        self.field = field
        fig = plt.figure()
        self.cs = CrossSection(fig)
        self.cs._fig.show()
        if fs is None:
            import filestore.api as fs
        self.fs = fs

    def event(self, doc):
        if 'filled' not in doc.keys() or \
                        doc['filled'].get(self.field, False) is False:
            uid = doc['data'][self.field]
            data = self.fs.retrieve(uid)
        else:
            data = doc['data'][self.field]
        self.update(data)
        super().event(doc)

    def update(self, data):
        self.cs.update_image(data)
        self.cs._fig.canvas.draw_idle()


def post_run(callback, db=None, fill=False):
    """Trigger a callback to process all the Documents from a run at the end.

    This function does not receive the Document stream during collection.
    It retrieves the complete set of Documents from the DataBroker after
    collection is complete.

    Parameters
    ----------
    callback : callable
        Expected signature ::

            def func(doc_name, doc):
                pass

{'fraction_non_alphanumeric': 0.09550215650030808, 'fraction_numerical': 0.0012322858903265558, 'mean_word_length': 2.0806451612903225, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '744998', 'n_tokens_mistral': 926, 'n_tokens_neox': 900, 'n_words': 309}