Merge pull request #507 from davidwaroquiers/writable_json_store

Added option for writable JSONStores (for single JSON files only).

src/maggma/stores/mongolike.py

@@ -606,13 +606,26 @@ class JSONStore(MemoryStore):
     A Store for access to a single or multiple JSON files
     """
 
-    def __init__(self, paths: Union[str, List[str]], **kwargs):
+    def __init__(self, paths: Union[str, List[str]], file_writable=False, **kwargs):
         """
         Args:
             paths: paths for json files to turn into a Store
+            file_writable: whether this JSONStore is "file-writable". When a JSONStore
+                is "file-writable", the json file will be automatically updated
+                everytime a write-like operation is performed. Note that only
+                JSONStore with a single JSON file is compatible with file_writable
+                True. Note also that when file_writable is False, the JSONStore
+                can still apply MongoDB-like writable operations (e.g. an update)
+                as it behaves like a MemoryStore, but it will not write those changes
+                to the file.
         """
         paths = paths if isinstance(paths, (list, tuple)) else [paths]
         self.paths = paths
+        if file_writable and len(paths) > 1:
+            raise RuntimeError(
+                'Cannot instantiate file-writable JSONStore with multiple JSON files.'
+            )
+        self.file_writable = file_writable
         self.kwargs = kwargs
         super().__init__(collection_name="collection", **kwargs)
 
@@ -629,6 +642,46 @@ def connect(self, force_reset=False):
                 objects = [objects] if not isinstance(objects, list) else objects
                 self.update(objects)
 
+    def update(self, docs: Union[List[Dict], Dict], key: Union[List, str, None] = None):
+        """
+        Update documents into the Store.
+
+        For a file-writable JSONStore, the json file is updated.
+
+        Args:
+            docs: the document or list of documents to update
+            key: field name(s) to determine uniqueness for a
+                 document, can be a list of multiple fields,
+                 a single field, or None if the Store's key
+                 field is to be used
+        """
+        super().update(docs=docs, key=key)
+        if self.file_writable:
+            self.update_json_file()
+
+    def remove_docs(self, criteria: Dict):
+        """
+        Remove docs matching the query dictionary.
+
+        For a file-writable JSONStore, the json file is updated.
+
+        Args:
+            criteria: query dictionary to match
+        """
+        super().remove_docs(criteria=criteria)
+        if self.file_writable:
+            self.update_json_file()
+
+    def update_json_file(self):
+        """
+        Updates the json file when a write-like operation is performed.
+        """
+        with zopen(self.paths[0], 'w') as f:
+            data = [d for d in self.query()]
+            for d in data:
+                d.pop('_id')
+            json.dump(data, f)
+
     def __hash__(self):
         return hash((*self.paths, self.last_updated_field))
 

tests/stores/test_mongolike.py

@@ -1,11 +1,15 @@
 import os
+import shutil
 from datetime import datetime
+from unittest import mock
 
 import mongomock.collection
+from monty.tempfile import ScratchDir
 import pymongo.collection
 import pytest
 from pymongo.errors import ConfigurationError, DocumentTooLarge, OperationFailure
 
+import maggma.stores
 from maggma.core import StoreError
 from maggma.stores import JSONStore, MemoryStore, MongoStore, MongoURIStore
 from maggma.validators import JSONSchemaValidator
@@ -278,6 +282,41 @@ def test_json_store_load(jsonstore, test_dir):
     assert len(list(jsonstore.query())) == 20
 
 
+def test_json_store_writeable(test_dir):
+    with ScratchDir("."):
+        shutil.copy(test_dir / "test_set" / "d.json", ".")
+        jsonstore = JSONStore("d.json", file_writable=True)
+        jsonstore.connect()
+        assert jsonstore.count() == 2
+        jsonstore.update({'new': 'hello', 'task_id': 2})
+        assert jsonstore.count() == 3
+        jsonstore.close()
+        jsonstore = JSONStore("d.json", file_writable=True)
+        jsonstore.connect()
+        assert jsonstore.count() == 3
+        jsonstore.remove_docs({'a': 5})
+        assert jsonstore.count() == 2
+        jsonstore.close()
+        jsonstore = JSONStore("d.json", file_writable=True)
+        jsonstore.connect()
+        assert jsonstore.count() == 2
+        jsonstore.close()
+        with mock.patch("maggma.stores.JSONStore.update_json_file") as update_json_file_mock:
+            jsonstore = JSONStore("d.json", file_writable=False)
+            jsonstore.connect()
+            jsonstore.update({'new': 'hello', 'task_id': 5})
+            assert jsonstore.count() == 3
+            jsonstore.close()
+            update_json_file_mock.assert_not_called()
+        with mock.patch("maggma.stores.JSONStore.update_json_file") as update_json_file_mock:
+            jsonstore = JSONStore("d.json", file_writable=False)
+            jsonstore.connect()
+            jsonstore.remove_docs({'task_id': 5})
+            assert jsonstore.count() == 2
+            jsonstore.close()
+            update_json_file_mock.assert_not_called()
+
+
 def test_eq(mongostore, memorystore, jsonstore):
     assert mongostore == mongostore
     assert memorystore == memorystore

tests/test_files/test_set/d.json

@@ -0,0 +1 @@
+[{"a": 1, "b": [10, 20], "task_id": 0}, {"a": 5, "b": [0, 1], "c": "hello", "task_id": 1}]