Merge pull request #930 from Andrew-S-Rosen/patch-1

Enable `recursive_msonable` in `jsanitize` calls
materialsproject · Mar 17, 2024 · 45f72fd · 45f72fd
2 parents 91872e5 + 844b309
commit 45f72fd
Show file tree

Hide file tree

Showing 4 changed files with 4 additions and 4 deletions.
diff --git a/src/maggma/cli/distributed.py b/src/maggma/cli/distributed.py
@@ -135,7 +135,7 @@ def manager(url: str, port: int, builders: List[Builder], num_chunks: int, num_w
                 if not chunk_dict["distributed"]:
                     temp_builder_dict = dict(**builder_dict)
                     temp_builder_dict.update(chunk_dict["chunk"])  # type: ignore
-                    temp_builder_dict = jsanitize(temp_builder_dict)
+                    temp_builder_dict = jsanitize(temp_builder_dict, recursive_msonable=True)
 
                     # Send work for available workers
                     for identity in workers:

diff --git a/src/maggma/cli/rabbitmq.py b/src/maggma/cli/rabbitmq.py
@@ -136,7 +136,7 @@ def manager(
                 if not chunk_dict["distributed"]:
                     temp_builder_dict = dict(**builder_dict)
                     temp_builder_dict.update(chunk_dict["chunk"])  # type: ignore
-                    temp_builder_dict = jsanitize(temp_builder_dict)
+                    temp_builder_dict = jsanitize(temp_builder_dict, recursive_msonable=True)
 
                     # Send work for available workers
                     for identity in workers:

diff --git a/src/maggma/stores/gridfs.py b/src/maggma/stores/gridfs.py
@@ -388,7 +388,7 @@ def update(
                 if has(d, k)
             }
             metadata.update(search_doc)
-            data = json.dumps(jsanitize(d)).encode("UTF-8")
+            data = json.dumps(jsanitize(d, recursive_msonable=True)).encode("UTF-8")
             if self.compression:
                 data = zlib.compress(data)
                 metadata["compression"] = "zlib"

diff --git a/src/maggma/stores/mongolike.py b/src/maggma/stores/mongolike.py
@@ -384,7 +384,7 @@ def update(self, docs: Union[List[Dict], Dict], key: Union[List, str, None] = No
         if not isinstance(docs, list):
             docs = [docs]
 
-        for d in (jsanitize(x, allow_bson=True) for x in docs):
+        for d in (jsanitize(x, allow_bson=True, recursive_msonable=True) for x in docs):
             # document-level validation is optional
             validates = True
             if self.validator: