Test json encoding for NaN, pd.NaT, pd.Index

chriddyp · chriddyp · commit bff27b811c74 · 2014-12-12T11:35:34.000-08:00
diff --git a/plotly/tests/test_optional/test_utils/test_utils.py b/plotly/tests/test_optional/test_utils/test_utils.py
@@ -13,9 +13,15 @@
 
 ## JSON encoding
 numeric_list = [1, 2, 3]
-np_list = np.array([1, 2, 3])
-mixed_list = [1, 'A', dt(2014, 1, 5)]
-pd = pd.DataFrame(columns=['col 1'], data=[1, 2, 3])
+np_list = np.array([1, 2, 3, np.NaN, np.NAN, np.Inf, dt(2014, 1, 5)])
+mixed_list = [1, 'A', dt(2014, 1, 5), dt(2014, 1, 5, 1, 1, 1),
+              dt(2014, 1, 5, 1, 1, 1, 1)]
+
+df = pd.DataFrame(columns=['col 1'],
+                  data=[1, 2, 3, dt(2014, 1, 5), pd.NaT, np.NaN, np.Inf])
+
+rng = pd.date_range('1/1/2011', periods=2, freq='H')
+ts = pd.Series([1.5, 2.5], index=rng)
 
 
 def test_column_json_encoding():
@@ -28,21 +34,47 @@ def test_column_json_encoding():
         columns, cls=utils._plotlyJSONEncoder, sort_keys=True
     )
     assert('[{"data": [1, 2, 3], "name": "col 1"}, '
-           '{"data": [1, "A", "2014-01-05"], "name": "col 2"}, '
-           '{"data": [1, 2, 3], "name": "col 3"}]' == json_columns)
+           '{"data": [1, "A", "2014-01-05", '
+           '"2014-01-05 01:01:01", '
+           '"2014-01-05 01:01:01.000001"], '
+           '"name": "col 2"}, '
+           '{"data": [1, 2, 3, NaN, NaN, Infinity, '
+           '"2014-01-05"], "name": "col 3"}]' == json_columns)
 
 
 def test_figure_json_encoding():
+    df = pd.DataFrame(columns=['col 1'], data=[1, 2, 3])
     s1 = Scatter3d(x=numeric_list, y=np_list, z=mixed_list)
-    s2 = Scatter(x=pd['col 1'])
+    s2 = Scatter(x=df['col 1'])
     data = Data([s1, s2])
     figure = Figure(data=data)
 
     js1 = json.dumps(s1, cls=utils._plotlyJSONEncoder, sort_keys=True)
     js2 = json.dumps(s2, cls=utils._plotlyJSONEncoder, sort_keys=True)
 
-    assert(js1 == '{"type": "scatter3d", "x": [1, 2, 3], '+
-                  '"y": [1, 2, 3], "z": [1, "A", "2014-01-05"]}')
+    assert(js1 == '{"type": "scatter3d", "x": [1, 2, 3], '
+                  '"y": [1, 2, 3, NaN, NaN, Infinity, "2014-01-05"], '
+                  '"z": [1, "A", "2014-01-05", '
+                  '"2014-01-05 01:01:01", "2014-01-05 01:01:01.000001"]}')
     assert(js2 == '{"type": "scatter", "x": [1, 2, 3]}')
     json.dumps(data, cls=utils._plotlyJSONEncoder, sort_keys=True)
     json.dumps(figure, cls=utils._plotlyJSONEncoder, sort_keys=True)
+
+
+def test_pandas_json_encoding():
+    j1 = json.dumps(df['col 1'], cls=utils._plotlyJSONEncoder)
+    assert(j1 == '[1, 2, 3, "2014-01-05", null, NaN, Infinity]')
+    j2 = json.dumps(df.index, cls=utils._plotlyJSONEncoder)
+    assert(j2 == '[0, 1, 2, 3, 4, 5, 6]')
+
+    j3 = json.dumps([pd.NaT], cls=utils._plotlyJSONEncoder)
+    assert(j3 == '[null]')
+
+    j4 = json.dumps(rng, cls=utils._plotlyJSONEncoder)
+    assert(j4 == '["2011-01-01", "2011-01-01 01:00:00"]')
+
+    j5 = json.dumps(ts, cls=utils._plotlyJSONEncoder)
+    assert(j5 == '[1.5, 2.5]')
+
+    j6 = json.dumps(ts.index, cls=utils._plotlyJSONEncoder)
+    assert(j6 == '["2011-01-01", "2011-01-01 01:00:00"]')
