Skip to content

Commit

Permalink
proper enum schema for pydantic
Browse files Browse the repository at this point in the history
  • Loading branch information
jbkoh committed Dec 9, 2024
1 parent ed7eefa commit 6b75cae
Showing 1 changed file with 10 additions and 1 deletion.
11 changes: 10 additions & 1 deletion src/betterproto/templates/template.py.j2
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,16 @@ class {{ enum.py_name }}(betterproto.Enum):
def __get_pydantic_core_schema__(cls, _source_type, _handler):
from pydantic_core import core_schema

return core_schema.int_schema(ge=0)
def validate(value: int) -> "{{ enum.py_name }}":
return cls(value)

# Return the schema for validation and serialization
return core_schema.chain_schema(
[
core_schema.int_schema(), # Validate as a string first
core_schema.no_info_plain_validator_function(validate), # Custom validation
]
)
{% endif %}

{% endfor %}
Expand Down

0 comments on commit 6b75cae

Please sign in to comment.