✨ improve registering implementation

Author: jvdd

plotly_resampler/__init__.py

@@ -2,7 +2,7 @@
 
 from .aggregation import LTTB, EfficientLTTB, EveryNthPoint
 from .figure_resampler import FigureResampler, FigureWidgetResampler
-from .module import register_plotly_resampler, unregister_plotly_resampler
+from .registering import register_plotly_resampler, unregister_plotly_resampler
 
 __docformat__ = "numpy"
 __author__ = "Jonas Van Der Donckt, Jeroen Van Der Donckt, Emiel Deprost"

plotly_resampler/figure_resampler/figure_resampler.py

@@ -48,7 +48,7 @@ def __init__(
         else:
             # Create a new figure object and make sure that the trace uid will not get
             # adjusted when they are added.
-            f = go.Figure()
+            f = self._get_figure_class(go.Figure)()
             f._data_validator.set_uid = False
 
             if isinstance(figure, BaseFigure):  # go.FigureWidget or AbstractFigureAggregator
@@ -92,6 +92,12 @@ def __init__(
         self._port: int | None = None
         self._host: str | None = None
 
+    # @staticmethod
+    # def _get_figure_class() -> type:
+    #     """Return the class of the underlying figure."""
+    #     from ..module import get_plotly_constr
+    #     return get_plotly_constr(go.Figure)
+
     def show_dash(
         self,
         mode=None,

plotly_resampler/figure_resampler/figure_resampler_interface.py

@@ -96,6 +96,9 @@ def __init__(
 
         self._global_downsampler = default_downsampler
 
+        # Given figure should always be a BaseFigure (that is not wrapped by a plotly-resampler class)
+        assert isinstance(figure, BaseFigure)
+        assert not issubclass(type(figure), AbstractFigureAggregator)
         self._figure_class = figure.__class__
 
         if convert_existing_traces:
@@ -422,6 +425,29 @@ def _check_update_figure_dict(
                 updated_trace_indices.append(idx)
         return updated_trace_indices
 
+    @staticmethod
+    def _get_figure_class(constr: type) -> type:
+        """Get the plotly figure class (constructor) for the given class (constructor).
+
+        .. Note::
+            This method will always return a plotly constructor, even when the given
+            `constr` is decorated (after executing the ``register_plotly_resampler`` 
+            function).
+
+        Parameters
+        ----------
+        constr: type
+            The constructor class for which we want to retrieve the plotly constructor.
+
+        Returns
+        -------
+        type:
+            The plotly figure class (constructor) of the given `constr`.
+
+        """
+        from ..registering import get_plotly_constr  # To avoid ImportError
+        return get_plotly_constr(constr)
+
     @staticmethod
     def _slice_time(
         hf_series: pd.Series,

plotly_resampler/figure_resampler/figurewidget_resampler.py

@@ -52,7 +52,7 @@ def __init__(
         verbose: bool = False,
     ):
         # Parse the figure input before calling `super`
-        f = go.FigureWidget()
+        f = self._get_figure_class(go.FigureWidget)()
         f._data_validator.set_uid = False
 
         if isinstance(figure, BaseFigure):  # go.Figure or go.FigureWidget or AbstractFigureAggregator
@@ -244,6 +244,12 @@ def _update_spike_ranges(self, layout, *showspikes, force_update=False):
             self._relayout_hist.append(["showspikes", "initial call or showspikes"])
             self._relayout_hist.append("-" * 40)
 
+    # @staticmethod
+    # def _get_figure_class() -> type:
+    #     """Return the class of the underlying figure."""
+    #     from ..module import get_plotly_constr
+    #     return get_plotly_constr(go.FigureWidget)
+
     def reset_axes(self):
         """Reset the axes of the FigureWidgetResampler.
 

plotly_resampler/module.py

@@ -0,0 +1,117 @@
+__author__ = "Jeroen Van Der Donckt, Jonas Van Der Donckt, Emiel Deprost"
+
+from plotly_resampler import FigureResampler, FigureWidgetResampler
+from plotly_resampler.figure_resampler.figure_resampler_interface import (
+    AbstractFigureAggregator,
+)
+from functools import wraps
+
+import plotly
+
+WRAPPED_PREFIX = "[Plotly-Resampler]__"
+PLOTLY_MODULES = [
+    plotly.graph_objs,
+    plotly.graph_objects,
+]  # wait for this PR https://github.com/plotly/plotly.py/pull/3779
+PLOTLY_CONSTRUCTOR_WRAPPER = {
+    "Figure": FigureResampler,
+    "FigureWidget": FigureWidgetResampler,
+}
+
+
+def _already_wrapped(constr):
+    return constr.__name__.startswith(WRAPPED_PREFIX)
+
+
+def get_plotly_constr(constr):
+    if _already_wrapped(constr):
+        return constr.__wrapped__  # get the original constructor
+    return constr
+
+
+### Registering the wrappers
+
+
+def _is_ipython_env():
+    """Check if we are in an IPython environment (with a kernel)."""
+    try:
+        from IPython import get_ipython
+
+        return "IPKernelApp" in get_ipython().config
+    except Exception:
+        return False
+
+
+def _register_wrapper(
+    module: type,
+    constr_name: str,
+    pr_class: AbstractFigureAggregator,
+    **aggregator_kwargs,
+):
+    constr = getattr(module, constr_name)
+    constr = get_plotly_constr(constr)  # get the original plotly constructor
+
+    # print(f"Wrapping {constr_name} with {pr_class}")
+
+    @wraps(constr)
+    def wrapped_constr(*args, **kwargs):
+        # print(f"Executing constructor wrapper for {constr_name}", constr)
+        return pr_class(constr(*args, **kwargs), **aggregator_kwargs)
+
+    wrapped_constr.__name__ = WRAPPED_PREFIX + constr_name
+    setattr(module, constr_name, wrapped_constr)
+
+
+def register_plotly_resampler(mode="auto", **aggregator_kwargs):  # TODO: show kwargs (e.g., port)?
+    """Register plotly-resampler to plotly.graph_objects.
+
+    This function results in the use of plotly-resampler under the hood.
+
+    .. Note::
+        We advise to
+
+    Parameters
+    ----------
+    mode : str, optional
+        The mode of the plotly-resampler.
+        Possible values are: 'auto', 'figure', 'widget', None.
+        If 'auto' is used, the mode is determined based on the environment; if it is in
+        an ipython environment, the mode is 'widget', otherwise it is 'figure'.
+        If 'figure' is used, all plotly figures are wrapped as FigureResampler objects.
+        If 'widget' is used, all plotly figure widgets are wrapped as 
+        FigureWidgetResampler objects (we advise to use this mode in ipython environment
+        with a kernel).
+        If None is used, wrapping is done as expected (go.Figure -> FigureResampler,
+        go.FigureWidget -> FigureWidgetResampler).
+    aggregator_kwargs : dict, optional
+        The keyword arguments to pass to the plotly-resampler decorator its constructor.
+        See more details in :class:`FigureResampler <FigureResampler>` and 
+        :class:`FigureWidgetResampler <FigureWidgetResampler>`.
+
+    """
+    for constr_name, pr_class in PLOTLY_CONSTRUCTOR_WRAPPER.items():
+        if (mode == "auto" and _is_ipython_env()) or mode == "widget":
+            pr_class = FigureWidgetResampler
+        elif mode == "figure":
+            pr_class = FigureResampler
+        # else: default mode -> wrap according to PLOTLY_CONSTRUCTOR_WRAPPER
+
+        for module in PLOTLY_MODULES:
+            _register_wrapper(module, constr_name, pr_class, **aggregator_kwargs)
+
+
+### Unregistering the wrappers
+
+
+def _unregister_wrapper(module: type, constr_name: str):
+    constr = getattr(module, constr_name)
+    if _already_wrapped(constr):
+        constr = constr.__wrapped__
+        setattr(module, constr_name, constr)
+
+
+def unregister_plotly_resampler():
+    """Unregister plotly-resampler from plotly.graph_objects."""
+    for constr in PLOTLY_CONSTRUCTOR_WRAPPER.keys():
+        for module in PLOTLY_MODULES:
+            _unregister_wrapper(module, constr)