added from_choice classmethod

This commit is contained in:
Brian Bjarke Jensen
2024-02-24 22:42:11 +01:00
parent 59ad0104ad
commit 5a147e9492
4 changed files with 45 additions and 6 deletions
+13 -1
View File
@@ -1,4 +1,4 @@
from pydantic import BaseModel
from pydantic import BaseModel, ValidationError
from typing import List
import random
@@ -20,6 +20,18 @@ class ModelOutput(BaseModel):
"""Instantiate with random numbers."""
kwargs = {field: random.random() for field in cls.list_fields()}
return cls(**kwargs)
@classmethod
def from_choice(cls, option: str):
"""Instantiate from choice."""
if option is None:
raise ValidationError()
assert isinstance(option, str)
allowed_options_list = cls.list_fields()
assert option in allowed_options_list, f"{option} is not among allowed fields {allowed_options_list}"
kwargs = {field: 0 for field in cls.list_fields()}
kwargs[option] = 1
return cls(**kwargs)
def __repr__(self) -> str:
model_dict = self.model_dump()