mlukac/xrf-explorer-dev

1from typing import Any, Callable, Union2 3import pandas as pd4import param5from bokeh.models import ColumnDataSource6 7from xrf_explorer.core.events import DatasetUpdated, EventWithForceableUpdate8from xrf_explorer.helpers.print import event_print, time_me9 10TData = Union[pd.DataFrame, dict[str, Any]]11 12 13class ColumnDataSourceWrapper(param.Parameterized):14 15    output = param.ClassSelector(class_=ColumnDataSource, default=None, allow_refs=True)16 17    def __init__(self, getter: Callable[[], TData | None], **kwargs) -> None:18        super().__init__(**kwargs)19        self._getter = getter20        self._last_data: TData | None = None21 22        self._time_me_meta: str | None = kwargs["name"] if "name" in kwargs else None23 24    @time_me()25    def handle(self, event: DatasetUpdated) -> None:26        self._handle(event, self._update)27 28    @time_me()29    def force_handle(self, event: DatasetUpdated) -> None:30        self._handle(event, self._force_update)31 32    def _handle(33        self, event: DatasetUpdated, updater: Callable[[DatasetUpdated, TData], None]34    ) -> None:35        data = self._getter()36 37        if data is None:38            return39 40        if self.output is None:41            self.output = ColumnDataSource(data)42        else:43            updater(event, data)44 45        self._last_data = data46 47    def _force_update(self, event: DatasetUpdated, data: TData) -> None:48        self.output.data = data49 50    def _update(self, event: DatasetUpdated, data: TData) -> None:51        if self._should_skip_update(event, data):52            event_print("   SKIPPING: CDS update")53            return54 55        self.output.data = data56 57    def _should_skip_update(self, event: DatasetUpdated, data: TData) -> bool:58        if self._should_force_update(event):59            return False60        return self._last_data is not None and data is self._last_data61 62    def _should_force_update(self, event: DatasetUpdated) -> bool:63        return isinstance(event, EventWithForceableUpdate) and event.force_update64