Merge pull request #3022 from plotly/numpy_date_serialization

Pandas and Numpy datetime serialization fixes

Author: nicolaskruchten

Diff for: packages/python/plotly/_plotly_utils/basevalidators.py

@@ -103,6 +103,13 @@ def copy_to_readonly_numpy_array(v, kind=None, force_numeric=False):
             else:
                 # DatetimeIndex
                 v = v.to_pydatetime()
+    elif pd and isinstance(v, pd.DataFrame) and len(set(v.dtypes)) == 1:
+        dtype = v.dtypes.tolist()[0]
+        if dtype.kind in numeric_kinds:
+            v = v.values
+        elif dtype.kind == "M":
+            v = [row.dt.to_pydatetime().tolist() for i, row in v.iterrows()]
+
     if not isinstance(v, np.ndarray):
         # v has its own logic on how to convert itself into a numpy array
         if is_numpy_convertable(v):
@@ -146,7 +153,7 @@ def copy_to_readonly_numpy_array(v, kind=None, force_numeric=False):
         # datatype. This works around cases like np.array([1, 2, '3']) where
         # numpy converts the integers to strings and returns array of dtype
         # '<U21'
-        if new_v.dtype.kind not in ["u", "i", "f", "O"]:
+        if new_v.dtype.kind not in ["u", "i", "f", "O", "M"]:
             new_v = np.array(v, dtype="object")
 
     # Set new array to be read-only

Diff for: packages/python/plotly/_plotly_utils/tests/validators/test_pandas_series_input.py

@@ -173,7 +173,9 @@ def test_color_validator_categorical(color_validator, color_categorical_pandas):
     np.testing.assert_array_equal(res, np.array(color_categorical_pandas))
 
 
-def test_data_array_validator_dates(data_array_validator, datetime_pandas, dates_array):
+def test_data_array_validator_dates_series(
+    data_array_validator, datetime_pandas, dates_array
+):
 
     res = data_array_validator.validate_coerce(datetime_pandas)
 
@@ -185,3 +187,20 @@ def test_data_array_validator_dates(data_array_validator, datetime_pandas, dates
 
     # Check values
     np.testing.assert_array_equal(res, dates_array)
+
+
+def test_data_array_validator_dates_dataframe(
+    data_array_validator, datetime_pandas, dates_array
+):
+
+    df = pd.DataFrame({"d": datetime_pandas})
+    res = data_array_validator.validate_coerce(df)
+
+    # Check type
+    assert isinstance(res, np.ndarray)
+
+    # Check dtype
+    assert res.dtype == "object"
+
+    # Check values
+    np.testing.assert_array_equal(res, dates_array.reshape(len(dates_array), 1))

Diff for: packages/python/plotly/_plotly_utils/utils.py

@@ -184,8 +184,13 @@ def encode_as_numpy(obj):
 
         if obj is numpy.ma.core.masked:
             return float("nan")
-        else:
-            raise NotEncodable
+        elif isinstance(obj, numpy.ndarray) and obj.dtype.kind == "M":
+            try:
+                return numpy.datetime_as_string(obj).tolist()
+            except TypeError:
+                pass
+
+        raise NotEncodable
 
     @staticmethod
     def encode_as_datetime(obj):

Diff for: packages/python/plotly/plotly/tests/test_optional/test_utils/test_utils.py

@@ -257,6 +257,63 @@ def test_pandas_json_encoding(self):
         j6 = _json.dumps(ts.index, cls=utils.PlotlyJSONEncoder)
         assert j6 == '["2011-01-01T00:00:00", "2011-01-01T01:00:00"]'
 
+    def test_encode_customdata_datetime_series(self):
+        df = pd.DataFrame(dict(t=pd.to_datetime(["2010-01-01", "2010-01-02"])))
+
+        # 1D customdata
+        fig = Figure(
+            Scatter(x=df["t"], customdata=df["t"]), layout=dict(template="none")
+        )
+        fig_json = _json.dumps(
+            fig, cls=utils.PlotlyJSONEncoder, separators=(",", ":"), sort_keys=True
+        )
+        self.assertTrue(
+            fig_json.startswith(
+                '{"data":[{"customdata":["2010-01-01T00:00:00","2010-01-02T00:00:00"]'
+            )
+        )
+
+    def test_encode_customdata_datetime_homogenous_dataframe(self):
+        df = pd.DataFrame(
+            dict(
+                t1=pd.to_datetime(["2010-01-01", "2010-01-02"]),
+                t2=pd.to_datetime(["2011-01-01", "2011-01-02"]),
+            )
+        )
+        # 2D customdata
+        fig = Figure(
+            Scatter(x=df["t1"], customdata=df[["t1", "t2"]]),
+            layout=dict(template="none"),
+        )
+        fig_json = _json.dumps(
+            fig, cls=utils.PlotlyJSONEncoder, separators=(",", ":"), sort_keys=True
+        )
+        self.assertTrue(
+            fig_json.startswith(
+                '{"data":[{"customdata":'
+                '[["2010-01-01T00:00:00","2011-01-01T00:00:00"],'
+                '["2010-01-02T00:00:00","2011-01-02T00:00:00"]'
+            )
+        )
+
+    def test_encode_customdata_datetime_inhomogenous_dataframe(self):
+        df = pd.DataFrame(
+            dict(t=pd.to_datetime(["2010-01-01", "2010-01-02"]), v=np.arange(2),)
+        )
+        # 2D customdata
+        fig = Figure(
+            Scatter(x=df["t"], customdata=df[["t", "v"]]), layout=dict(template="none")
+        )
+        fig_json = _json.dumps(
+            fig, cls=utils.PlotlyJSONEncoder, separators=(",", ":"), sort_keys=True
+        )
+        self.assertTrue(
+            fig_json.startswith(
+                '{"data":[{"customdata":'
+                '[["2010-01-01T00:00:00",0],["2010-01-02T00:00:00",1]]'
+            )
+        )
+
     def test_numpy_masked_json_encoding(self):
         l = [1, 2, np.ma.core.masked]
         j1 = _json.dumps(l, cls=utils.PlotlyJSONEncoder)
@@ -277,6 +334,15 @@ def test_datetime_dot_date(self):
         j1 = _json.dumps(a, cls=utils.PlotlyJSONEncoder)
         assert j1 == '["2014-01-01", "2014-01-02"]'
 
+    def test_numpy_datetime64(self):
+        a = pd.date_range("2011-07-11", "2011-07-13", freq="D").values
+        j1 = _json.dumps(a, cls=utils.PlotlyJSONEncoder)
+        assert (
+            j1 == '["2011-07-11T00:00:00.000000000", '
+            '"2011-07-12T00:00:00.000000000", '
+            '"2011-07-13T00:00:00.000000000"]'
+        )
+
     def test_pil_image_encoding(self):
         import _plotly_utils