Merge pull request #1914 from jakevdp/serialization

ENH: support numpy types within chart specifications

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
 

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],
+    }

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
 

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],
+    }