Merge branch 'master' of github.com:materialsvirtuallab/monty

materialsvirtuallab · Sep 7, 2023 · 7e52e52 · 7e52e52
2 parents 5960521 + dc61002
commit 7e52e52
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 7 deletions.
diff --git a/monty/json.py b/monty/json.py
@@ -277,11 +277,7 @@ def __get_pydantic_core_schema__(cls, source_type, handler):
 
         s = core_schema.general_plain_validator_function(cls.validate_monty_v2)
 
-        return core_schema.json_or_python_schema(
-            json_schema=s,
-            python_schema=s,
-            serialization=core_schema.plain_serializer_function_ser_schema(lambda instance: instance.as_dict()),
-        )
+        return core_schema.json_or_python_schema(json_schema=s, python_schema=s)
 
     @classmethod
     def _generic_json_schema(cls):

diff --git a/tests/test_json.py b/tests/test_json.py
@@ -652,7 +652,7 @@ class ModelWithMSONable(BaseModel):
         test_dict_object = ModelWithMSONable(a=test_object.a.as_dict())
         assert test_dict_object.a.a == test_object.a.a
 
-        assert test_object.schema() == {
+        assert test_object.model_json_schema() == {
             "title": "ModelWithMSONable",
             "type": "object",
             "properties": {
@@ -670,7 +670,7 @@ class ModelWithMSONable(BaseModel):
             "required": ["a"],
         }
 
-        d = jsanitize(test_object, strict=True)
+        d = jsanitize(test_object, strict=True, enum_values=True, allow_bson=True)
         assert d == {
             "a": {
                 "@module": "tests.test_json",