use enum for type

core/data/tests/anonymous_struct_with_rename/output.py

@@ -3,6 +3,7 @@
 """
 from __future__ import annotations
 
+from enum import Enum
 from pydantic import BaseModel, ConfigDict, Field
 from typing import List, Literal, Union
 
@@ -33,16 +34,21 @@ class AnonymousStructWithRenameKebabCaseInner(BaseModel):
     camel_case_string_field: str = Field(alias="camelCaseStringField")
     something_else: bool = Field(alias="something-else")
 
+class AnonymousStructWithRenameTypes(str, Enum):
+    LIST = "list"
+    LONG_FIELD_NAMES = "longFieldNames"
+    KEBAB_CASE = "kebabCase"
+
 class AnonymousStructWithRenameList(BaseModel):
-    AnonymousStructWithRenameTypes: Literal["list"] = "list"
+    type: Literal[AnonymousStructWithRenameTypes.LIST] = AnonymousStructWithRenameTypes.LIST
     content: AnonymousStructWithRenameListInner
 
 class AnonymousStructWithRenameLongFieldNames(BaseModel):
-    AnonymousStructWithRenameTypes: Literal["longFieldNames"] = "longFieldNames"
+    type: Literal[AnonymousStructWithRenameTypes.LONG_FIELD_NAMES] = AnonymousStructWithRenameTypes.LONG_FIELD_NAMES
     content: AnonymousStructWithRenameLongFieldNamesInner
 
 class AnonymousStructWithRenameKebabCase(BaseModel):
-    AnonymousStructWithRenameTypes: Literal["kebabCase"] = "kebabCase"
+    type: Literal[AnonymousStructWithRenameTypes.KEBAB_CASE] = AnonymousStructWithRenameTypes.KEBAB_CASE
     content: AnonymousStructWithRenameKebabCaseInner
 
 AnonymousStructWithRename = Union[AnonymousStructWithRenameList, AnonymousStructWithRenameLongFieldNames, AnonymousStructWithRenameKebabCase]

core/data/tests/can_apply_prefix_correctly/output.py

@@ -3,35 +3,44 @@
 """
 from __future__ import annotations
 
+from enum import Enum
 from pydantic import BaseModel
 from typing import Dict, List, Literal, Union
 
 
 class ItemDetailsFieldValue(BaseModel):
     hello: str
 
+class AdvancedColorsTypes(str, Enum):
+    STRING = "String"
+    NUMBER = "Number"
+    NUMBER_ARRAY = "NumberArray"
+    REALLY_COOL_TYPE = "ReallyCoolType"
+    ARRAY_REALLY_COOL_TYPE = "ArrayReallyCoolType"
+    DICTIONARY_REALLY_COOL_TYPE = "DictionaryReallyCoolType"
+
 class AdvancedColorsString(BaseModel):
-    AdvancedColorsTypes: Literal["String"] = "String"
+    t: Literal[AdvancedColorsTypes.STRING] = AdvancedColorsTypes.STRING
     c: str
 
 class AdvancedColorsNumber(BaseModel):
-    AdvancedColorsTypes: Literal["Number"] = "Number"
+    t: Literal[AdvancedColorsTypes.NUMBER] = AdvancedColorsTypes.NUMBER
     c: int
 
 class AdvancedColorsNumberArray(BaseModel):
-    AdvancedColorsTypes: Literal["NumberArray"] = "NumberArray"
+    t: Literal[AdvancedColorsTypes.NUMBER_ARRAY] = AdvancedColorsTypes.NUMBER_ARRAY
     c: List[int]
 
 class AdvancedColorsReallyCoolType(BaseModel):
-    AdvancedColorsTypes: Literal["ReallyCoolType"] = "ReallyCoolType"
+    t: Literal[AdvancedColorsTypes.REALLY_COOL_TYPE] = AdvancedColorsTypes.REALLY_COOL_TYPE
     c: ItemDetailsFieldValue
 
 class AdvancedColorsArrayReallyCoolType(BaseModel):
-    AdvancedColorsTypes: Literal["ArrayReallyCoolType"] = "ArrayReallyCoolType"
+    t: Literal[AdvancedColorsTypes.ARRAY_REALLY_COOL_TYPE] = AdvancedColorsTypes.ARRAY_REALLY_COOL_TYPE
     c: List[ItemDetailsFieldValue]
 
 class AdvancedColorsDictionaryReallyCoolType(BaseModel):
-    AdvancedColorsTypes: Literal["DictionaryReallyCoolType"] = "DictionaryReallyCoolType"
+    t: Literal[AdvancedColorsTypes.DICTIONARY_REALLY_COOL_TYPE] = AdvancedColorsTypes.DICTIONARY_REALLY_COOL_TYPE
     c: Dict[str, ItemDetailsFieldValue]
 
 AdvancedColors = Union[AdvancedColorsString, AdvancedColorsNumber, AdvancedColorsNumberArray, AdvancedColorsReallyCoolType, AdvancedColorsArrayReallyCoolType, AdvancedColorsDictionaryReallyCoolType]

