From 71bdd2fdb00ccd3ffbee218e67f74f47426824a4 Mon Sep 17 00:00:00 2001
From: OE <olivier.ernst@edf.fr>
Date: Fri, 3 May 2024 10:34:25 +0200
Subject: [PATCH] WIP

---
 dash_tooltip/__init__.py | 25 +++++++++++++++++--------
 dash_tooltip/utils.py    | 28 ++++++++++++++++------------
 2 files changed, 33 insertions(+), 20 deletions(-)

diff --git a/dash_tooltip/__init__.py b/dash_tooltip/__init__.py
index db2bd28..30ad627 100644
--- a/dash_tooltip/__init__.py
+++ b/dash_tooltip/__init__.py
@@ -5,7 +5,8 @@
 import plotly.graph_objs as go
 from dash import Input, Output, State, dash
 
-from .config import DEFAULT_ANNOTATION_CONFIG
+from .config import DEFAULT_ANNOTATION_CONFIG, DEFAULT_TEMPLATE
+from .custom_figure import CustomFigure
 from .utils import _display_click_data, _find_all_graph_ids, add_annotation_store
 
 # Logger setup
@@ -29,8 +30,6 @@
 # Now, you can log messages
 logger.debug("dash_tooltip log active")
 
-DEFAULT_TEMPLATE = "x: %{x},<br>y: %{y}"
-
 registered_callbacks = set()
 
 
@@ -80,6 +79,7 @@ def tooltip(
     for graph_id in graph_ids:
         callback_identifier = (graph_id, "figure")
         if callback_identifier in registered_callbacks:
+            # update figure here
             continue  # Skip reattaching if already registered
 
         registered_callbacks.add(callback_identifier)
@@ -93,11 +93,20 @@ def tooltip(
             State(component_id=graph_id, component_property="figure"),
         )
         def display_click_data(
-            clickData: Dict[str, Any], figure: go.Figure
-        ) -> go.Figure:
-            return _display_click_data(
-                clickData, figure, app, template, style, apply_log_fix, debug
-            )
+            clickData: Dict[str, Any], figure: Union[CustomFigure, Dict[str, Any]]
+        ) -> CustomFigure:
+            if figure is None:
+                # Initialize the figure
+                figure = CustomFigure()
+            if isinstance(figure, CustomFigure):
+                return _display_click_data(
+                    clickData, figure, app, style, template, apply_log_fix, debug
+                )
+            else:
+                figure = CustomFigure(figure)
+                return _display_click_data(
+                    clickData, figure, app, style, template, apply_log_fix, debug
+                )
 
         dbg_str = "console.log(relayoutData);"
 
diff --git a/dash_tooltip/utils.py b/dash_tooltip/utils.py
index 303ebaa..8bd4468 100644
--- a/dash_tooltip/utils.py
+++ b/dash_tooltip/utils.py
@@ -9,7 +9,7 @@
 from dash import dcc
 from dash.html import Div
 
-from dash_tooltip import DEFAULT_ANNOTATION_CONFIG
+from dash_tooltip import DEFAULT_ANNOTATION_CONFIG, CustomFigure
 
 logger = logging.getLogger("dash_tooltip")
 
@@ -152,21 +152,17 @@ def get_axis_type(fig: go.Figure, axis: str) -> str:
 
 def _display_click_data(
     clickData: Dict[str, Any],
-    figure: Union[go.Figure, Dict[str, Any]],  # Allow both go.Figure and dictionary
+    figure: Union[CustomFigure, Dict[str, Any]],  # Allow both go.Figure and dictionary
     app: dash.Dash,
-    template: str,
     config: Dict[Any, Any],
+    template: str,
     apply_log_fix: bool = True,
     debug: bool = False,
-) -> go.Figure:
+) -> CustomFigure:
     """Displays the tooltip on the graph when a data point is clicked."""
 
     xaxis, yaxis = "x", "y"  # Default values
 
-    if figure is None:
-        # Initialize the figure
-        figure = go.Figure()
-
     # Check if figure is a dictionary
     if isinstance(figure, dict):
         # Extract data and layout from the figure dictionary
@@ -181,9 +177,11 @@ def _display_click_data(
             data.append(trace_class(**trace))
 
         # Construct the go.Figure using data and layout
-        fig = go.Figure(data=data, layout=layout)
+        fig = CustomFigure(data=data, layout=layout)
+        fig.update_template(template)
     else:
         fig = figure
+        fig.update_template(template)
 
     merged_config = deep_merge_dicts(DEFAULT_ANNOTATION_CONFIG.copy(), config)
 
@@ -245,7 +243,7 @@ def _display_click_data(
                 ),
             )
 
-        placeholders = re.findall(r"%{(.*?)}", template)
+        placeholders = re.findall(r"%{(.*?)}", fig.layout._tooltip_template)
 
         template_data = {}
         for placeholder in placeholders:
@@ -269,12 +267,18 @@ def _display_click_data(
                 else:
                     template_data[placeholder] = str(value)
 
+        tooltip_template = fig.layout._tooltip_template
         for placeholder, value in template_data.items():
-            template = template.replace(f"%{{{placeholder}}}", value)
+            tooltip_template = tooltip_template.replace(f"%{{{placeholder}}}", value)
 
         try:
             fig.add_annotation(
-                x=x_val, y=y_val, xref=xaxis, yref=yaxis, text=template, **merged_config
+                x=x_val,
+                y=y_val,
+                xref=xaxis,
+                yref=yaxis,
+                text=tooltip_template,
+                **merged_config,
             )
         except ValueError as e:
             logger.error(