change symbol to use str object instead of bytes

This commit is contained in:
stuart committed 2026-02-26 10:02:56 +01:00
1 parent a0f3fa492f
commit 3e07027aa5
1 file changed
+3 -49
+3 -49
View File
@@ -1,6 +1,7 @@
from dataclasses import dataclass
from typing import Any
import numpy as np
from pydantic import BaseModel
from pydantic_tensor import Tensor
from typing import Any, Literal
from pydantic_core import core_schema
@@ -8,54 +9,7 @@ import base64
@dataclass
class Symbol:
class Symbol(BaseModel):
id: int
label: str
image: bytes
@classmethod
def from_dict(cls, data: dict) -> "Symbol":
if isinstance(data["image"], str):
image_bytes = base64.b64decode(data["image"])
else:
image_bytes = bytes(data["image"])
return cls(
id=data["id"],
label=data["label"],
image=image_bytes,
)
@classmethod
def __get_pydantic_core_schema__(cls, source_type: Any, handler):
from pydantic_core import core_schema
def validate_symbol(value):
if isinstance(value, cls):
return value
if isinstance(value, dict):
return cls.from_dict(value)
raise ValueError(f"Cannot convert {type(value)} to Symbol")
python_schema = core_schema.with_info_plain_validator_function(
lambda value, _: validate_symbol(value),
serialization=core_schema.plain_serializer_function_ser_schema(
lambda x: {
"id": x.id,
"label": x.label,
"image": base64.b64encode(x.image).decode("utf-8"),
},
when_used="json",
),
)
return core_schema.json_or_python_schema(
json_schema=core_schema.typed_dict_schema(
{
"id": core_schema.typed_dict_field(core_schema.int_schema()),
"label": core_schema.typed_dict_field(core_schema.str_schema()),
"image": core_schema.typed_dict_field(core_schema.str_schema()),
}
),
python_schema=python_schema,
)
image: str # Base64-encoded PNG image data