Refactor schema.py and add VertexBuildResponseModel

This commit is contained in:
Gabriel Luiz Freitas Almeida 2024-02-16 16:44:08 -03:00
commit d6ebc60462

View file

@ -7,7 +7,9 @@ from pydantic import BaseModel, Field, field_serializer, validator
class TransactionModel(BaseModel): class TransactionModel(BaseModel):
id: Optional[int] = Field(default=None, alias="id") id: Optional[int] = Field(default=None, alias="id")
timestamp: Optional[datetime] = Field(default_factory=datetime.now, alias="timestamp") timestamp: Optional[datetime] = Field(
default_factory=datetime.now, alias="timestamp"
)
source: str source: str
target: str target: str
target_args: dict target_args: dict
@ -89,18 +91,13 @@ class VertexBuildModel(BaseModel):
return v return v
# create a function that turns dicts into a class VertexBuildResponseModel(VertexBuildModel):
# dict like this:
# my_map_dict = { @field_serializer("data", "artifacts")
# "key": [ def serialize_dict(v):
# 1, 2, 3 return v
# ],
# "value": [
# "one", "two", "three"
# ]
# }
# so map has a "key" and a "value" list
# containing the keys and values of the dict
def to_map(value: dict): def to_map(value: dict):
keys = list(value.keys()) keys = list(value.keys())
values = list(value.values()) values = list(value.values())
@ -108,13 +105,13 @@ def to_map(value: dict):
class VertexBuildMapModel(BaseModel): class VertexBuildMapModel(BaseModel):
vertex_builds: dict[str, list[VertexBuildModel]] vertex_builds: dict[str, list[VertexBuildResponseModel]]
@classmethod @classmethod
def from_list_of_dicts(cls, vertex_build_dicts): def from_list_of_dicts(cls, vertex_build_dicts):
vertex_build_map = {} vertex_build_map = {}
for vertex_build_dict in vertex_build_dicts: for vertex_build_dict in vertex_build_dicts:
vertex_build = VertexBuildModel(**vertex_build_dict) vertex_build = VertexBuildResponseModel(**vertex_build_dict)
if vertex_build.id not in vertex_build_map: if vertex_build.id not in vertex_build_map:
vertex_build_map[vertex_build.id] = [] vertex_build_map[vertex_build.id] = []
vertex_build_map[vertex_build.id].append(vertex_build) vertex_build_map[vertex_build.id].append(vertex_build)