change logic so ModelField of type ENUM requires type_enum to be set

Author: mslaursen

fastapi_forge/data_type_registry.py

@@ -58,12 +58,17 @@ def get(self, enum_name: str) -> DataTypeInfo:
     def all(self) -> list[DataTypeInfo]:
         return list(self._registry.values())
 
+    def __repr__(self) -> str:
+        return str(list(self._registry.values()))
 
-registry = DataTypeInfoRegistry()
+
+# enums are dynamically registered when a `CustomEnum` model is instantiated
+# and should not be registered manually
 enum_registry = DataTypeInfoEnumRegistry()
 
-faker_placeholder = "factory.Faker({placeholder})"
 
+registry = DataTypeInfoRegistry()
+faker_placeholder = "factory.Faker({placeholder})"
 
 registry.register(
     FieldDataTypeEnum.STRING,

fastapi_forge/dtos.py

@@ -62,6 +62,7 @@ class CustomEnum(_Base):
 
     def __init__(self, **kwargs: Any):
         super().__init__(**kwargs)
+        # dynamically register in the enum registry on instantiation
         enum_repr = f"enums.{self.name}"
         enum_registry.register(
             self.name,
@@ -104,8 +105,7 @@ class ModelField(_Base):
     """Represents a field in a model with validation and computed properties."""
 
     name: FieldName
-    # one of `type` or `type_enum` HAS to be set
-    type: FieldDataTypeEnum | None = None
+    type: FieldDataTypeEnum
     type_enum: CustomEnum | None = None
     primary_key: bool = False
     nullable: bool = False
@@ -125,19 +125,26 @@ def name_cc(self) -> str:
     @computed_field
     @property
     def type_info(self) -> DataTypeInfo:
-        if self.type:
-            return registry.get(self.type)
         if self.type_enum:
             return enum_registry.get(self.type_enum.name)
-
-        msg = f"Neither `type` nor `type_enum` was set for model '{self.name}'"
-        raise ValueError(msg)
+        return registry.get(self.type)
 
     @model_validator(mode="after")
     def _validate_type(self) -> Self:
-        if sum([self.type is None, self.type_enum is None]) != 1:
-            msg = "Exactly one of the fields 'type' or 'type_enum' has to be set."
+        if self.type == FieldDataTypeEnum.ENUM and self.type_enum is None:
+            msg = (
+                f"ModelField '{self.name}' has field type 'ENUM', "
+                "but is missing 'type_enum'."
+            )
+            raise ValueError(msg)
+
+        if self.type_enum and self.type != FieldDataTypeEnum.ENUM:
+            msg = (
+                f"ModelField '{self.name}' has 'type_enum' set, "
+                "but is not field type 'ENUM'."
+            )
             raise ValueError(msg)
+
         return self
 
     @model_validator(mode="after")

fastapi_forge/enums.py

@@ -10,7 +10,7 @@ class FieldDataTypeEnum(StrEnum):
     DATETIME = "DateTime"
     UUID = "UUID"
     JSONB = "JSONB"
-    CUSTOM_ENUM = "ENUM"
+    ENUM = "ENUM"
 
     @classmethod
     @lru_cache

fastapi_forge/frontend/panels/model_panel.py

@@ -3,6 +3,7 @@
 from nicegui import ui
 from pydantic import ValidationError
 
+from fastapi_forge.data_type_registry import enum_registry
 from fastapi_forge.frontend import ModelCreate, ModelRow
 from fastapi_forge.frontend.constants import SELECTED_MODEL_TEXT_COLOR
 from fastapi_forge.frontend.notifications import notify_validation_error
@@ -23,6 +24,7 @@ def _build(self) -> None:
         with self, ui.column().classes("items-align content-start w-full"):
             ModelCreate()
             self._render_models()
+            ui.button(text="YO", on_click=lambda: print(enum_registry))
 
             ui.button(
                 "Export",

fastapi_forge/jinja.py

@@ -481,6 +481,7 @@ def render_custom_enums_to_enums(custom_enums: list[CustomEnum]) -> str:
             ),
             ModelField(
                 name="my_enum",
+                type=FieldDataTypeEnum.ENUM,
                 type_enum=enum0,
             ),
         ],

tests/test_dtos.py

@@ -84,17 +84,13 @@ def test_factory_field_value(
     assert model_field.type_info.faker_field_value == expected_factory_value
 
 
-def test_type_not_set() -> None:
+def test_type_missing_type_enum() -> None:
     with pytest.raises(ValidationError) as exc_info:
-        ModelField(
-            name="test",
-        )
-    assert "Exactly one of the fields 'type' or 'type_enum' has to be set." in str(
-        exc_info.value
-    )
+        ModelField(name="test", type=FieldDataTypeEnum.ENUM)
+    assert "has field type 'ENUM'" in str(exc_info.value)
 
 
-def test_type_both_set() -> None:
+def test_type_incorrect_type() -> None:
     with pytest.raises(ValidationError) as exc_info:
         ModelField(
             name="test",
@@ -109,9 +105,7 @@ def test_type_both_set() -> None:
                 ],
             ),
         )
-    assert "Exactly one of the fields 'type' or 'type_enum' has to be set." in str(
-        exc_info.value
-    )
+    assert "but is not field type 'ENUM'" in str(exc_info.value)
 
 
 ###############################