added from_choice classmethod
This commit is contained in:
@@ -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()
|
||||
|
||||
Reference in New Issue
Block a user