diff --git a/altair/utils/schemapi.py b/altair/utils/schemapi.py index ebb463601..abfaaf682 100644 --- a/altair/utils/schemapi.py +++ b/altair/utils/schemapi.py @@ -8,6 +8,8 @@ import json import jsonschema +import numpy as np +import pandas as pd # If DEBUG_MODE is True, then schema objects are converted to dict and @@ -54,13 +56,17 @@ def _todict(obj, validate, context): """Convert an object to a dict representation.""" if isinstance(obj, SchemaBase): return obj.to_dict(validate=validate, context=context) - elif isinstance(obj, (list, tuple)): + elif isinstance(obj, (list, tuple, np.ndarray)): return [_todict(v, validate, context) for v in obj] elif isinstance(obj, dict): return {k: _todict(v, validate, context) for k, v in obj.items() if v is not Undefined} elif hasattr(obj, 'to_dict'): return obj.to_dict() + elif isinstance(obj, np.number): + return float(obj) + elif isinstance(obj, (pd.Timestamp, np.datetime64)): + return pd.Timestamp(obj).isoformat() else: return obj diff --git a/altair/utils/tests/test_schemapi.py b/altair/utils/tests/test_schemapi.py index 53d9ec12f..1206d1d39 100644 --- a/altair/utils/tests/test_schemapi.py +++ b/altair/utils/tests/test_schemapi.py @@ -4,10 +4,13 @@ # tools/generate_schema_wrapper.py. Do not modify directly. import copy import io +import json import jsonschema import pickle import pytest +import numpy as np + from ..schemapi import (UndefinedType, SchemaBase, Undefined, _FromDict, SchemaValidationError) @@ -324,3 +327,17 @@ def test_schema_validation_error(): assert "validating {!r}".format(the_err.validator) in message assert the_err.message in message + +def test_serialize_numpy_types(): + m = MySchema( + a={'date': np.datetime64('2019-01-01')}, + a2={'int64': np.int64(1), 'float64': np.float64(2)}, + b2=np.arange(4), + ) + out = m.to_json() + dct = json.loads(out) + assert dct == { + 'a': {'date': '2019-01-01T00:00:00'}, + 'a2': {'int64': 1, 'float64': 2}, + 'b2': [0, 1, 2, 3], + } diff --git a/tools/schemapi/schemapi.py b/tools/schemapi/schemapi.py index 80f1ce994..69b222de8 100644 --- a/tools/schemapi/schemapi.py +++ b/tools/schemapi/schemapi.py @@ -4,6 +4,8 @@ import json import jsonschema +import numpy as np +import pandas as pd # If DEBUG_MODE is True, then schema objects are converted to dict and @@ -50,13 +52,17 @@ def _todict(obj, validate, context): """Convert an object to a dict representation.""" if isinstance(obj, SchemaBase): return obj.to_dict(validate=validate, context=context) - elif isinstance(obj, (list, tuple)): + elif isinstance(obj, (list, tuple, np.ndarray)): return [_todict(v, validate, context) for v in obj] elif isinstance(obj, dict): return {k: _todict(v, validate, context) for k, v in obj.items() if v is not Undefined} elif hasattr(obj, 'to_dict'): return obj.to_dict() + elif isinstance(obj, np.number): + return float(obj) + elif isinstance(obj, (pd.Timestamp, np.datetime64)): + return pd.Timestamp(obj).isoformat() else: return obj diff --git a/tools/schemapi/tests/test_schemapi.py b/tools/schemapi/tests/test_schemapi.py index d33df97ce..de669a1b2 100644 --- a/tools/schemapi/tests/test_schemapi.py +++ b/tools/schemapi/tests/test_schemapi.py @@ -1,9 +1,12 @@ import copy import io +import json import jsonschema import pickle import pytest +import numpy as np + from ..schemapi import (UndefinedType, SchemaBase, Undefined, _FromDict, SchemaValidationError) @@ -320,3 +323,17 @@ def test_schema_validation_error(): assert "validating {!r}".format(the_err.validator) in message assert the_err.message in message + +def test_serialize_numpy_types(): + m = MySchema( + a={'date': np.datetime64('2019-01-01')}, + a2={'int64': np.int64(1), 'float64': np.float64(2)}, + b2=np.arange(4), + ) + out = m.to_json() + dct = json.loads(out) + assert dct == { + 'a': {'date': '2019-01-01T00:00:00'}, + 'a2': {'int64': 1, 'float64': 2}, + 'b2': [0, 1, 2, 3], + }