chore: Replace dataclasses and dataclasses_json with pydantic

Author: kulikthebird

packages/cw-schema-codegen/playground/playground.py

@@ -1,84 +1,46 @@
-from dataclasses import dataclass, field
-from dataclasses_json import dataclass_json, config
-from typing import Optional, Iterable
 import sys
-import json
+from typing import Literal, Union, Tuple
+from pydantic import BaseModel, RootModel
 
 
-# TODO tkulik: try to get rid of the `dataclasses_json` dependency
+class SomeEnum(RootModel):
+    class Field1(RootModel[Literal['Field1']]):
+        pass
 
+    class Field2(BaseModel):
+        Field2: Tuple[int, int]
 
-enum_field = lambda: field(default=None, metadata=config(exclude=lambda x: x is None))
+    class Field3(BaseModel):
+        class __InnerStruct(BaseModel):
+            a: str
+            b: int
+        Field3: __InnerStruct
 
-@dataclass_json
-@dataclass
-class SomeEnum:
-    class VariantIndicator:
-        pass
+    class Field4(BaseModel):
+        Field4: 'SomeEnum'
 
-    class Field3Type:
-        a: str
-        b: int
-
-    class Field5Type:
-        a: Iterable['SomeEnum']
-
-    Field1: Optional[VariantIndicator] = enum_field()
-    Field2: Optional[tuple[int, int]] = enum_field()
-    Field3: Optional[Field3Type] = enum_field()
-    Field4: Optional[Iterable['SomeEnum']] = enum_field()
-    Field5: Optional[Field5Type] = enum_field()
-    
-    def deserialize(json):
-        if not ":" in json:
-            if json == '"Field1"':
-                return SomeEnum(Field1=SomeEnum.VariantIndicator())
-            else:
-                raise Exception(f"Deserialization error, undefined variant: {json}")
-        else:
-            return SomeEnum.from_json(json)
-        
-    def serialize(self):
-        if self.Field1 is not None:
-            return '"Field1"'
-        else:
-            return SomeEnum.to_json(self)
-        
-@dataclass_json
-@dataclass
-class UnitStructure:
-    def deserialize(json):
-        if json == "null":
-            return UnitStructure()
-        else:
-            Exception(f"Deserialization error, undefined value: {json}")
-        
-    def serialize(self):
-        return 'null'
-
-@dataclass_json
-@dataclass
-class TupleStructure:
-    Tuple: tuple[int, str, int]
-
-    def deserialize(json):
-        return TupleStructure.from_json(f'{{ "Tuple": {json} }}')
-        
-    def serialize(self):
-        return json.dumps(self.Tuple)
-
-@dataclass_json
-@dataclass
-class NamedStructure:
+    class Field5(BaseModel):
+        class __InnerStruct(BaseModel):
+            a: 'SomeEnum'
+        Field5: __InnerStruct
+
+    root: Union[Field1, Field2, Field3, Field4, Field5]
+
+
+class UnitStructure(RootModel):
+    root: None
+
+
+class TupleStructure(RootModel):
+    root: Tuple[int, str, int]
+
+
+class NamedStructure(BaseModel):
     a: str
     b: int
-    c: Iterable['SomeEnum']
+    c: SomeEnum
+
 
-    def deserialize(json):
-        return NamedStructure.from_json(json)
-        
-    def serialize(self):
-        return self.to_json()
 
 ###
 ### TESTS:
@@ -88,32 +50,36 @@ def serialize(self):
     input = input.rstrip()
     try:
         if index < 5:
-            deserialized = SomeEnum.deserialize(input)
+            deserialized = SomeEnum.model_validate_json(input)
         elif index == 5:
-            deserialized = UnitStructure.deserialize(input)
+            deserialized = UnitStructure.model_validate_json(input)
         elif index == 6:
-            deserialized = TupleStructure.deserialize(input)
+            deserialized = TupleStructure.model_validate_json(input)
         else:
-            deserialized = NamedStructure.deserialize(input)
+            deserialized = NamedStructure.model_validate_json(input)
     except:
         raise(Exception(f"This json can't be deserialized: {input}"))
-    serialized = deserialized.serialize()
+    serialized = deserialized.model_dump_json()
     print(serialized)
 
 
 # def handle_msg(json):
