Optional flags allow to convert python NaN to null in JSON writer

Author: maciekrb

mapreduceutils/__init__.py

@@ -310,6 +310,7 @@ def record_map(data_record):
   ctx = context.get()
   property_map = ctx.mapreduce_spec.mapper.params.get('property_map')
   output_format = ctx.mapreduce_spec.mapper.params.get('output_format', 'JSON')
+  writer_args = ctx.mapreduce_spec.mapper.params.get('writer_args', dict())
 
   record = MapperRecord.create(data_record)
   if record:
@@ -324,7 +325,7 @@ def record_map(data_record):
           writer = OutputWriter.get_writer(output_format)
           if 'mapper_key_spec' in map_rule:
             key = record.mapper_key(map_rule.get('mapper_key_spec'))
-            data = writer.write(row)
+            data = writer.write(row, **writer_args)
             #logging.warn("Mapper pre yield MR: {}:{}".format(key, data))
             yield (key, data)
           else:

mapreduceutils/tests/test_writers.py

@@ -1,5 +1,6 @@
 #!/usr/bin/env python
 from collections import OrderedDict
+import datetime
 from mapreduceutils.writers import OutputWriter
 import unittest
 
@@ -66,3 +67,21 @@ def test_escaped_output(self):
       '}\r\n'
     )
     self.assertEqual(expected, out)
+
+  def test_nan_as_null_output(self):
+    record = OrderedDict([
+      ("sprop_abc", float('NaN')),
+      ("sprop_bcd", 'test "message"'),
+      ("sprop_cde", datetime.datetime(2014, 8, 1, 0, 0, 0))
+    ])
+    writer = OutputWriter.get_writer('json')
+    out = writer.write(record, nan_to_null=True)
+
+    expected = (
+      '{'
+      '"sprop_cde": "2014-08-01 00:00:00", '
+      '"sprop_bcd": "test \\"message\\"", '
+      '"sprop_abc": null'
+      '}\r\n'
+    )
+    self.assertEqual(expected, out)

mapreduceutils/writers.py

@@ -2,6 +2,7 @@
 from cStringIO import StringIO
 import datetime
 from json import JSONEncoder
+import math
 
 
 class OutputWriter:
@@ -63,5 +64,18 @@ def write(cls, data_obj):
 
 class JSONWriter(OutputWriter):
   @classmethod
-  def write(cls, data_obj):
+  def write(cls, data_obj, nan_to_null=False):
+    """
+    Encodes given python object to JSON
+
+    Args:
+      data_obj:      Python object to encode to JSON
+      nan_to_null:   (Bool) If True NaN values will be converted to None
+                     before encoding to JSON
+    Returns:
+      JSON string
+    """
+    if nan_to_null:
+      data_obj = {k: None if isinstance(v, float) and math.isnan(v) else v
+                  for k, v in data_obj.items()}
     return "{}\r\n".format(MapperJSONEncoder().encode(data_obj))