Gabriel Petersson
Gabriel Petersson

Reputation: 10492

How to require predefined string values in python pydantic basemodels?

Is there any in-built way in pydantic to specify options? For example, let's say I want a string value that must either have the value "foo" or "bar".

I know I can use regex validation to do this, but since I use pydantic with FastAPI, the users will only see the required input as a string, but when they enter something, it will give a validation error. All in-built validations of pydantic are displayed in the api interface, so would be great if there was something like

class Input(BaseModel):
     option: "foo" || "bar"

Upvotes: 64

Views: 65028

Answers (4)

Yaakov Bressler
Yaakov Bressler

Reputation: 12168

Easiest: Use typing.Literal

Literal is a good option when you can hardcode your values:

class Input(BaseModel):
    option: Literal["foo", "bar"]

To use a dynamic list, wrap your list with a tuple:

allowed_values = ["foo", "bar"]

class Input(BaseModel):
    option: Literal[tuple(allowed_values)]

Alt: Use Validator

Assuming it is not possible to transcode into regex (say you have objects, not only strings), you would then want to use a field validator:

allowed_values = ["foo", "bar"]

class Input(BaseModel):
    option: str

    @field_validator("option")
    def validate_option(cls, v):
        assert v in allowed_values
        return v

Best: Reusable Field

Let's say this field is going to be reused in your codebase. A better approach would be to create a "custom field type" as such:

from typing import Annotated, Literal


allowed_values = ["foo", "bar"]
custom_option = Annotated[Literal[tuple(allowed_values)]]

class Input(BaseModel):
    option: custom_option

Reusable field with Annotated Validator

Previously, I had this as the answer. Will keep for posterity. You can add validators1 to reusable fields, which are neat:

from typing import Annotated
from pydantic import BaseModel, AfterValidator


allowed_values = ["foo", "bar"]

def option_in_allowed_values(v):
    """Ensures option is allowed"""
    assert v in allowed_values
    return v

custom_option = Annotated[str, AfterValidator(option_in_allowed_values)]


class Input(BaseModel):
    option: custom_option

Upvotes: 11

Samuel Colvin
Samuel Colvin

Reputation: 13349

Yes, you can either use an enum:

class Choices(Enum):
    foo = 'foo'
    bar = 'bar'

class Input(BaseModel):
     option: Choices

see here

Or you can use Literal:

from typing import Literal

class Input(BaseModel):
     option: Literal['foo', 'bar']

see here

Upvotes: 130

Shaun
Shaun

Reputation: 3895

I'd like to expand on Yaakov's answer. You can use functools.partial to bake in your values list

from functools import partial
from typing import Annotated
from pydantic import BaseModel, AfterValidator


def allowed_values(v, values):
    assert v in values
    return v

class Input(BaseModel):
    option: Annotated[str, AfterValidator(partial(allowed_values, values=["a", "b"]))]

Upvotes: 3

Steven Staley
Steven Staley

Reputation: 337

Wanted to add another option here. You could also use a Regex. Worked better for me since Literal isn't available until python 3.8 (which is unfortunately not an easy upgrade for me) and since I'm only expecting a single string for each, enum didn't really fit.

class YourClass(pydantic.BaseModel):
  your_attribute: pydantic.constr(regex="^yourvalwith\.escapes\/abcd$")

Upvotes: 5

Related Questions