core/data/tests/can_generate_algebraic_enum/output.py

@@ -3,6 +3,7 @@
 """
 from __future__ import annotations
 
+from enum import Enum
 from pydantic import BaseModel
 from typing import List, Literal, Union
 
@@ -12,54 +13,67 @@ class ItemDetailsFieldValue(BaseModel):
     Struct comment
     """
     pass
+class AdvancedColorsTypes(str, Enum):
+    STRING = "String"
+    NUMBER = "Number"
+    UNSIGNED_NUMBER = "UnsignedNumber"
+    NUMBER_ARRAY = "NumberArray"
+    REALLY_COOL_TYPE = "ReallyCoolType"
+
 class AdvancedColorsString(BaseModel):
     """
     This is a case comment
     """
-    AdvancedColorsTypes: Literal["String"] = "String"
+    type: Literal[AdvancedColorsTypes.STRING] = AdvancedColorsTypes.STRING
     content: str
 
 class AdvancedColorsNumber(BaseModel):
-    AdvancedColorsTypes: Literal["Number"] = "Number"
+    type: Literal[AdvancedColorsTypes.NUMBER] = AdvancedColorsTypes.NUMBER
     content: int
 
 class AdvancedColorsUnsignedNumber(BaseModel):
-    AdvancedColorsTypes: Literal["UnsignedNumber"] = "UnsignedNumber"
+    type: Literal[AdvancedColorsTypes.UNSIGNED_NUMBER] = AdvancedColorsTypes.UNSIGNED_NUMBER
     content: int
 
 class AdvancedColorsNumberArray(BaseModel):
-    AdvancedColorsTypes: Literal["NumberArray"] = "NumberArray"
+    type: Literal[AdvancedColorsTypes.NUMBER_ARRAY] = AdvancedColorsTypes.NUMBER_ARRAY
     content: List[int]
 
 class AdvancedColorsReallyCoolType(BaseModel):
     """
     Comment on the last element
     """
-    AdvancedColorsTypes: Literal["ReallyCoolType"] = "ReallyCoolType"
+    type: Literal[AdvancedColorsTypes.REALLY_COOL_TYPE] = AdvancedColorsTypes.REALLY_COOL_TYPE
     content: ItemDetailsFieldValue
 
 # Enum comment
 AdvancedColors = Union[AdvancedColorsString, AdvancedColorsNumber, AdvancedColorsUnsignedNumber, AdvancedColorsNumberArray, AdvancedColorsReallyCoolType]
+class AdvancedColors2Types(str, Enum):
+    STRING = "string"
+    NUMBER = "number"
+    NUMBER_ARRAY = "number-array"
+    REALLY_COOL_TYPE = "really-cool-type"
+
 class AdvancedColors2String(BaseModel):
     """
     This is a case comment
     """
-    AdvancedColors2Types: Literal["string"] = "string"
+    type: Literal[AdvancedColors2Types.STRING] = AdvancedColors2Types.STRING
     content: str
 
 class AdvancedColors2Number(BaseModel):
-    AdvancedColors2Types: Literal["number"] = "number"
+    type: Literal[AdvancedColors2Types.NUMBER] = AdvancedColors2Types.NUMBER
     content: int
 
 class AdvancedColors2NumberArray(BaseModel):
-    AdvancedColors2Types: Literal["number-array"] = "number-array"
+    type: Literal[AdvancedColors2Types.NUMBER_ARRAY] = AdvancedColors2Types.NUMBER_ARRAY
     content: List[int]
 
 class AdvancedColors2ReallyCoolType(BaseModel):
     """
     Comment on the last element
     """
-    AdvancedColors2Types: Literal["really-cool-type"] = "really-cool-type"
+    type: Literal[AdvancedColors2Types.REALLY_COOL_TYPE] = AdvancedColors2Types.REALLY_COOL_TYPE
     content: ItemDetailsFieldValue
 
 AdvancedColors2 = Union[AdvancedColors2String, AdvancedColors2Number, AdvancedColors2NumberArray, AdvancedColors2ReallyCoolType]

core/data/tests/can_generate_algebraic_enum_with_skipped_variants/output.py

