Serialize common numpy data types in write_json

CHANGES.md

@@ -6,7 +6,12 @@
 
 * titers: Support parsing of thresholded values (e.g., "<80" or ">2560"). [#1118][] (@huddlej)
 
+### Bug Fixes
+
+* utils: Serialize common numpy data types in `write_json`. [#1119][] (@victorlin)
+
 [#1118]: https://github.com/nextstrain/augur/pull/1118
+[#1119]: https://github.com/nextstrain/augur/pull/1119
 
 ## 19.2.0 (19 December 2022)
 

augur/utils.py

@@ -2,6 +2,7 @@
 import Bio
 import Bio.Phylo
 import gzip
+import numpy as np
 import os, json, sys
 import pandas as pd
 import subprocess
@@ -124,7 +125,19 @@ 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)
+        json.dump(data, handle, indent=indent, sort_keys=sort_keys, cls=NumpyJSONEncoder)
+
+
+class NumpyJSONEncoder(json.JSONEncoder):
+    """A custom JSONEncoder subclass to serialize additional numpy data types."""
+    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()
+        return super().default(obj)
 
 
 def load_features(reference, feature_names=None):

tests/test_utils.py

@@ -1,3 +1,5 @@
+import json
+import numpy as np
 from pathlib import Path
 from unittest.mock import patch
 
@@ -90,3 +92,19 @@ def test_read_strains(self, tmpdir):
         strains = utils.read_strains(strains1, strains2)
         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."""
+        data = {
+            'int': np.int64(1),
+            'float': np.float64(2.0),
+            'array': np.array([3,4,5])
+        }
+        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]
+            }