Javier SG
03/05/2025, 11:51 AMMarvin
03/05/2025, 11:51 AMpython
from prefect import flow
from pydantic import BaseModel, Field
from typing import Union, Literal
# Define your subclasses
class Cat(BaseModel):
type: Literal['cat']
meow_volume: int
class Dog(BaseModel):
type: Literal['dog']
bark_volume: int
# Define the main model using a discriminator
class Pet(BaseModel):
pet: Union[Cat, Dog] = Field(..., discriminator='type')
# Define your Prefect flow
@flow
def my_flow(pet: Pet):
if pet.pet.type == 'cat':
print(f"A cat with meow volume {pet.pet.meow_volume}")
elif pet.pet.type == 'dog':
print(f"A dog with bark volume {pet.pet.bark_volume}")
# Example usage
my_flow(pet=Pet(pet={"type": "cat", "meow_volume": 5}))
### Key Points:
- Pydantic Models: You can use Pydantic models to enforce and validate flow parameters.
- Union with Discriminator: The discriminator
field allows Pydantic to determine which subclass to use based on a specified field (type
in this example).
For more information on using Pydantic with Prefect, you can refer to the Prefect documentation.