-#     a = SomeEnum.deserialize(json)
-#     if a.Field1 is not None:
+#     a = SomeEnum.model_validate_json(json)
+#     if isinstance(a.root, SomeEnum.Field1):
 #         print("SomeEnum::Field1")
-#     elif a.Field2 is not None:
-#         print(a.Field2[0])
-#         print(a.Field2[1])
-#     elif a.Field3 is not None:
-#         print(a.Field3)
-#     elif a.Field4 is not None:
-#         print(a.Field4)
-#     elif a.Field5 is not None:
-#         print(a.Field5)
+#     elif isinstance(a.root, SomeEnum.Field2):
+#         print(a.root.Field2[0])
+#         print(a.root.Field2[1])
+#     elif isinstance(a.root, SomeEnum.Field3):
+#         print(a.root.Field3)
+#     elif isinstance(a.root, SomeEnum.Field4):
+#         print(a.root.Field4)
+#     elif isinstance(a.root, SomeEnum.Field5):
+#         print(a.root.Field5)
 
 # handle_msg('"Field1"')
-# handle_msg('{"Field2": [10, 12]}')
+# handle_msg('{"Field2": [10, 12]}')
+# handle_msg('{"Field3": { "a": "10", "b": 12 } }')
+# handle_msg('{"Field4": { "Field4": "Field1" } }')
+# handle_msg('{"Field5": { "a": "Field1" } }')
+# handle_msg('{"Field5": { "a": { "Field5": { "a": "Field1" } } } }')

packages/cw-schema-codegen/templates/python/enum.tpl.py

@@ -1,44 +1,38 @@
 # This code is @generated by cw-schema-codegen. Do not modify this manually.
-from dataclasses import dataclass, field
-from dataclasses_json import dataclass_json, config
-from typing import Optional, Iterable
 
-enum_field = lambda: field(default=None, metadata=config(exclude=lambda x: x is None))
+import typing
+from pydantic import BaseModel, RootModel
 
