Merge pull request #1213: utils.write_json: Serialize Pandas Series

CHANGES.md

@@ -2,6 +2,11 @@
 
 ## __NEXT__
 
+### Bug fixes
+
+* utils: Serialize pandas Series in `write_json`. [#1213][] (@victorlin)
+
+[#1213]: https://github.com/nextstrain/augur/pull/1213
 
 ## 22.0.2 (26 May 2023)
 

augur/utils.py

@@ -127,18 +127,23 @@ def write_json(data, file_name, indent=(None if os.environ.get("AUGUR_MINIFY_JSO
         data["generated_by"] = {"program": "augur", "version": get_augur_version()}
     with open(file_name, 'w', encoding='utf-8') as handle:
         sort_keys = False if isinstance(data, OrderedDict) else True
-        json.dump(data, handle, indent=indent, sort_keys=sort_keys, cls=NumpyJSONEncoder)
+        json.dump(data, handle, indent=indent, sort_keys=sort_keys, cls=AugurJSONEncoder)
 
 
-class NumpyJSONEncoder(json.JSONEncoder):
-    """A custom JSONEncoder subclass to serialize additional numpy data types."""
+class AugurJSONEncoder(json.JSONEncoder):
+    """
+    A custom JSONEncoder subclass to serialize data types used for various data
+    stored in dictionary format.
+    """
     def default(self, obj):
         if isinstance(obj, np.integer):
             return int(obj)
         if isinstance(obj, np.floating):
             return float(obj)
         if isinstance(obj, np.ndarray):
             return obj.tolist()
+        if isinstance(obj, pd.Series):
+            return obj.tolist()
         return super().default(obj)
 
 

tests/test_utils.py

@@ -2,6 +2,7 @@
 import numpy as np
 from pathlib import Path
 from unittest.mock import patch
+import pandas as pd
 
 import pytest
 
@@ -93,18 +94,20 @@ def test_read_strains(self, tmpdir):
         assert len(strains) == 3
         assert "strain1" in strains
 
-    def test_write_json_numpy_types(self, tmpdir):
-        """write_json should be able to serialize numpy data types."""
+    def test_write_json_data_types(self, tmpdir):
+        """write_json should be able to serialize various data types."""
         data = {
             'int': np.int64(1),
             'float': np.float64(2.0),
-            'array': np.array([3,4,5])
+            'array': np.array([3,4,5]),
+            'series': pd.Series([6,7,8])
         }
         file = Path(tmpdir) / Path("data.json")
         utils.write_json(data, file, include_version=False)
         with open(file) as f:
             assert json.load(f) == {
                 'int': 1,
                 'float': 2.0,
-                'array': [3,4,5]
+                'array': [3,4,5],
+                'series': [6,7,8]
             }