mypy compliant

This commit is contained in:
Brian Bjarke Jensen
2024-02-25 00:04:01 +01:00
parent 4719117d71
commit 8e107b785e
3 changed files with 8 additions and 9 deletions
+2 -2
View File
@@ -52,7 +52,7 @@ class ModelOutputs(BaseModel):
def from_random(cls) -> ModelOutputs: def from_random(cls) -> ModelOutputs:
"""Instantiate with random numbers.""" """Instantiate with random numbers."""
kwargs = { kwargs = {
field: field_info.annotation.from_random() field: field_info.annotation.from_random() # type: ignore
for field, field_info for field, field_info
in cls.model_fields.items() in cls.model_fields.items()
} }
@@ -124,7 +124,7 @@ class VisualCommunication(BaseModel):
return VisualCommunication(name=name, image=image) return VisualCommunication(name=name, image=image)
@field_serializer("image") @field_serializer("image")
def serialize_image(image: Image.Image) -> bytes: def serialize_image(image: Image.Image) -> bytes: # type: ignore
buffer = BytesIO() buffer = BytesIO()
image.save(buffer, format="JPEG") image.save(buffer, format="JPEG")
return buffer.getvalue() return buffer.getvalue()
+4 -5
View File
@@ -36,7 +36,7 @@ def get_visual_communication(
if with_annotation: if with_annotation:
query["annotation"] = {"$ne": None} query["annotation"] = {"$ne": None}
else: else:
query["annotation"] = None query["annotation"] = {"$eq": None}
data = collection.aggregate([ data = collection.aggregate([
{ {
"$match": query # find using filters "$match": query # find using filters
@@ -47,13 +47,12 @@ def get_visual_communication(
} }
} }
]) ])
data = list(data) # read data from cursor object data_list = list(data) # read data from cursor object
if len(data) == 0: if len(data_list) == 0:
logging.error("failed getting visual communication") logging.error("failed getting visual communication")
raise NoDocumentFoundException() raise NoDocumentFoundException()
data = data[0]
logging.info("finished") logging.info("finished")
return VisualCommunication.model_validate(data) return VisualCommunication.model_validate(data_list[0])
def upsert_predictions( def upsert_predictions(
+2 -2
View File
@@ -103,13 +103,13 @@ def cycle_visual_communication_data(
for elem for elem
in annotation_values in annotation_values
] ]
annotations = { annotation_map = {
key: value key: value
for key, value for key, value
in zip(annotation_keys, annotation_values) in zip(annotation_keys, annotation_values)
} }
# instantiate ModelOutputs object # instantiate ModelOutputs object
annotations = ModelOutputs.from_annotations(**annotations) annotations = ModelOutputs.from_annotations(**annotation_map)
# save data to database # save data to database
upsert_annotations( upsert_annotations(
collection=collection, collection=collection,