-@dataclass_json
-@dataclass
-class {{ name }}:
-    '''{% for doc in docs %}
+class {{ name }}(RootModel):
+    """{% for doc in docs %}
     {{ doc }}
-    {% endfor %}'''
-
-    class VariantIndicator:
-        '''
-        This structure is an indicator of the simple enum variant that is currently contained
-        in this enum structure. It's used only for the enum variants that does not contain
-        any inner structure. It's constructed automatically and it's not intend to be manually
-        used by the user.
-        '''
-        pass
-
+    {% endfor %}"""
 
 {% for variant in variants %}
-    '''{% for doc in variant.docs %}
-    {{ doc }}
-    {% endfor %}'''
-    {% match variant.ty %}
-        {% when TypeTemplate::Unit %}
-            {{ variant.name }}: Optional[VariantIndicator] = enum_field()
-        {% when TypeTemplate::Tuple with (types) %}
-            {{ variant.name }}: Optional[tuple[{{ types|join(", ") }}]] = enum_field()
-        {% when TypeTemplate::Named with { fields } %}
-            class {{ variant.name }}Type:
-                {% for field in fields %}
-                    '''{% for doc in field.docs %}
-                    # {{ doc }}
-                    {% endfor %}'''
-                    {{ field.name }}: {{ field.ty }}
-                {% endfor %}
-            {{ variant.name }}: Optional[{{ variant.name }}Type] = enum_field()
-    {% endmatch %}
+{% match variant.ty %}
+{% when TypeTemplate::Unit %}
+    class {{ variant.name }}(RootModel):
+        """{% for doc in variant.docs %}
+        {{ doc }}
+        {% endfor %}"""
+        root: None
+{% when TypeTemplate::Tuple with (types) %}
+    class {{ variant.name }}(BaseModel):
+        """{% for doc in variant.docs %}
+        {{ doc }}
+        {% endfor %}"""
+        {{ variant.name }}: typing.Tuple[{{ types|join(", ") }}]
+{% when TypeTemplate::Named with { fields } %}
+    class __Inner:
+        """{% for doc in variant.docs %}
+        {{ doc }}
+        {% endfor %}"""
+        {% for field in fields %}
+        {{ field.name }}: {{ field.ty }}
+        """{% for doc in field.docs %}
+        # {{ doc }}
+        {% endfor %}"""
+        {% endfor %}
+        {{ variant.name }}: __Inner
+{% endmatch %}
 {% endfor %}

packages/cw-schema-codegen/templates/python/struct.tpl.py

@@ -1,55 +1,31 @@
 # This code is @generated by cw-schema-codegen. Do not modify this manually.
 
-/**
-{% for doc in docs %}
-    * {{ doc }}
-{% endfor %}
- */
-
-type {{ name }} = 
-{% match ty %}
-    {% when TypeTemplate::Unit %}
-        void
-    {% when TypeTemplate::Tuple with (types) %}
-        [{{ types|join(", ") }}]
-    {% when TypeTemplate::Named with { fields } %}
-        {
-            {% for field in fields %}
-                /**
-                {% for doc in field.docs %}
-                    * {{ doc }}
-                {% endfor %}
-                 */
-
-                {{ field.name }}: {{ field.ty }};
-            {% endfor %}
-        }
-{% endmatch %}
-
+import typing
+from pydantic import BaseModel, RootModel
 
 
-# This code is @generated by cw-schema-codegen. Do not modify this manually.
-from dataclasses import dataclass, field
-from dataclasses_json import dataclass_json, config
-from typing import Optional, Iterable
-
-@dataclass_json
-@dataclass
-class {{ name }}:
+{% match ty %}
+{% when TypeTemplate::Unit %}
+class {{ name }}(RootModel):
     '''{% for doc in docs %}
     {{ doc }}
     {% endfor %}'''
-
-    {% match ty %}
-        {% when TypeTemplate::Unit %}
-            pass
-        {% when TypeTemplate::Tuple with (types) %}
-            {{ variant.name }}: tuple[{{ types|join(", ") }}]
-        {% when TypeTemplate::Named with { fields } %}
-            {% for field in fields %}
-                '''{% for doc in field.docs %}
-                # {{ doc }}
-                {% endfor %}'''
-                {{ field.name }}: {{ field.ty }}
-            {% endfor %}
-    {% endmatch %}
+        root: None
+{% when TypeTemplate::Tuple with (types) %}
+class {{ name }}(RootModel):
+    '''{% for doc in docs %}
+    {{ doc }}
+    {% endfor %}'''
+    root: typing.Tuple[{{ types|join(", ") }}]
+{% when TypeTemplate::Named with { fields } %}
+class {{ name }}(BaseModel):
+    '''{% for doc in docs %}
+    {{ doc }}
+    {% endfor %}'''
+    {% for field in fields %}
+    {{ field.name }}: {{ field.ty }}
+    '''{% for doc in field.docs %}
+    # {{ doc }}
+    {% endfor %}'''
+    {% endfor %}
+{% endmatch %}

packages/cw-schema-codegen/tests/snapshots/python_tpl__simple_enum.snap

@@ -4,39 +4,27 @@ expression: rendered
 snapshot_kind: text
 ---
 # This code is @generated by cw-schema-codegen. Do not modify this manually.
-from dataclasses import dataclass, field
-from dataclasses_json import dataclass_json, config
-from typing import Optional, Iterable
 
-enum_field = lambda: field(default=None, metadata=config(exclude=lambda x: x is None))
+import typing
+from pydantic import BaseModel, RootModel
 
-@dataclass_json
-@dataclass
-class Simple:
-    '''
+class Simple(RootModel):
+    """
     Simple enum
-    '''
-
-    class VariantIndicator:
-        '''
-        This structure is an indicator of the simple enum variant that is currently contained
-        in this enum structure. It's used only for the enum variants that does not contain
-        any inner structure. It's constructed automatically and it's not intend to be manually
-        used by the user.
-        '''
-        pass
-
-
-
-    '''
-    One variant
-    '''
-    
-            One: Optional[VariantIndicator] = enum_field()
-        
-
-    '''
-    Two variant
-    '''
-    
-            Two: Optional[VariantIndicator] = enum_field()
+    """
+
+
+
+    class One(RootModel):
+        """
+        One variant
+        """
+        root: None
+
+
+
+    class Two(RootModel):
+        """
+        Two variant
+        """
+        root: None