From 3c34ff932be757b186cd2367bff9649a91ccb11a Mon Sep 17 00:00:00 2001
From: Jason Munro <jmunro@lbl.gov>
Date: Wed, 30 Aug 2023 16:15:26 -0700
Subject: [PATCH] Factor out generic json schema

---
 monty/json.py | 29 ++++++++++++-----------------
 1 file changed, 12 insertions(+), 17 deletions(-)

diff --git a/monty/json.py b/monty/json.py
index bfb05832c..4886b786f 100644
--- a/monty/json.py
+++ b/monty/json.py
@@ -262,6 +262,11 @@ def validate_monty_v2(cls, __input_value, _):
         """
         return cls._validate_monty(__input_value)
 
+    @classmethod
+    def __get_validators__(cls):
+        """Return validators for use in pydantic"""
+        yield cls.validate_monty_v1
+
     @classmethod
     def __get_pydantic_core_schema__(cls, source_type, handler):
         """
@@ -279,9 +284,7 @@ def __get_pydantic_core_schema__(cls, source_type, handler):
         )
 
     @classmethod
-    def __get_pydantic_json_schema__(cls, core_schema, handler):
-        """JSON schema for MSONable pattern"""
-
+    def _generic_json_schema(cls):
         return {
             "type": "object",
             "properties": {
@@ -293,24 +296,16 @@ def __get_pydantic_json_schema__(cls, core_schema, handler):
         }
 
     @classmethod
-    def __get_validators__(cls):
-        """Return validators for use in pydantic"""
-        yield cls.validate_monty_v1
+    def __get_pydantic_json_schema__(cls, core_schema, handler):
+        """JSON schema for MSONable pattern"""
+        return cls._generic_json_schema()
+        
 
     @classmethod
     def __modify_schema__(cls, field_schema):
         """JSON schema for MSONable pattern"""
-        field_schema.update(
-            {
-                "type": "object",
-                "properties": {
-                    "@class": {"enum": [cls.__name__], "type": "string"},
-                    "@module": {"enum": [cls.__module__], "type": "string"},
-                    "@version": {"type": "string"},
-                },
-                "required": ["@class", "@module"],
-            }
-        )
+        custom_schema = cls._generic_json_schema()
+        field_schema.update(custom_schema)
 
 
 class MontyEncoder(json.JSONEncoder):