44

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"
Gabriel Petersson
  • 8,434
  • 4
  • 32
  • 41

3 Answers3

98

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:

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

see here

SColvin
  • 11,584
  • 6
  • 57
  • 71
  • 7
    Thanks this worked. I also want to add that in python 3.8+ its `from typing import Literal`, while <3.8 you need to use `from typing_extensions import Literal` – Gabriel Petersson Apr 18 '20 at 14:37
  • 19
    Or for simplicity you can use `from pydantic.typing import Literal` which will work on both. – SColvin Apr 18 '20 at 15:33
3

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$")
Steven Staley
  • 182
  • 1
  • 1
  • 10
0

Easiest: Use typing.Literal

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

class Input(BaseModel):
    option: Litera["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
Yaakov Bressler
  • 9,056
  • 2
  • 45
  • 69