ENH #137 add utils.export_json() and test

Author: prjemian

apstools/utils.py

@@ -8,6 +8,7 @@
    ~EmailNotifications
    ~ExcelDatabaseFileBase
    ~ExcelDatabaseFileGeneric
+   ~export_json
    ~ipython_profile_name
    ~pairwise
    ~print_snapshot_list
@@ -31,6 +32,9 @@
 from collections import OrderedDict
 import datetime
 from email.mime.text import MIMEText
+from event_model import NumpyEncoder
+from io import StringIO
+import json
 import logging
 import math
 import os
@@ -40,6 +44,7 @@
 import smtplib
 import subprocess
 import time
+import zipfile
 
 from .plans import run_in_thread
 
@@ -535,3 +540,50 @@ def print_snapshot_list(db, **search_criteria):
         n = len(list(h.start.keys()))
         t.addRow((i, uid, h.start["iso8601"], n, h.start["purpose"]))
     print(t) 
+
+
+def export_json(headers, filename, zipfilename=None):
+    """
+    write a list of headers (all documents) to a file
+
+    PARAMETERS
+
+    headers : list(headers) or `databroker._core.Results` object
+        list of databroker headers as returned from `db(...search criteria...)`
+    filename : str
+        name of file into which to write JSON
+    zipfilename : str or None
+        name of ZIP file container of `filename` 
+        (if None, do not ZIP `filename`)
+    
+    EXAMPLE::
+
+        from databroker import Broker
+        db = Broker.named("mongodb_config")
+        headers = db(plan_name="count", since="2019-04-01")
+
+        export_json(
+            headers, 
+            "data.json", 
+            zipfilename="bluesky_data.zip")
+    
+    EXAMPLE: READ THE ZIP FILE::
+
+        with zipfile.ZipFile(zipfilename, "r") as fp:
+            buf = fp.read(filename).decode("utf-8")
+            header_dict = json.loads(buf)
+
+    """
+    datasets = [list(h.documents()) for h in headers]
+#     for h in headers:
+#         uid = h.start["uid"]
+#         datasets[uid] = list(h.documents())
+
+    json_docs = json.dumps(datasets, cls=NumpyEncoder, indent=2)
+    if zipfilename is None:
+        with open(filename, "w") as fp:
+            fp.write(json_docs)
+    else:
+        with zipfile.ZipFile(zipfilename, "w", allowZip64=True) as fp:
+            fp.writestr(filename, json_docs, compress_type=zipfile.ZIP_LZMA)
+                

tests/__main__.py

@@ -14,10 +14,12 @@ def suite(*args, **kw):
 
     import test_simple
     import test_filewriter
+    import test_export_json
     # import test_excel
     test_list = [
         test_simple,
         test_filewriter,
+        test_export_json,
         # test_excel
         ]
 

tests/test_export_json.py

@@ -0,0 +1,80 @@
+
+"""
+unit tests for the SPEC filewriter
+"""
+
+import json
+import os
+import shutil
+import sys
+import tempfile
+import unittest
+import zipfile
+
+
+_test_path = os.path.dirname(__file__)
+_path = os.path.join(_test_path, '..')
+if _path not in sys.path:
+    sys.path.insert(0, _path)
+
+from apstools.utils import export_json
+
+
+class Test_ExportZippedJson(unittest.TestCase):
+
+    def setUp(self):
+        self.tempdir = tempfile.mkdtemp()
+
+    def tearDown(self):
+        if os.path.exists(self.tempdir):
+            shutil.rmtree(self.tempdir, ignore_errors=True)
+    
+    def test_writer_default_name(self):
+        from databroker import Broker
+        db = Broker.named("mongodb_config")
+        headers = db(plan_name="count")
+        headers = list(headers)[0:1]
+
+        filename = os.path.join(self.tempdir, "export1.txt")
+        export_json(headers, filename=filename)
+        self.assertTrue(os.path.exists(filename), f"wrote to requested {filename}")
+        
+        filename = "export2.txt"
+        zipfilename = os.path.join(self.tempdir, "export2.zip")
+        export_json(headers, filename, zipfilename=zipfilename)
+        self.assertFalse(os.path.exists(filename), f"did not write to {filename}")
+        self.assertTrue(
+            os.path.exists(zipfilename), 
+            f"wrote to requested ZIP {zipfilename}")
+        with zipfile.ZipFile(zipfilename, "r") as fp:
+            self.assertIn(filename, fp.namelist(), "found JSON test data")
+            buf = fp.read(filename).decode("utf-8")
+            testdata = json.loads(buf)
+        self.assertEqual(len(testdata), 1, "ZIP file contains one dataset")
+        dataset = testdata[0]
+        self.assertGreater(len(dataset), 1, "dataset contains more than one document")
+        tag, doc = dataset[0]
+        self.assertEqual(tag, "start", "found start document")
+        self.assertNotEqual(doc.get("plan_name"), None, "found a start document by duck type")
+        self.assertNotEqual(doc.get("uid"), None, "found a uid document")
+        self.assertEqual(
+            doc["uid"],
+            headers[0].start["uid"],
+            "found matching start document"
+            )
+
+
+def suite(*args, **kw):
+    test_list = [
+        Test_ExportZippedJson,
+        ]
+
+    test_suite = unittest.TestSuite()
+    for test_case in test_list:
+        test_suite.addTest(unittest.makeSuite(test_case))
+    return test_suite
+
+
+if __name__ == "__main__":
+    runner=unittest.TextTestRunner()
+    runner.run(suite())