class Report(Document): best_number_cluster: IntField() cluster_name: EnumField(ClusterModel) classification_model_name: EnumField(ClassificationModel) pca: FloatField() start_time = DateTimeField(default=datetime.utcnow) end_time = DateTimeField(default=datetime.utcnow)
def test_wrong_choices(self): with pytest.raises(ValueError, match="Invalid choices"): EnumField(Status, choices=["my", "custom", "options"]) with pytest.raises(ValueError, match="Invalid choices"): EnumField(Status, choices=[Color.RED]) with pytest.raises(ValueError, match="Invalid choices"): EnumField(Status, choices=[Status.DONE, Color.RED])
class TextCovidSymptomAnnotation(EmbeddedDocument): start = IntField(required=True) length = IntField(required=True) text = StringField(required=True) confidence = FloatField(required=True) condition = EnumField(Condition, required=True) certainty = EnumField(Certainty, required=True) def to_dict(self): doc = self.to_mongo().to_dict() return doc
class Identity(Document): process = LazyReferenceField("Process") name = StringField(required=True) card_number = StringField() type = EnumField(IdentityType, default=IdentityType.KNOWN, required=True) status = EnumField(IdentityStatus, default=IdentityStatus.TRACKING, required=True) matching_face_ids = ListField(EmbeddedDocumentField(MatchingFace)) clustering_face_ids = ListField( ListField(EmbeddedDocumentField(ClusteringFace))) created_at = DateTimeField(default=datetime.datetime.utcnow, required=True) meta = {"collection": "identities"}
class MetricsReport(Document): name: EnumField(MetricClustering) best_value_cluster: IntField() best_value_score: FloatField() value: ListField(field=FloatField()) rank: ListField(field=IntField()) created_at = DateTimeField(default=datetime.utcnow)
class ClusterElement(Document): cluster = LazyReferenceField("Cluster", required=True) object = LazyReferenceField("Object", required=True) ref_object = LazyReferenceField("Object") type = EnumField(ClusterType, default=ClusterType.FACE, required=True) created_at = DateTimeField(default=datetime.utcnow, required=True) meta = {"collection": "cluster_elements"}
class TextLocationAnnotation(EmbeddedDocument): start = IntField(required=True) length = IntField(required=True) text = StringField(required=True) confidence = FloatField(required=True) locationType = EnumField(LocationType, required=True) def to_dict(self): doc = self.to_mongo().to_dict() return doc
def test_partial_choices(self): partial = [Status.DONE] enum_field = EnumField(Status, choices=partial) assert enum_field.choices == partial class FancyDoc(Document): z = enum_field FancyDoc(z=Status.DONE).validate() with pytest.raises(ValidationError, match=r"Value must be one of .*Status.DONE"): FancyDoc(z=Status.NEW).validate()
class User(BaseDocument): name = StringField() last_name = StringField() email = StringField(unique=True) phone = IntField() phone_prefix = StringField() institution = StringField() institution_role = StringField() profession = StringField() gender = StringField() birthday = DateField() role = EnumField(UserRoles) is_deleted = BooleanField(default=False) is_valid = BooleanField(default=False) is_enabled = BooleanField(default=True) notify_removal = BooleanField(default=True) notify_simulation_done = BooleanField(default=True)
class ModelWithColor(Document): color = EnumField(Color, default=Color.RED)
def test_user_is_informed_when_tries_to_set_choices(self): with pytest.raises(ValueError, match="'choices' can't be set on EnumField"): EnumField(Status, choices=["my", "custom", "options"])
class ModelWithDefault(Document): status = EnumField(Status, default=Status.DONE)
class ModelWithEnum(Document): status = EnumField(Status)
class ModelWithFunkyColor(Document): color = EnumField(FunkyColor)
def test_partial_choices(self): partial = [Status.DONE] assert EnumField(Status, choices=partial).choices == partial