Fixes #210 (#211)

* 💪 add tests for #208

* 🙏 fix for #208

* ✨ tests for #210

* 💪 code-fix for #210

* 🔧 tests for setting hf_x dynamically for #210

* 🔥 fix for setting hf_series x to a tz-aware pd.Series

* 🖊️ review

* 🔍 review code

* 🔍 fix helper method

* 🙏

---------

Co-authored-by: jvdd <[email protected]>

Author: jonasvdd

plotly_resampler/figure_resampler/figure_resampler_interface.py

@@ -326,7 +326,11 @@ def _check_update_trace_data(
             ):
                 # is faster to escape the loop here than check inside the hasattr if
                 continue
-            if hasattr(hf_trace_data[k], "values"):
+            elif pd.core.dtypes.common.is_datetime64tz_dtype(hf_trace_data[k]):
+                # When we use the .values method, timezone information is lost
+                # so convert it to pd.DatetimeIndex, which preserves the tz-info
+                hf_trace_data[k] = pd.Index(hf_trace_data[k])
+            elif hasattr(hf_trace_data[k], "values"):
                 # when not a range index or datetime index
                 hf_trace_data[k] = hf_trace_data[k].values
 
@@ -582,13 +586,17 @@ def _parse_get_trace_props(
 
         """
         hf_x: np.ndarray = (
+            # fmt: off
             (np.asarray(trace["x"]) if trace["x"] is not None else None)
             if hasattr(trace, "x") and hf_x is None
-            else hf_x.values
-            if isinstance(hf_x, pd.Series)
-            else hf_x
-            if isinstance(hf_x, pd.Index)
+            # If we cast a tz-aware datetime64 array to `.values` we lose the tz-info 
+            # and the UTC time will be displayed instead of the tz-localized time, 
+            # hence we cast to a pd.DatetimeIndex, which preserves the tz-info
+            else pd.Index(hf_x) if pd.core.dtypes.common.is_datetime64tz_dtype(hf_x)
+            else hf_x.values if isinstance(hf_x, pd.Series)
+            else hf_x if isinstance(hf_x, pd.Index)
             else np.asarray(hf_x)
+            # fmt: on
         )
 
         hf_y = (

tests/test_figure_resampler.py

@@ -529,6 +529,7 @@ def test_hf_text_and_hf_hovertext():
 def test_multiple_timezones():
     n = 5_050
 
+    # NOTE: date-range returns a (tz-aware) DatetimeIndex
     dr = pd.date_range("2022-02-14", freq="s", periods=n, tz="UTC")
     dr_v = np.random.randn(n)
 
@@ -538,6 +539,12 @@ def test_multiple_timezones():
         dr.tz_convert("Europe/Brussels"),
         dr.tz_convert("Australia/Perth"),
         dr.tz_convert("Australia/Canberra"),
+        # NOTE: this pd.Series tests the functionality of a Pandas series with (tz-aware) DatetimeIndex
+        pd.Series(dr),
+        pd.Series(dr.tz_localize(None).tz_localize("Europe/Amsterdam")),
+        pd.Series(dr.tz_convert("Europe/Brussels")),
+        pd.Series(dr.tz_convert("Australia/Perth")),
+        pd.Series(dr.tz_convert("Australia/Canberra")),
     ]
 
     plain_plotly_fig = make_subplots(rows=len(cs), cols=1, shared_xaxes=True)
@@ -564,7 +571,60 @@ def test_multiple_timezones():
             col=1,
         )
         # Assert that the time parsing is exactly the same
-        assert plain_plotly_fig.data[0].x[0] == fr_fig.data[0].x[0]
+        assert plain_plotly_fig.data[i - 1].x[0] == fr_fig.data[i - 1].x[0]
+
+
+def test_set_hfx_tz_aware_series():
+    df = pd.DataFrame(
+        {
+            "timestamp": pd.date_range(
+                "2020-01-01", "2020-01-02", freq="1s"
+            ).tz_localize("Asia/Seoul")
+        }
+    )
+    df["value"] = np.random.randn(len(df))
+
+    fr = FigureResampler()
+    fr.add_trace({}, hf_x=pd.Index(df.timestamp), hf_y=df.value)
+    assert isinstance(fr.hf_data[0]["x"], pd.DatetimeIndex)
+    # Now we set the pd.Series as hf_x
+    fr.hf_data[0]["x"] = df.timestamp
+    assert not isinstance(fr.hf_data[0]["x"], pd.DatetimeIndex)
+    # perform an update
+    out = fr.construct_update_data({"xaxis.autorange": True, "xaxis.showspikes": False})
+    assert len(out) == 2
+    # assert that the update was performed correctly
+    assert isinstance(fr.hf_data[0]["x"], pd.DatetimeIndex)
+    assert all(fr.hf_data[0]["x"] == pd.DatetimeIndex(df.timestamp))
+
+
+def test_datetime_hf_x_no_index_():
+    df = pd.DataFrame(
+        {"timestamp": pd.date_range("2020-01-01", "2020-01-02", freq="1s")}
+    )
+    df["value"] = np.random.randn(len(df))
+
+    # add via hf_x kwargs
+    fr = FigureResampler()
+    fr.add_trace({}, hf_x=df.timestamp, hf_y=df.value)
+    output = fr.construct_update_data(
+        {
+            "xaxis.range[0]": "2020-01-01 00:00:00",
+            "xaxis.range[1]": "2020-01-01 00:00:20",
+        }
+    )
+    assert len(output) == 2
+
+    # add via scatter kwargs
+    fr = FigureResampler()
+    fr.add_trace(go.Scatter(x=df.timestamp, y=df.value))
+    output = fr.construct_update_data(
+        {
+            "xaxis.range[0]": "2020-01-01 00:00:00",
+            "xaxis.range[1]": "2020-01-01 00:00:20",
+        }
+    )
+    assert len(output) == 2
 
 
 def test_datetime_hf_x_no_index():

tests/utils.py

@@ -27,7 +27,10 @@ def construct_hf_data_dict(hf_x, hf_y, **kwargs):
     hf_data_dict = {
         "x": hf_x,
         "y": hf_y,
-        "axis_type": "date" if isinstance(hf_x, pd.DatetimeIndex) else "linear",
+        "axis_type": "date"
+        if isinstance(hf_x, pd.DatetimeIndex)
+        or pd.core.dtypes.common.is_datetime64_any_dtype(hf_x)
+        else "linear",
         "downsampler": MinMaxLTTB(),
         "gap_handler": MedDiffGapHandler(),
         "max_n_samples": 1_000,