Remove config dict and allow unit enums to be serializble

1Password · Nov 7, 2024 · 92f6922 · 92f6922
1 parent 9810e7b
commit 92f6922
Show file tree

Hide file tree

Showing 7 changed files with 10 additions and 16 deletions.
diff --git a/core/data/tests/can_generate_bare_string_enum/output.py b/core/data/tests/can_generate_bare_string_enum/output.py
@@ -7,8 +7,7 @@
 from pydantic import ConfigDict
 
 
-class Colors(Enum):
-    model_config = ConfigDict(use_enum_values=True)
+class Colors(str, Enum):
     RED = "Red"
     BLUE = "Blue"
     GREEN = "Green"
diff --git a/core/data/tests/can_generate_simple_enum/output.py b/core/data/tests/can_generate_simple_enum/output.py
@@ -7,8 +7,7 @@
 from pydantic import ConfigDict
 
 
-class Colors(Enum):
-    model_config = ConfigDict(use_enum_values=True)
+class Colors(str, Enum):
     RED = "Red"
     BLUE = "Blue"
     GREEN = "Green"
diff --git a/core/data/tests/can_handle_quote_in_serde_rename/output.py b/core/data/tests/can_handle_quote_in_serde_rename/output.py
@@ -7,6 +7,5 @@
 from pydantic import ConfigDict
 
 
-class Colors(Enum):
-    model_config = ConfigDict(use_enum_values=True)
+class Colors(str, Enum):
     GREEN = "Green\""
diff --git a/core/data/tests/enum_is_properly_named_with_serde_overrides/output.py b/core/data/tests/enum_is_properly_named_with_serde_overrides/output.py
@@ -7,8 +7,7 @@
 from pydantic import ConfigDict
 
 
-class Colors(Enum):
-    model_config = ConfigDict(use_enum_values=True)
+class Colors(str, Enum):
     RED = "red"
     BLUE = "blue"
     GREEN = "green-like"
diff --git a/core/data/tests/excluded_by_target_os/output.py b/core/data/tests/excluded_by_target_os/output.py
@@ -29,7 +29,7 @@ class NestedNotTarget1(BaseModel):
 class OtherExcluded(BaseModel):
     pass
 
-class SomeEnum(Enum):
+class SomeEnum(str, Enum):
     pass
 class TestEnumVariant7(BaseModel):
     """

diff --git a/core/data/tests/test_simple_enum_case_name_support/output.py b/core/data/tests/test_simple_enum_case_name_support/output.py
@@ -7,8 +7,7 @@
 from pydantic import ConfigDict
 
 
-class Colors(Enum):
-    model_config = ConfigDict(use_enum_values=True)
+class Colors(str, Enum):
     RED = "red"
     BLUE = "blue-ish"
     GREEN = "Green"
diff --git a/core/src/language/python.rs b/core/src/language/python.rs
@@ -248,7 +248,7 @@ impl Language for Python {
 
         self.write_comments(w, true, &rs.comments, 1)?;
 
-        handle_model_config(w, self, rs);
+        handle_model_config(w, self, &rs.fields);
 
         rs.fields
             .iter()
@@ -275,11 +275,10 @@ impl Language for Python {
             // Write all the unit variants out (there can only be unit variants in
             // this case)
             RustEnum::Unit(shared) => {
-                writeln!(w, "class {}(Enum):", shared.id.renamed)?;
+                writeln!(w, "class {}(str, Enum):", shared.id.renamed)?;
                 if shared.variants.is_empty() {
                     writeln!(w, "    pass")?;
                 } else {
-                    writeln!(w, "    model_config = ConfigDict(use_enum_values=True)")?;
                     shared.variants.iter().try_for_each(|v| {
                         writeln!(
                             w,
@@ -765,8 +764,8 @@ fn python_property_aware_rename(name: &str) -> String {
 }
 
 // If at least one field from within a class is changed when the serde rename is used (a.k.a the field has 2 words) then we must use aliasing and we must also use a config dict at the top level of the class.
-fn handle_model_config(w: &mut dyn Write, python_module: &mut Python, rs: &RustStruct) {
-    let visibly_renamed_field = rs.fields.iter().find(|f| {
+fn handle_model_config(w: &mut dyn Write, python_module: &mut Python, fields: &[RustField]) {
+    let visibly_renamed_field = fields.iter().find(|f| {
         let python_field_name = python_property_aware_rename(&f.id.original);
         python_field_name != f.id.renamed
     });