Merge pull request #5204 from chu11/python_bindings_kvs

python: support non-JSON formatted data in KVS

Author: mergify[bot]

src/bindings/python/flux/kvs.py

@@ -35,7 +35,12 @@ def get_key_direct(flux_handle, key):
     if valp[0] == ffi.NULL:
         return None
 
-    ret = json.loads(ffi.string(valp[0]).decode("utf-8"))
+    try:
+        ret = json.loads(ffi.string(valp[0]).decode("utf-8"))
+    except json.decoder.JSONDecodeError:
+        ret = ffi.string(valp[0]).decode("utf-8")
+    except UnicodeDecodeError:
+        ret = ffi.string(valp[0])
     RAW.flux_future_destroy(future)
     return ret
 
@@ -78,10 +83,17 @@ def get(flux_handle, key):
 
 
 def put(flux_handle, key, value):
-    json_str = json.dumps(value)
     if flux_handle.aux_txn is None:
         flux_handle.aux_txn = RAW.flux_kvs_txn_create()
-    return RAW.flux_kvs_txn_put(flux_handle.aux_txn, 0, key, json_str)
+    try:
+        json_str = json.dumps(value)
+        return RAW.flux_kvs_txn_put(flux_handle.aux_txn, 0, key, json_str)
+    except TypeError:
+        if isinstance(value, bytes):
+            return RAW.flux_kvs_txn_put_raw(
+                flux_handle.aux_txn, 0, key, value, len(value)
+            )
+        raise TypeError
 
 
 def put_mkdir(flux_handle, key):

t/python/t0005-kvs.py

@@ -31,35 +31,51 @@ def test_kvs_dir_open(self):
         with flux.kvs.get_dir(self.f) as d:
             self.assertIsNotNone(d)
 
-    def set_and_check_context(self, key, value, msg=""):
+    def set_and_check_context(self, key, value, type):
         kd = flux.kvs.KVSDir(self.f)
         kd[key] = value
         kd.commit()
-        nv = kd[key]
-        self.assertEqual(value, nv)
-        self.assertFalse(isinstance(nv, bytes))
-        return kd
+
+        kd2 = flux.kvs.KVSDir(self.f)
+        nv = kd2[key]
+        if isinstance(value, bytes) and type is str:
+            self.assertEqual(value.decode('utf-8'), nv)
+        else:
+            self.assertEqual(value, nv)
+        if type is not None:
+            self.assertTrue(isinstance(nv, type))
+
+        return kd2
 
     def test_set_int(self):
-        self.set_and_check_context("int", 10)
+        self.set_and_check_context("int", 10, int)
 
     def test_set_float(self):
-        self.set_and_check_context("float", 10.5)
+        self.set_and_check_context("float", 10.5, float)
 
     def test_set_string(self):
-        self.set_and_check_context("string", "stuff")
+        self.set_and_check_context("string", "stuff", str)
+
+    def test_set_none(self):
+        self.set_and_check_context("none", None, None)
 
     def test_set_unicode(self):
-        self.set_and_check_context("unicode", "\u32db \u263a \u32e1")
+        self.set_and_check_context("unicode", "\u32db \u263a \u32e1", str)
+
+    def test_set_bytes(self):
+        self.set_and_check_context("bytes", bytes.fromhex("deadbeef"), bytes)
 
     def test_set_list(self):
-        self.set_and_check_context("list", [1, 2, 3, 4])
+        self.set_and_check_context("list", [1, 2, 3, 4], list)
 
     def test_set_dict(self):
         self.set_and_check_context(
-            "dict", {"thing": "stuff", "other thing": "more stuff"}
+            "dict", {"thing": "stuff", "other thing": "more stuff"}, dict
         )
 
+    def test_set_not_legal_json(self):
+        self.set_and_check_context("badjson", b"{", str)
+
     def test_exists_dir(self):
         with flux.kvs.get_dir(self.f) as kd:
             kd.mkdir("pytestdir")
@@ -79,7 +95,7 @@ def test_commit_flags(self):
         self.assertTrue(flux.kvs.exists(self.f, "flagcheck"))
 
     def test_remove(self):
-        kd = self.set_and_check_context("todel", "things to delete")
+        kd = self.set_and_check_context("todel", "things to delete", str)
         del kd["todel"]
         kd.commit()
         with self.assertRaises(KeyError):
@@ -96,7 +112,7 @@ def test_fill(self):
             self.assertEqual(kd["dir"]["other_thing"], "dirstuff")
 
     def test_set_deep(self):
-        self.set_and_check_context("a.b.c.e.f.j.k", 5)
+        self.set_and_check_context("a.b.c.e.f.j.k", 5, int)
 
     def test_bad_init(self):
         with self.assertRaises(ValueError):