From 084295787d34d2d7a60fc894a00f36842cd55561 Mon Sep 17 00:00:00 2001
From: Seth Michael Larson <seth.larson@elastic.co>
Date: Mon, 29 Jun 2020 10:10:05 -0500
Subject: [PATCH] [7.x] Serialize only 'pd.NA' to 'None'

---
 elasticsearch/serializer.py           |  4 ++--
 test_elasticsearch/test_serializer.py | 10 ++++++++++
 2 files changed, 12 insertions(+), 2 deletions(-)

diff --git a/elasticsearch/serializer.py b/elasticsearch/serializer.py
index 50f98c0f..7c3d470a 100644
--- a/elasticsearch/serializer.py
+++ b/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 @@ class JSONSerializer(object):
         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)))
diff --git a/test_elasticsearch/test_serializer.py b/test_elasticsearch/test_serializer.py
index 78200845..41990422 100644
--- a/test_elasticsearch/test_serializer.py
+++ b/test_elasticsearch/test_serializer.py
@@ -93,6 +93,11 @@ class TestJSONSerializer(TestCase):
             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 @@ class TestJSONSerializer(TestCase):
             '{"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(