add nested json read/write

mpstore/store.py

@@ -18,20 +18,41 @@ def read_store(key: str) -> any:
     _create_store_if_not_exists()
     with open(STORE_PATH) as f:
         json_content = json.load(f)
-        return json_content.get(key)  # if key does not exist, returns None
+        return _get_nested_json_value(json_content, key)  # if key does not exist, returns None
 
 
 def write_store(key: str, value: any):
     """Writes the key, value pair in json store, overwriting any previous value."""
     _create_store_if_not_exists()
     with open(STORE_PATH, "rb") as f:
         json_content = json.load(f)
-        json_content[key] = value
+        _set_nested_json_value(json_content, key, value)
     with open(TEMP_STORE_PATH, "wb") as temp_f:
         json.dump(json_content, temp_f)  # write to temp file
     os.rename(TEMP_STORE_PATH, STORE_PATH)  # atomic operation
 
 
+def _get_nested_json_value(data: dict, key: str) -> any:
+    """Gets the value from a nested dict given a dot-separated key."""
+    keys = key.split(".")
+    for k in keys:
+        data = data.get(k, None)
+        if data is None:
+            return None
+    return data
+
+
+def _set_nested_json_value(data: dict, key: str, value: any):
+    """Sets a value in a nested dict given a dot-separated key."""
+    keys = key.split(".")
+    d = data
+    for k in keys[:-1]:
+        if k not in d:
+            d[k] = {}
+        d = d[k]
+    d[keys[-1]] = value
+
+
 def _create_store_if_not_exists():
     try:
         with open(STORE_PATH):

test.py

@@ -1,5 +1,6 @@
 import unittest
 import os
+import json
 
 from mpstore import load_store, read_store, write_store
 from mpstore.store import STORE_PATH, TEMP_STORE_PATH
@@ -51,6 +52,23 @@ def test_store_read(self):
 
         assert read_store("foo") == "bar"
 
+    def test_nested_write_and_read(self):
+        write_store("item1.nested1", 2)
+        write_store("item1.nested2", "value2")
+
+        assert read_store("item1.nested1") == 2
+        assert read_store("item1.nested2") == "value2"
+
+    def test_nested_read_nonexistent_key(self):
+        write_store("item1.nested1", 2)
+
+        assert read_store("item1.nested3") is None
+
+    def test_deeply_nested_write_and_read(self):
+        write_store("level1.level2.level3", "deep_value")
+
+        assert read_store("level1.level2.level3") == "deep_value"
+
     def _assert_store_exists_and_temp_store_dne(self):
         assert os.stat(STORE_PATH)  # will raise error if does not exist
         with self.assertRaises(OSError):