Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
24 changes: 24 additions & 0 deletions src/lean_spec/types/collections.py
Original file line number Diff line number Diff line change
Expand Up @@ -45,6 +45,25 @@ class Uint16Vector2(SSZVector):
data: Tuple[SSZType, ...] = Field(default_factory=tuple)
"""The immutable data stored in the vector."""

@field_serializer("data", when_used="json")
def _serialize_data(self, value: Tuple[SSZType, ...]) -> list[Any]:
"""Serialize vector elements to JSON, preserving custom type serialization."""
from lean_spec.subspecs.koalabear import Fp

result: list[Any] = []
for item in value:
# For BaseBytes subclasses, manually add 0x prefix
if isinstance(item, BaseBytes):
result.append("0x" + item.hex())
# For Fp field elements, extract the value attribute
elif isinstance(item, Fp):
result.append(item.value)
else:
# For other types (Uint, etc.), convert to int
# BaseUint inherits from int, so this cast is safe
result.append(item)
return result

@field_validator("data", mode="before")
@classmethod
def _validate_vector_data(cls, v: Any) -> Tuple[SSZType, ...]:
Expand Down Expand Up @@ -188,11 +207,16 @@ class Uint64List32(SSZList):
@field_serializer("data", when_used="json")
def _serialize_data(self, value: Tuple[SSZType, ...]) -> list[Any]:
"""Serialize list elements to JSON, preserving custom type serialization."""
from lean_spec.subspecs.koalabear import Fp

result: list[Any] = []
for item in value:
# For BaseBytes subclasses, manually add 0x prefix
if isinstance(item, BaseBytes):
result.append("0x" + item.hex())
# For Fp field elements, extract the value attribute
elif isinstance(item, Fp):
result.append(item.value)
else:
# For other types (Uint, etc.), convert to int
# BaseUint inherits from int, so this cast is safe
Expand Down
27 changes: 26 additions & 1 deletion tests/lean_spec/types/test_collections.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@
from pydantic import ValidationError, create_model
from typing_extensions import Type

from lean_spec.subspecs.koalabear import Fp
from lean_spec.types.boolean import Boolean
from lean_spec.types.collections import SSZList, SSZVector
from lean_spec.types.container import Container
Expand Down Expand Up @@ -149,6 +150,13 @@ class Uint8Vector2(SSZVector):
LENGTH = 2


class FpVector8(SSZVector):
"""A vector of exactly 8 Fp values."""

ELEMENT_TYPE = Fp
LENGTH = 8


# Additional List classes for tests
class Uint8List32(SSZList):
"""A list with up to 32 Uint8 values."""
Expand Down Expand Up @@ -178,6 +186,13 @@ class BooleanList4(SSZList):
LIMIT = 4


class FpList8(SSZList):
"""A list with up to 8 Fp values."""

ELEMENT_TYPE = Fp
LIMIT = 8


# Test data for the 'sig' vector test case
sig_test_data_list = [0] * 96
for i, v in {0: 1, 32: 2, 64: 3, 95: 0xFF}.items():
Expand Down Expand Up @@ -335,6 +350,11 @@ class TestSSZVectorSerialization:
(FixedContainer(a=Uint8(1), b=Uint16(2)), FixedContainer(a=Uint8(3), b=Uint16(4))),
"010200030400", # 010200 for first element, 030400 for second
),
(
FpVector8,
(10, 20, 30, 40, 50, 60, 70, 80),
"0a000000140000001e00000028000000320000003c0000004600000050000000",
),
],
)
def test_fixed_size_element_vector_serialization(
Expand Down Expand Up @@ -372,7 +392,7 @@ def test_variable_size_element_vector_serialization(self) -> None:
assert decoded == instance


class TestListSerialization:
class TestSSZListSerialization:
"""Tests SSZ serialization and deserialization for the List type."""

@pytest.mark.parametrize(
Expand All @@ -397,6 +417,11 @@ class TestListSerialization:
tuple(range(1, 20)),
"".join(i.to_bytes(32, "little").hex() for i in range(1, 20)),
),
(
FpList8,
(10, 20, 30),
"0a000000140000001e000000",
),
],
)
def test_fixed_size_element_list_serialization(
Expand Down
Loading