Skip to content

Commit 974fcba

Browse files
committed
Rework date and string encoding, add and fix tests
1 parent 58b7192 commit 974fcba

File tree

2 files changed

+44
-28
lines changed

2 files changed

+44
-28
lines changed

packages/python/plotly/plotly/io/_json.py

Lines changed: 27 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -3,8 +3,7 @@
33
from six import string_types
44
import json
55
import decimal
6-
import os
7-
6+
import datetime
87

98
from plotly.io._utils import validate_coerce_fig_to_dict, validate_coerce_output_type
109
from _plotly_utils.optional_imports import get_module
@@ -474,12 +473,19 @@ def clean_to_json_compatible(obj, **kwargs):
474473
if np is not None:
475474
if obj is np.ma.core.masked:
476475
return float("nan")
477-
elif (
478-
numpy_allowed
479-
and isinstance(obj, np.ndarray)
480-
and obj.dtype.kind in ("b", "i", "u", "f")
481-
):
482-
return np.ascontiguousarray(obj)
476+
elif isinstance(obj, np.ndarray):
477+
if numpy_allowed and obj.dtype.kind in ("b", "i", "u", "f"):
478+
return np.ascontiguousarray(obj)
479+
elif obj.dtype.kind == "M":
480+
# datetime64 array
481+
return np.datetime_as_string(obj).tolist()
482+
elif obj.dtype.kind == "U":
483+
return obj.tolist()
484+
elif obj.dtype.kind == "O":
485+
# Treat object array as a lists, continue processing
486+
obj = obj.tolist()
487+
elif isinstance(obj, np.datetime64):
488+
return str(obj)
483489

484490
# pandas
485491
if pd is not None:
@@ -496,35 +502,29 @@ def clean_to_json_compatible(obj, **kwargs):
496502

497503
if not datetime_allowed:
498504
# Note: We don't need to handle dropping timezones here because
499-
# numpy's datetime64 doesn't support them and pandas's tolist()
500-
# doesn't preserve them.
505+
# numpy's datetime64 doesn't support them and pandas's tz_localize
506+
# above drops them.
501507
for i in range(len(dt_values)):
502508
dt_values[i] = dt_values[i].isoformat()
503509

504510
return dt_values
505511

506512
# datetime and date
507-
if not datetime_allowed:
508-
try:
509-
# Need to drop timezone for scalar datetimes
510-
return obj.replace(tzinfo=None).isoformat()
511-
except (TypeError, AttributeError):
512-
pass
513+
try:
514+
# Need to drop timezone for scalar datetimes. Don't need to convert
515+
# to string since engine can do that
516+
obj = obj.replace(tzinfo=None)
517+
obj = obj.to_pydatetime()
518+
except(TypeError, AttributeError):
519+
pass
513520

514-
if np and isinstance(obj, np.datetime64):
515-
return str(obj)
516-
else:
517-
res = None
521+
if not datetime_allowed:
518522
try:
519-
# Need to drop timezone for scalar datetimes. Don't need to convert
520-
# to string since engine can do that
521-
res = obj.replace(tzinfo=None)
522-
res = res.to_pydatetime()
523+
return obj.isoformat()
523524
except(TypeError, AttributeError):
524525
pass
525-
526-
if res is not None:
527-
return res
526+
elif isinstance(obj, datetime.datetime):
527+
return obj
528528

529529
# Try .tolist() convertible, do not recurse inside
530530
try:

packages/python/plotly/plotly/tests/test_io/test_to_from_plotly_json.py

Lines changed: 17 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -96,6 +96,11 @@ def object_numpy_array(request):
9696
return np.array(["a", 1, [2, 3]])
9797

9898

99+
@pytest.fixture(scope="module")
100+
def numpy_unicode_array(request):
101+
return np.array(["A", "BB", "CCC"], dtype="U")
102+
103+
99104
@pytest.fixture(
100105
scope="module",
101106
params=[
@@ -118,6 +123,7 @@ def datetime_value(request):
118123
lambda a: pd.DatetimeIndex(a), # Pandas DatetimeIndex
119124
lambda a: pd.Series(pd.DatetimeIndex(a)), # Pandas Datetime Series
120125
lambda a: pd.DatetimeIndex(a).values, # Numpy datetime64 array
126+
lambda a: np.array(a, dtype="object"), # Numpy object array of datetime
121127
]
122128
)
123129
def datetime_array(request, datetime_value):
@@ -143,6 +149,16 @@ def test_numeric_numpy_encoding(numeric_numpy_array, engine, pretty):
143149
check_roundtrip(result, engine=engine, pretty=pretty)
144150

145151

152+
def test_numpy_unicode_encoding(numpy_unicode_array, engine, pretty):
153+
value = build_test_dict(numpy_unicode_array)
154+
result = pio.to_json_plotly(value, engine=engine, pretty=pretty)
155+
156+
array_str = to_json_test(numpy_unicode_array.tolist())
157+
expected = build_test_dict_string(array_str)
158+
assert result == expected
159+
check_roundtrip(result, engine=engine, pretty=pretty)
160+
161+
146162
def test_object_numpy_encoding(object_numpy_array, engine, pretty):
147163
value = build_test_dict(object_numpy_array)
148164
result = pio.to_json_plotly(value, engine=engine, pretty=pretty)
@@ -191,7 +207,7 @@ def to_str(v):
191207
elif isinstance(datetime_array, pd.DatetimeIndex):
192208
dt_values = [to_str(d) for d in datetime_array.to_pydatetime().tolist()]
193209
else: # numpy datetime64 array
194-
dt_values = datetime_array.tolist()
210+
dt_values = [to_str(d) for d in datetime_array]
195211

196212
array_str = to_json_test(dt_values)
197213
expected = build_test_dict_string(array_str)

0 commit comments

Comments
 (0)