Skip to content

Commit

Permalink
Update unit enum to use enum values
Browse files Browse the repository at this point in the history
  • Loading branch information
MOmarMiraj committed Nov 7, 2024
1 parent 03d600b commit 61b6afd
Show file tree
Hide file tree
Showing 7 changed files with 17 additions and 16 deletions.
3 changes: 2 additions & 1 deletion core/data/tests/can_generate_bare_string_enum/output.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,10 +4,11 @@
from __future__ import annotations

from enum import Enum
from pydantic import ConfigDict


class Colors(Enum):
model_config = ConfigDict(use_enum_values=True)
RED = "Red"
BLUE = "Blue"
GREEN = "Green"

3 changes: 2 additions & 1 deletion core/data/tests/can_generate_simple_enum/output.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,10 +4,11 @@
from __future__ import annotations

from enum import Enum
from pydantic import ConfigDict


class Colors(Enum):
model_config = ConfigDict(use_enum_values=True)
RED = "Red"
BLUE = "Blue"
GREEN = "Green"

3 changes: 2 additions & 1 deletion core/data/tests/can_handle_quote_in_serde_rename/output.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,8 +4,9 @@
from __future__ import annotations

from enum import Enum
from pydantic import ConfigDict


class Colors(Enum):
model_config = ConfigDict(use_enum_values=True)
GREEN = "Green\""

Original file line number Diff line number Diff line change
Expand Up @@ -4,10 +4,11 @@
from __future__ import annotations

from enum import Enum
from pydantic import ConfigDict


class Colors(Enum):
model_config = ConfigDict(use_enum_values=True)
RED = "red"
BLUE = "blue"
GREEN = "green-like"

1 change: 0 additions & 1 deletion core/data/tests/excluded_by_target_os/output.py
Original file line number Diff line number Diff line change
Expand Up @@ -31,7 +31,6 @@ class OtherExcluded(BaseModel):

class SomeEnum(Enum):
pass

class TestEnumVariant7(BaseModel):
"""
Generated type representing the anonymous struct variant `Variant7` of the `TestEnum` Rust enum
Expand Down
3 changes: 2 additions & 1 deletion core/data/tests/test_simple_enum_case_name_support/output.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,10 +4,11 @@
from __future__ import annotations

from enum import Enum
from pydantic import ConfigDict


class Colors(Enum):
model_config = ConfigDict(use_enum_values=True)
RED = "red"
BLUE = "blue-ish"
GREEN = "Green"

17 changes: 7 additions & 10 deletions core/src/language/python.rs
Original file line number Diff line number Diff line change
Expand Up @@ -270,20 +270,21 @@ impl Language for Python {
// Generate named types for any anonymous struct variants of this enum
self.write_types_for_anonymous_structs(w, e, &make_anonymous_struct_name)?;
self.add_import("enum".to_string(), "Enum".to_string());
self.add_import("pydantic".to_string(), "ConfigDict".to_string());
match e {
// 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)?;

let fields = if shared.variants.is_empty() {
" pass".to_string()
if shared.variants.is_empty(){
writeln!(w, " pass")?;
} else {
writeln!(w, " model_config = ConfigDict(use_enum_values=True)")?;
shared
.variants
.iter()
.map(|v| {
format!(
.try_for_each(|v| {
writeln!(w,
" {} = \"{}\"",
v.shared().id.original.to_uppercase(),
match v {
Expand All @@ -293,11 +294,8 @@ impl Language for Python {
_ => panic!(),
}
)
})
.collect::<Vec<String>>()
.join("\n")
})?
};
writeln!(w, "{fields}\n")?;
}
// Write all the algebraic variants out (all three variant types are possible
// here)
Expand Down Expand Up @@ -590,7 +588,6 @@ impl Python {
.cloned()
.for_each(|v| self.add_type_var(v));
let mut variants: Vec<(String, Vec<String>)> = Vec::new();
self.add_import("pydantic".to_string(), "ConfigDict".to_string());
self.add_import("pydantic".to_string(), "BaseModel".to_string());
// write "types" class: a union of all the enum variants
writeln!(w, "class {}Types(str, Enum):", shared.id.renamed)?;
Expand Down

0 comments on commit 61b6afd

Please sign in to comment.