Repository: itamarst/eliot

Issue:
Support for serializing numpy objects in log messages
Scientific computing often involves objects like `numpy.int64`. We should support serializing them to JSON transparently.

* [x] Support JSON serialization of Numpy ints/floats/bools/arrays, automatically.
* [x] Numpy should not be required.

Comments:

Patch:
diff --git a/docs/source/news.rst b/docs/source/news.rst
index 0c07ee9..e7b8885 100644
--- a/docs/source/news.rst
+++ b/docs/source/news.rst
@@ -1,6 +1,13 @@
 What's New
 ==========
 
+1.6.0
+^^^^^
+
+Features:
+
+* NumPy integers, floats, bools and arrays are now automatically serialized to JSON.
+
 1.5.0
 ^^^^^
 
diff --git a/eliot/_json.py b/eliot/_json.py
new file mode 100644
index 0000000..845ed3b
--- /dev/null
+++ b/eliot/_json.py
@@ -0,0 +1,27 @@
+"""Custom JSON encoding support."""
+
+import json
+import sys
+
+
+class EliotJSONEncoder(json.JSONEncoder):
+    """JSON encoder with additional functionality.
+
+    In particular, supports NumPy types.
+    """
+
+    def default(self, o):
+        numpy = sys.modules.get("numpy", None)
+        if numpy is not None:
+            if isinstance(o, numpy.floating):
+                return float(o)
+            if isinstance(o, numpy.integer):
+                return int(o)
+            if isinstance(o, (numpy.bool, numpy.bool_)):
+                return bool(o)
+            if isinstance(o, numpy.ndarray):
+                return o.tolist()
+        return json.JSONEncoder.default(self, o)
+
+__all__ = ["EliotJSONEncoder"]
+
diff --git a/eliot/_output.py b/eliot/_output.py
index b365973..56c7bb2 100644
--- a/eliot/_output.py
+++ b/eliot/_output.py
@@ -22,6 +22,7 @@ from ._message import (
     REASON_FIELD,
 )
 from ._util import saferepr, safeunicode
+from ._json import EliotJSONEncoder
 
 
 class _DestinationsSendError(Exception):
@@ -375,7 +376,7 @@ class FileDestination(PClass):
     _dumps = field(mandatory=True)
     _linebreak = field(mandatory=True)
 
-    def __new__(cls, file, encoder=pyjson.JSONEncoder):
+    def __new__(cls, file, encoder=EliotJSONEncoder):
         unicodeFile = False
         if PY3:
             try:
@@ -408,7 +409,7 @@ class FileDestination(PClass):
         self.file.flush()
 
 
-def to_file(output_file, encoder=pyjson.JSONEncoder):
+def to_file(output_file, encoder=EliotJSONEncoder):
     """
     Add a destination that writes a JSON message per line to the given file.
 
diff --git a/tox.ini b/tox.ini
index 8a524f2..d4dc529 100644
--- a/tox.ini
+++ b/tox.ini
@@ -1,5 +1,5 @@
 [tox]
-envlist = pypy, py27, py34, py35, py36, py37,
+envlist = pypy, py27, py34, py35-numpy, py36, py37,
           py27-twisted-latest, py35-twisted-latest,
           linters3, sphinx
 
@@ -8,7 +8,7 @@ python =
        2.7: py27, py27-twisted-latest
        pypy: pypy
        3.4: py34
-       3.5: py35, py35-twisted-latest, linters3, sphinx
+       3.5: py35-numpy, py35-twisted-latest, linters3, sphinx
        3.6: py36
        3.7: py37
 
@@ -30,9 +30,9 @@ basepython = python2.7
 basepython = python3.4
 deps = cffi
 
-[testenv:py35]
+[testenv:py35-numpy]
 basepython = python3.5
-deps = cffi
+deps = numpy
 
 [testenv:py36]
 basepython = python3.6


Test patch:
diff --git a/eliot/tests/test_json.py b/eliot/tests/test_json.py
new file mode 100644
index 0000000..87c0a80
--- /dev/null
+++ b/eliot/tests/test_json.py
@@ -0,0 +1,45 @@
+"""
+Tests for L{eliot._json}.
+"""
+
+from __future__ import unicode_literals, absolute_import
+
+from unittest import TestCase, skipUnless, skipIf
+from json import loads, dumps
+
+try:
+    import numpy as np
+except ImportError:
+    np = None
+
+from eliot._json import EliotJSONEncoder
+
+
+class EliotJSONEncoderTests(TestCase):
+    """Tests for L{EliotJSONEncoder}."""
+
+    @skipUnless(np, "NumPy not installed.")
+    def test_numpy(self):
+        """NumPy objects get serialized to readable JSON."""
+        l = [np.float32(12.5), np.float64(2.0), np.float16(0.5),
+             np.bool(True), np.bool(False), np.bool_(True),
+             np.unicode_("hello"),
+             np.byte(12), np.short(12), np.intc(-13), np.int_(0),
+             np.longlong(100), np.intp(7),
+             np.ubyte(12), np.ushort(12), np.uintc(13),
+             np.ulonglong(100), np.uintp(7),
+             np.int8(1), np.int16(3), np.int32(4), np.int64(5),
+             np.uint8(1), np.uint16(3), np.uint32(4), np.uint64(5)]
+        l2 = [l, np.array([1, 2, 3])]
+        roundtripped = loads(dumps(l2, cls=EliotJSONEncoder))
+        self.assertEqual([l, [1, 2, 3]], roundtripped)
+
+    @skipIf(np, "NumPy is installed.")
+    def test_numpy_not_imported(self):
+        """If NumPy is not available, EliotJSONEncoder continues to work.
+
+        This ensures NumPy isn't a hard dependency.
+        """
+        with self.assertRaises(TypeError):
+            dumps([object()], cls=EliotJSONEncoder)
+        self.assertEqual(dumps(12, cls=EliotJSONEncoder), "12")
diff --git a/eliot/tests/test_output.py b/eliot/tests/test_output.py
index 54a7b23..568204d 100644
--- a/eliot/tests/test_output.py
+++ b/eliot/tests/test_output.py
@@ -5,7 +5,7 @@ Tests for L{eliot._output}.
 from __future__ import unicode_literals
 
 from sys import stdout
-from unittest import TestCase, skipIf
+from unittest import TestCase, skipIf, skipUnless
 # Make sure to use StringIO that only accepts unicode:
 from io import BytesIO, StringIO
 import json as pyjson
@@ -14,7 +14,10 @@ from time import time
 from uuid import UUID
 
 from six import PY3, PY2
-
+try:
+    import numpy as np
+except ImportError:
+    np = None
 from zope.interface.verify import verifyClass
 
 from .._output import (
@@ -686,6 +689,21 @@ class ToFileTests(TestCase):
             json.loads(line)
             for line in bytes_f.getvalue().splitlines()], [{"x": "abc"}])
 
+    @skipUnless(np, "NumPy is not installed.")
+    def test_default_encoder_is_EliotJSONEncoder(self):
+        """The default encoder if none are specified is EliotJSONEncoder."""
+        message = {"x": np.int64(3)}
+        f = StringIO()
+        destination = FileDestination(file=f)
+        destination(message)
+        self.assertEqual(
+            [
+                json.loads(line)
+                for line in f.getvalue().splitlines()
+            ],
+            [{"x": 3}]
+        )
+
     def test_filedestination_writes_json_bytes(self):
         """
         L{FileDestination} writes JSON-encoded messages to a file that accepts