feat: support custom postgres enums

Author: mslaursen

fastapi_forge/dtos.py

@@ -15,6 +15,10 @@
 
 BoundedStr = Annotated[str, Field(..., min_length=1, max_length=100)]
 SnakeCaseStr = Annotated[BoundedStr, Field(..., pattern=r"^[a-z][a-z0-9_]*$")]
+PascalCaseStr = Annotated[
+    BoundedStr,
+    Field(..., pattern=r"^[A-Z][a-zA-Z0-9]*$"),
+]
 ModelName = SnakeCaseStr
 FieldName = SnakeCaseStr
 BackPopulates = Annotated[str, Field(..., pattern=r"^[a-z][a-z0-9_]*$")]
@@ -36,6 +40,51 @@ class ModelFieldMetadata(_Base):
     is_foreign_key: bool = False
 
 
+class CustomEnumValue(_Base):
+    """Represents a single name/value pair in a custom enum."""
+
+    name: Annotated[
+        BoundedStr,
+        Field(
+            ...,
+            pattern=r"^[a-zA-Z][a-zA-Z0-9_]*$",
+        ),
+    ]
+    value: BoundedStr
+
+
+class CustomEnum(_Base):
+    """Represents a custom PostgreSQL ENUM type."""
+
+    name: PascalCaseStr
+    values: Annotated[list[CustomEnumValue], Field(..., min_length=1)]
+
+    @model_validator(mode="after")
+    def _validate_enum(self) -> Self:
+        names = [v.name for v in self.values]
+        values = [v.value for v in self.values]
+
+        if len(names) != len(set(names)):
+            raise ValueError(f"Enum '{self.name}' has duplicate names.")
+        if len(values) != len(set(values)):
+            raise ValueError(f"Enum '{self.name}' has duplicate values.")
+        return self
+
+    @computed_field
+    @property
+    def class_definition(self) -> str:
+        """Returns a string representing the Python Enum class definition."""
+        lines: list[str] = []
+        lines.extend([f"class {self.name}(StrEnum):"])
+
+        value_lines: list[str] = []
+        for v in self.values:
+            value_lines.extend([f'    {v.name} = "{v.value}"'])
+
+        lines.extend(value_lines)
+        return "\n".join(lines)
+
+
 class ModelField(_Base):
     """Represents a field in a model with validation and computed properties."""
 

tests/test_dtos.py

@@ -3,7 +3,14 @@
 import pytest
 from pydantic import ValidationError
 
-from fastapi_forge.dtos import Model, ModelField, ModelRelationship, ProjectSpec
+from fastapi_forge.dtos import (
+    CustomEnum,
+    CustomEnumValue,
+    Model,
+    ModelField,
+    ModelRelationship,
+    ProjectSpec,
+)
 from fastapi_forge.enums import FieldDataTypeEnum
 
 ########################
@@ -126,3 +133,45 @@ def test_project_spec_non_existing_target_model() -> None:
         "'restaurant' has a relationship to 'non_existing', which does not exist."
         in str(exc_info.value)
     )
+
+
+##############
+# CustomEnum #
+##############
+
+
+def test_custom_enum_not_unique_values() -> None:
+    with pytest.raises(ValidationError) as exc_info:
+        CustomEnum(
+            name="MyEnum",
+            values=[
+                CustomEnumValue(name="HELLO", value="hello"),
+                CustomEnumValue(name="HI", value="hello"),
+            ],
+        )
+    assert "Enum 'MyEnum' has duplicate values." in str(exc_info.value)
+
+
+def test_custom_enum_not_unique_names() -> None:
+    with pytest.raises(ValidationError) as exc_info:
+        CustomEnum(
+            name="MyEnum",
+            values=[
+                CustomEnumValue(name="HELLO", value="hello"),
+                CustomEnumValue(name="HELLO", value="hi"),
+            ],
+        )
+    assert "Enum 'MyEnum' has duplicate names." in str(exc_info.value)
+
+
+def test_custom_enum_valid() -> None:
+    enum = CustomEnum(
+        name="MyEnum",
+        values=[
+            CustomEnumValue(name="FoO", value="foo"),
+            CustomEnumValue(name="BAR", value="bar"),
+        ],
+    )
+    assert enum.class_definition == (
+        'class MyEnum(StrEnum):\n    FoO = "foo"\n    BAR = "bar"'
+    )