@@ -3,15 +3,20 @@
 """
 from __future__ import annotations
 
+from enum import Enum
 from pydantic import BaseModel
 from typing import Literal, Union
 
 
+class SomeEnumTypes(str, Enum):
+    A = "A"
+    C = "C"
+
 class SomeEnumA(BaseModel):
-    SomeEnumTypes: Literal["A"] = "A"
+    type: Literal[SomeEnumTypes.A] = SomeEnumTypes.A
 
 class SomeEnumC(BaseModel):
-    SomeEnumTypes: Literal["C"] = "C"
+    type: Literal[SomeEnumTypes.C] = SomeEnumTypes.C
     content: int
 
 SomeEnum = Union[SomeEnumA, SomeEnumC]

core/data/tests/can_generate_anonymous_struct_with_skipped_fields/output.py

@@ -3,6 +3,7 @@
 """
 from __future__ import annotations
 
+from enum import Enum
 from pydantic import BaseModel
 from typing import Literal, Union
 
@@ -25,18 +26,22 @@ class AutofilledBySomethingElseInner(BaseModel):
     The UUID for the fill
     """
 
+class AutofilledByTypes(str, Enum):
+    US = "Us"
+    SOMETHING_ELSE = "SomethingElse"
+
 class AutofilledByUs(BaseModel):
     """
     This field was autofilled by us
     """
-    AutofilledByTypes: Literal["Us"] = "Us"
+    type: Literal[AutofilledByTypes.US] = AutofilledByTypes.US
     content: AutofilledByUsInner
 
 class AutofilledBySomethingElse(BaseModel):
     """
     Something else autofilled this field
     """
-    AutofilledByTypes: Literal["SomethingElse"] = "SomethingElse"
+    type: Literal[AutofilledByTypes.SOMETHING_ELSE] = AutofilledByTypes.SOMETHING_ELSE
     content: AutofilledBySomethingElseInner
 
 # Enum keeping track of who autofilled a field

core/data/tests/can_generate_empty_algebraic_enum/output.py

@@ -3,17 +3,22 @@
 """
 from __future__ import annotations
 
+from enum import Enum
 from pydantic import BaseModel
 from typing import Literal, Union
 
 
 class AddressDetails(BaseModel):
     pass
+class AddressTypes(str, Enum):
+    FIXED_ADDRESS = "FixedAddress"
+    NO_FIXED_ADDRESS = "NoFixedAddress"
+
 class AddressFixedAddress(BaseModel):
-    AddressTypes: Literal["FixedAddress"] = "FixedAddress"
+    type: Literal[AddressTypes.FIXED_ADDRESS] = AddressTypes.FIXED_ADDRESS
     content: AddressDetails
 
 class AddressNoFixedAddress(BaseModel):
-    AddressTypes: Literal["NoFixedAddress"] = "NoFixedAddress"
+    type: Literal[AddressTypes.NO_FIXED_ADDRESS] = AddressTypes.NO_FIXED_ADDRESS
 
 Address = Union[AddressFixedAddress, AddressNoFixedAddress]

core/data/tests/can_handle_anonymous_struct/output.py

@@ -3,6 +3,7 @@
 """
 from __future__ import annotations
 
+from enum import Enum
 from pydantic import BaseModel
 from typing import Literal, Union
 
@@ -29,18 +30,22 @@ class AutofilledBySomethingElseInner(BaseModel):
     Some other thing
     """
 
+class AutofilledByTypes(str, Enum):
+    US = "Us"
+    SOMETHING_ELSE = "SomethingElse"
+
 class AutofilledByUs(BaseModel):
     """
     This field was autofilled by us
     """
-    AutofilledByTypes: Literal["Us"] = "Us"
+    type: Literal[AutofilledByTypes.US] = AutofilledByTypes.US
     content: AutofilledByUsInner
 
 class AutofilledBySomethingElse(BaseModel):
     """
     Something else autofilled this field
     """
-    AutofilledByTypes: Literal["SomethingElse"] = "SomethingElse"
+    type: Literal[AutofilledByTypes.SOMETHING_ELSE] = AutofilledByTypes.SOMETHING_ELSE
     content: AutofilledBySomethingElseInner
 
 # Enum keeping track of who autofilled a field
@@ -58,26 +63,34 @@ class EnumWithManyVariantsAnotherAnonVariantInner(BaseModel):
     uuid: str
     thing: int
 
+class EnumWithManyVariantsTypes(str, Enum):
+    UNIT_VARIANT = "UnitVariant"
+    TUPLE_VARIANT_STRING = "TupleVariantString"
+    ANON_VARIANT = "AnonVariant"
+    TUPLE_VARIANT_INT = "TupleVariantInt"
+    ANOTHER_UNIT_VARIANT = "AnotherUnitVariant"
+    ANOTHER_ANON_VARIANT = "AnotherAnonVariant"
+
 class EnumWithManyVariantsUnitVariant(BaseModel):
-    EnumWithManyVariantsTypes: Literal["UnitVariant"] = "UnitVariant"
+    type: Literal[EnumWithManyVariantsTypes.UNIT_VARIANT] = EnumWithManyVariantsTypes.UNIT_VARIANT
 
 class EnumWithManyVariantsTupleVariantString(BaseModel):
