diff --git a/src/lean_spec/types/collections.py b/src/lean_spec/types/collections.py index 71db3925..95a6d5d3 100644 --- a/src/lean_spec/types/collections.py +++ b/src/lean_spec/types/collections.py @@ -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, ...]: @@ -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 diff --git a/tests/lean_spec/types/test_collections.py b/tests/lean_spec/types/test_collections.py index 2fbc6c74..c81a9cc2 100644 --- a/tests/lean_spec/types/test_collections.py +++ b/tests/lean_spec/types/test_collections.py @@ -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 @@ -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.""" @@ -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(): @@ -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( @@ -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( @@ -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(