pydantic_version.py

from typing import Union, Dict
from enum import Enum
from pydantic.dataclasses import dataclass
from pydantic import ValidationError, validator
from pydantic import confloat, conint

class ProcedureType(Enum):
    cancer = "cancer"
    flu = "flu"
    
class HealthOutcome(Enum):
    healthy = "healthy"
    at_risk = "at risk"
    unknown = "unknown"
    
@dataclass
class Response:
    outcome: HealthOutcome
    probability: confloat(ge=0, le=1)
        
@dataclass
class UserAssessment:
    procedure: ProcedureType
    age: conint(ge=1, le=100)
    occupation: str
    doctor_approved: bool
        
    @validator("doctor_approved")
    def procedure_type(cls, v: bool, values: Dict) -> Union[bool, dict]:
        if v == False and values["procedure"] == ProcedureType.cancer:
            raise ValueError(
                "Automatic model based diagnosis is not available without doctor approval"
            )
        return v
    
response = Union[ValueError, Response]

def apply_treatment(user_data: UserAssessment) -> response:
    return Response(HealthOutcome.healthy, 0.05)