CODEWITHSUNDEEP
pythonpydantic
Gabriel Petersson
Gabriel Petersson

Reputation: 10412

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: 61

Views: 62314

Answers (4)

Samuel Colvin
Samuel Colvin

Reputation: 13269

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: 127

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

Yaakov Bressler
Yaakov Bressler

Reputation: 12008

Easiest: Use typing.Literal

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

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

It will fail if your list of strings are dynamic:

allowed_values = ["foo", "bar"]

class Input(BaseModel):
    option: Literal[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 with Annotated Validator

Let's say this field (and validator) are going to be reused in your codebase. A better approach would be to create a "custom field type" with an annotated validator, as such:

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: 9

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