From cbe94c384a014e4bae2e6baace5df1db4dd437b0 Mon Sep 17 00:00:00 2001
From: drewoldag <47493171+drewoldag@users.noreply.github.com>
Date: Mon, 4 Dec 2023 15:53:36 -0800
Subject: [PATCH] Adding basic support for a data handle that reads/writes json
 to/from a dict.

---
 src/rail/core/data.py | 26 ++++++++++++++++++++++++++
 1 file changed, 26 insertions(+)

diff --git a/src/rail/core/data.py b/src/rail/core/data.py
index 73739ea1..f401e80e 100644
--- a/src/rail/core/data.py
+++ b/src/rail/core/data.py
@@ -3,6 +3,7 @@
 import os
 import tables_io
 import pickle
+import json
 import qp
 
 
@@ -333,6 +334,31 @@ def default_model_write(model, path):
         pickle.dump(obj=model, file=fout, protocol=pickle.HIGHEST_PROTOCOL)
 
 
+class JsonHandle(DataHandle):
+    """Generic handle for reading in a whole file in json format as a dictionary.
+    """
+
+    suffix = 'json'
+
+    @classmethod
+    def _open(cls, path, **kwargs):
+        """Open and return the contents of the json file."""
+        return cls.read(path, **kwargs)
+
+    @classmethod
+    def _read(cls, path, **kwargs):
+        """Read the contents of the json file into a python dictionary"""
+        with open(path, "r") as infile:
+            data = json.loads(infile)
+        return data
+
+    @classmethod
+    def _write(cls, data, path, **kwargs):
+        """Write out the provided data dictionary as json."""
+        with open(path, "w") as outfile:
+            json.dumps(data, outfile)
+
+
 class ModelDict(dict):
     """
     A specialized dict to keep track of individual estimation models objects: this is just a dict these additional features