fix: generate proper default values when reusing Enums types (#796)

marcosschroh · Nov 13, 2024 · 78eab46 · 78eab46
1 parent 081d897
commit 78eab46
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 0 deletions.
diff --git a/dataclasses_avroschema/model_generator/lang/python/base.py b/dataclasses_avroschema/model_generator/lang/python/base.py
@@ -126,6 +126,11 @@ def get_field_default(self) -> str:
             python_type = func(self.default)
             template_func = avro_to_python_utils.LOGICAL_TYPE_TEMPLATES[self.avro_type]
             default_repr = template_func(python_type)
+        elif self.avro_type == field_utils.NULL:
+            default_repr = "None"
+        elif self.avro_type == self.type_hint:
+            # this is when a Custom type is been reused which is probably an Enum
+            default_repr = f"{casefy.pascalcase(self.type_hint)}.{casefy.uppercase(self.default)}"
         else:
             default_repr = str(self.default)
 

diff --git a/tests/model_generator/conftest.py b/tests/model_generator/conftest.py
@@ -269,6 +269,11 @@ def schema_with_enum_types() -> Dict:
                 ],
                 "default": "MIN_LIMIT",
             },
+            {
+                "name": "another_limit_type",
+                "type": "LimitTypes",
+                "default": "EXACT_LIMIT",
+            },
         ],
     }
 

diff --git a/tests/model_generator/test_model_generator.py b/tests/model_generator/test_model_generator.py
@@ -246,6 +246,7 @@ class User(AvroModel):
     superheros: Superheros = Superheros.BATMAN
     cars: typing.Optional[Cars] = None
     limit_type: typing.Optional[LimitTypes] = LimitTypes.MIN_LIMIT
+    another_limit_type: LimitTypes = LimitTypes.EXACT_LIMIT
 """
     model_generator = ModelGenerator()
     result = model_generator.render(schema=schema_with_enum_types)