Merge pull request #1818 from dvora-h/set-file-and-set-path

JSON set_file and set_path support

2 files changed, 90 insertions, 0 deletions

diff --git a/redis/commands/json/commands.py b/redis/commands/json/commands.py
index e7f07b6..a132b8e 100644
--- a/redis/commands/json/commands.py
+++ b/redis/commands/json/commands.py
@@ -1,3 +1,6 @@
+import os
+from json import JSONDecodeError, loads
+
 from deprecated import deprecated
 
 from redis.exceptions import DataError
@@ -213,6 +216,54 @@ class JSONCommands:
             pieces.append("XX")
         return self.execute_command("JSON.SET", *pieces)
 
+    def set_file(self, name, path, file_name, nx=False, xx=False, decode_keys=False):
+        """
+        Set the JSON value at key ``name`` under the ``path`` to the content
+        of the json file ``file_name``.
+
+        ``nx`` if set to True, set ``value`` only if it does not exist.
+        ``xx`` if set to True, set ``value`` only if it exists.
+        ``decode_keys`` If set to True, the keys of ``obj`` will be decoded
+        with utf-8.
+
+        """
+
+        with open(file_name, "r") as fp:
+            file_content = loads(fp.read())
+
+        return self.set(name, path, file_content, nx=nx, xx=xx, decode_keys=decode_keys)
+
+    def set_path(self, json_path, root_folder, nx=False, xx=False, decode_keys=False):
+        """
+        Iterate over ``root_folder`` and set each JSON file to a value
+        under ``json_path`` with the file name as the key.
+
+        ``nx`` if set to True, set ``value`` only if it does not exist.
+        ``xx`` if set to True, set ``value`` only if it exists.
+        ``decode_keys`` If set to True, the keys of ``obj`` will be decoded
+        with utf-8.
+
+        """
+        set_files_result = {}
+        for root, dirs, files in os.walk(root_folder):
+            for file in files:
+                file_path = os.path.join(root, file)
+                try:
+                    file_name = file_path.rsplit(".")[0]
+                    self.set_file(
+                        file_name,
+                        json_path,
+                        file_path,
+                        nx=nx,
+                        xx=xx,
+                        decode_keys=decode_keys,
+                    )
+                    set_files_result[file_path] = True
+                except JSONDecodeError:
+                    set_files_result[file_path] = False
+
+        return set_files_result
+
     def strlen(self, name, path=None):
         """Return the length of the string JSON value under ``path`` at key
         ``name``.
diff --git a/tests/test_json.py b/tests/test_json.py
index a99547d..8bd1d77 100644
--- a/tests/test_json.py
+++ b/tests/test_json.py
@@ -1392,3 +1392,42 @@ def test_custom_decoder(client):
     assert client.exists("foo") == 0
     assert not isinstance(cj.__encoder__, json.JSONEncoder)
     assert not isinstance(cj.__decoder__, json.JSONDecoder)
+
+
+@pytest.mark.redismod
+def test_set_file(client):
+    import json
+    import tempfile
+
+    obj = {"hello": "world"}
+    jsonfile = tempfile.NamedTemporaryFile(suffix=".json")
+    with open(jsonfile.name, "w+") as fp:
+        fp.write(json.dumps(obj))
+
+    nojsonfile = tempfile.NamedTemporaryFile()
+    nojsonfile.write(b"Hello World")
+
+    assert client.json().set_file("test", Path.rootPath(), jsonfile.name)
+    assert client.json().get("test") == obj
+    with pytest.raises(json.JSONDecodeError):
+        client.json().set_file("test2", Path.rootPath(), nojsonfile.name)
+
+
+@pytest.mark.redismod
+def test_set_path(client):
+    import json
+    import tempfile
+
+    root = tempfile.mkdtemp()
+    sub = tempfile.mkdtemp(dir=root)
+    jsonfile = tempfile.mktemp(suffix=".json", dir=sub)
+    nojsonfile = tempfile.mktemp(dir=root)
+
+    with open(jsonfile, "w+") as fp:
+        fp.write(json.dumps({"hello": "world"}))
+    open(nojsonfile, "a+").write("hello")
+
+    result = {jsonfile: True, nojsonfile: False}
+    print(result)
+    assert client.json().set_path(Path.rootPath(), root) == result
+    assert client.json().get(jsonfile.rsplit(".")[0]) == {"hello": "world"}