[7.x] Serialize only 'pd.NA' to 'None'

Author: sethmlarson

elasticsearch/serializer.py

@@ -67,7 +67,7 @@ class JSONSerializer(object):
     mimetype = "application/json"
 
     def default(self, data):
-        if isinstance(data, TIME_TYPES):
+        if isinstance(data, TIME_TYPES) and getattr(pd, "NaT", None) is not data:
             return data.isoformat()
         elif isinstance(data, uuid.UUID):
             return str(data)
@@ -87,7 +87,7 @@ def default(self, data):
         if pd:
             if isinstance(data, (pd.Series, pd.Categorical)):
                 return data.tolist()
-            elif hasattr(pd, "NA") and pd.isna(data):
+            elif data is getattr(pd, "NA", None):
                 return None
 
         raise TypeError("Unable to serialize %r (type: %s)" % (data, type(data)))

test_elasticsearch/test_serializer.py

@@ -93,6 +93,11 @@ def test_serializes_numpy_ndarray(self):
             JSONSerializer().dumps({"d": np.zeros((2, 2), dtype=np.uint8)}),
         )
 
+    def test_serializes_numpy_nan_to_nan(self):
+        self.assertEqual(
+            '{"d":NaN}', JSONSerializer().dumps({"d": np.nan}),
+        )
+
     def test_serializes_pandas_timestamp(self):
         self.assertEqual(
             '{"d":"2010-10-01T02:30:00"}',
@@ -112,6 +117,11 @@ def test_serializes_pandas_na(self):
             '{"d":null}', JSONSerializer().dumps({"d": pd.NA}),
         )
 
+    def test_raises_serialization_error_pandas_nat(self):
+        if not hasattr(pd, "NaT"):
+            raise SkipTest("pandas.NaT required")
+        self.assertRaises(SerializationError, JSONSerializer().dumps, {"d": pd.NaT})
+
     def test_serializes_pandas_category(self):
         cat = pd.Categorical(["a", "c", "b", "a"], categories=["a", "b", "c"])
         self.assertEqual(