-    EnumWithManyVariantsTypes: Literal["TupleVariantString"] = "TupleVariantString"
+    type: Literal[EnumWithManyVariantsTypes.TUPLE_VARIANT_STRING] = EnumWithManyVariantsTypes.TUPLE_VARIANT_STRING
     content: str
 
 class EnumWithManyVariantsAnonVariant(BaseModel):
-    EnumWithManyVariantsTypes: Literal["AnonVariant"] = "AnonVariant"
+    type: Literal[EnumWithManyVariantsTypes.ANON_VARIANT] = EnumWithManyVariantsTypes.ANON_VARIANT
     content: EnumWithManyVariantsAnonVariantInner
 
 class EnumWithManyVariantsTupleVariantInt(BaseModel):
-    EnumWithManyVariantsTypes: Literal["TupleVariantInt"] = "TupleVariantInt"
+    type: Literal[EnumWithManyVariantsTypes.TUPLE_VARIANT_INT] = EnumWithManyVariantsTypes.TUPLE_VARIANT_INT
     content: int
 
 class EnumWithManyVariantsAnotherUnitVariant(BaseModel):
-    EnumWithManyVariantsTypes: Literal["AnotherUnitVariant"] = "AnotherUnitVariant"
+    type: Literal[EnumWithManyVariantsTypes.ANOTHER_UNIT_VARIANT] = EnumWithManyVariantsTypes.ANOTHER_UNIT_VARIANT
 
 class EnumWithManyVariantsAnotherAnonVariant(BaseModel):
-    EnumWithManyVariantsTypes: Literal["AnotherAnonVariant"] = "AnotherAnonVariant"
+    type: Literal[EnumWithManyVariantsTypes.ANOTHER_ANON_VARIANT] = EnumWithManyVariantsTypes.ANOTHER_ANON_VARIANT
     content: EnumWithManyVariantsAnotherAnonVariantInner
 
 # This is a comment (yareek sameek wuz here)

core/data/tests/can_handle_unit_type/output.py

@@ -3,6 +3,7 @@
 """
 from __future__ import annotations
 
+from enum import Enum
 from pydantic import BaseModel, ConfigDict, Field
 from typing import Literal
 
@@ -15,8 +16,11 @@ class StructHasVoidType(BaseModel):
 
     this_is_a_unit: None = Field(alias="thisIsAUnit")
 
+class EnumHasVoidTypeTypes(str, Enum):
+    HAS_A_UNIT = "hasAUnit"
+
 class EnumHasVoidTypeHasAUnit(BaseModel):
-    EnumHasVoidTypeTypes: Literal["hasAUnit"] = "hasAUnit"
+    type: Literal[EnumHasVoidTypeTypes.HAS_A_UNIT] = EnumHasVoidTypeTypes.HAS_A_UNIT
     content: None
 
 # This enum has a variant associated with unit data

core/data/tests/excluded_by_target_os/output.py

@@ -41,21 +41,28 @@ class TestEnumVariant9Inner(BaseModel):
 
     field_2: str = Field(alias="field2")
 
+class TestEnumTypes(str, Enum):
+    VARIANT_1 = "Variant1"
+    VARIANT_5 = "Variant5"
+    VARIANT_7 = "Variant7"
+    VARIANT_8 = "Variant8"
+    VARIANT_9 = "Variant9"
+
 class TestEnumVariant1(BaseModel):
-    TestEnumTypes: Literal["Variant1"] = "Variant1"
+    type: Literal[TestEnumTypes.VARIANT_1] = TestEnumTypes.VARIANT_1
 
 class TestEnumVariant5(BaseModel):
-    TestEnumTypes: Literal["Variant5"] = "Variant5"
+    type: Literal[TestEnumTypes.VARIANT_5] = TestEnumTypes.VARIANT_5
 
 class TestEnumVariant7(BaseModel):
-    TestEnumTypes: Literal["Variant7"] = "Variant7"
+    type: Literal[TestEnumTypes.VARIANT_7] = TestEnumTypes.VARIANT_7
     content: TestEnumVariant7Inner
 
 class TestEnumVariant8(BaseModel):
-    TestEnumTypes: Literal["Variant8"] = "Variant8"
+    type: Literal[TestEnumTypes.VARIANT_8] = TestEnumTypes.VARIANT_8
 
 class TestEnumVariant9(BaseModel):
-    TestEnumTypes: Literal["Variant9"] = "Variant9"
+    type: Literal[TestEnumTypes.VARIANT_9] = TestEnumTypes.VARIANT_9
     content: TestEnumVariant9Inner
 
 TestEnum = Union[TestEnumVariant1, TestEnumVariant5, TestEnumVariant7, TestEnumVariant8, TestEnumVariant9]