Move model_cast into helpers

This commit is contained in:
James Ketr 2025-06-18 12:31:17 -07:00
parent e5ac267935
commit aefc14c610

View File

@ -0,0 +1,53 @@
from typing import Type, TypeVar
from pydantic import BaseModel
import copy
from models import Candidate, CandidateAI, Employer, Guest, BaseUserWithType
# Ensure all user models inherit from BaseUserWithType
assert issubclass(Candidate, BaseUserWithType), "Candidate must inherit from BaseUserWithType"
assert issubclass(CandidateAI, BaseUserWithType), "CandidateAI must inherit from BaseUserWithType"
assert issubclass(Employer, BaseUserWithType), "Employer must inherit from BaseUserWithType"
assert issubclass(Guest, BaseUserWithType), "Guest must inherit from BaseUserWithType"
T = TypeVar('T', bound=BaseModel)
def cast_to_model(model_cls: Type[T], source: BaseModel) -> T:
data = {field: getattr(source, field) for field in model_cls.__fields__}
return model_cls(**data)
def cast_to_model_safe(model_cls: Type[T], source: BaseModel) -> T:
data = {field: copy.deepcopy(getattr(source, field)) for field in model_cls.__fields__}
return model_cls(**data)
def cast_to_base_user_with_type(user) -> BaseUserWithType:
"""
Casts a Candidate, CandidateAI, Employer, or Guest to BaseUserWithType.
This is useful for FastAPI dependencies that expect a common user type.
"""
if isinstance(user, BaseUserWithType):
return user
# If it's a dict, try to detect type
if isinstance(user, dict):
user_type = user.get("user_type") or user.get("type")
if user_type == "candidate":
if user.get("is_AI"):
return CandidateAI.model_validate(user)
return Candidate.model_validate(user)
elif user_type == "employer":
return Employer.model_validate(user)
elif user_type == "guest":
return Guest.model_validate(user)
else:
raise ValueError(f"Unknown user_type: {user_type}")
# If it's a model, check its type
if hasattr(user, "user_type"):
if getattr(user, "user_type", None) == "candidate":
if getattr(user, "is_AI", False):
return CandidateAI.model_validate(user.model_dump())
return Candidate.model_validate(user.model_dump())
elif getattr(user, "user_type", None) == "employer":
return Employer.model_validate(user.model_dump())
elif getattr(user, "user_type", None) == "guest":
return Guest.model_validate(user.model_dump())
raise TypeError(f"Cannot cast object of type {type(user)} to BaseUserWithType")