Merge pull request #376 from mcgalcode/bugfix/datetime-inputs

Also allow_bson in job serialization
materialsproject · Jul 27, 2023 · c626e16 · c626e16
2 parents 53f0c76 + 409861b
commit c626e16
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 2 deletions.
diff --git a/src/jobflow/core/job.py b/src/jobflow/core/job.py
@@ -1068,7 +1068,7 @@ def as_dict(self) -> dict:
 
         # fireworks can't serialize functions and classes, so explicitly serialize to
         # the job recursively using monty to avoid issues
-        return jsanitize(d, strict=True, enum_values=True)
+        return jsanitize(d, strict=True, enum_values=True, allow_bson=True)
 
     def __setattr__(self, key, value):
         """Handle setting attributes. Implements a special case for job name."""

diff --git a/tests/core/test_job.py b/tests/core/test_job.py
@@ -558,12 +558,15 @@ class MySchema(BaseModel):
 
 def test_serialization():
     import json
+    from datetime import datetime
 
     from monty.json import MontyDecoder, MontyEncoder
 
     from jobflow import Job
 
-    test_job = Job(function=add, function_args=(1,), function_kwargs={"b": 2})
+    test_job = Job(
+        function=add, function_args=(1, datetime.now()), function_kwargs={"b": 2}
+    )
 
     uuid = test_job.uuid