dataclass.py

from typing import Optional, Union
from enum import Enum
from dataclasses import dataclass
from pydantic import ValidationError


class ProcedureType(Enum):
    cancer = "cancer"
    flu = "flu"


class HealthOutcome(Enum):
    healthy = "healthy"
    at_risk = "at risk"
    unknown = "unknown"


@dataclass
class Response:
    outcome: HealthOutcome
    probability: float


@dataclass
class UserAssessment:
    procedure: ProcedureType
    age: int
    occupation: str
    doctor_approved: Optional[bool] = None


response = Union[ValidationError, Response]


def apply_treatment(user_data: UserAssessment) -> response:
    if user_data.procedure == "cancer" and user_data.doctor_approved == False:
        return ValidationError(
            "Automatic model based diagnosis is not available without doctor approval"
        )
    return Response(HealthOutcome.healthy, 0.05)