From d2eb4c728c7d298d60fd8b660abf43ea969b58dd Mon Sep 17 00:00:00 2001 From: unnawut Date: Fri, 5 Dec 2025 21:45:09 +0700 Subject: [PATCH 1/6] feat: make Fp and its vector json serializable --- src/lean_spec/subspecs/koalabear/field.py | 37 ++++++++++++++++++++++- src/lean_spec/types/collections.py | 24 +++++++++++++++ 2 files changed, 60 insertions(+), 1 deletion(-) diff --git a/src/lean_spec/subspecs/koalabear/field.py b/src/lean_spec/subspecs/koalabear/field.py index 0c5ea9b0f..0cd8bdafc 100644 --- a/src/lean_spec/subspecs/koalabear/field.py +++ b/src/lean_spec/subspecs/koalabear/field.py @@ -1,6 +1,9 @@ """Core definition of the KoalaBear prime field Fp.""" -from typing import IO, Self +from typing import IO, Any, Self + +from pydantic.annotated_handlers import GetCoreSchemaHandler +from pydantic_core import core_schema from lean_spec.types import SSZType @@ -97,6 +100,38 @@ def __init__(self, value: int) -> None: # Normalize to [0, P) - handles negative values correctly self.value: int = value % P + @classmethod + def __get_pydantic_core_schema__( + cls, source_type: Any, handler: GetCoreSchemaHandler + ) -> core_schema.CoreSchema: + """ + Hook into Pydantic's validation system for strict field element validation. + + This schema ensures that only values in the range [0, P) are accepted + during Pydantic model validation. + """ + # Validator that takes an integer and returns an Fp instance + from_int_validator = core_schema.no_info_plain_validator_function(cls) + + # Schema that first validates the input as an int in the range [0, P), + # then calls our validator + python_schema = core_schema.chain_schema( + [core_schema.int_schema(ge=0, lt=P, strict=True), from_int_validator] + ) + + return core_schema.union_schema( + [ + # Case 1: The value is already our field element type + core_schema.is_instance_schema(cls), + # Case 2: The value is a standard int and needs to be parsed and validated + python_schema, + ], + # For serialization (e.g., to JSON), convert the instance back to a plain int. + serialization=core_schema.plain_serializer_function_ser_schema( + lambda x: x.value, return_schema=core_schema.int_schema() + ), + ) + @classmethod def is_fixed_size(cls) -> bool: """Fp elements are fixed-size (4 bytes).""" diff --git a/src/lean_spec/types/collections.py b/src/lean_spec/types/collections.py index 71db39250..7f0f0bd9a 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_vector_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 From 0c62ec518e9785e03852f12e05db456666bf8e24 Mon Sep 17 00:00:00 2001 From: unnawut Date: Fri, 5 Dec 2025 22:10:12 +0700 Subject: [PATCH 2/6] fix: remove unneeded Fp.__get_pydantic_core_schema() --- src/lean_spec/subspecs/koalabear/field.py | 32 ----------------------- 1 file changed, 32 deletions(-) diff --git a/src/lean_spec/subspecs/koalabear/field.py b/src/lean_spec/subspecs/koalabear/field.py index 0cd8bdafc..cbb2b49c4 100644 --- a/src/lean_spec/subspecs/koalabear/field.py +++ b/src/lean_spec/subspecs/koalabear/field.py @@ -100,38 +100,6 @@ def __init__(self, value: int) -> None: # Normalize to [0, P) - handles negative values correctly self.value: int = value % P - @classmethod - def __get_pydantic_core_schema__( - cls, source_type: Any, handler: GetCoreSchemaHandler - ) -> core_schema.CoreSchema: - """ - Hook into Pydantic's validation system for strict field element validation. - - This schema ensures that only values in the range [0, P) are accepted - during Pydantic model validation. - """ - # Validator that takes an integer and returns an Fp instance - from_int_validator = core_schema.no_info_plain_validator_function(cls) - - # Schema that first validates the input as an int in the range [0, P), - # then calls our validator - python_schema = core_schema.chain_schema( - [core_schema.int_schema(ge=0, lt=P, strict=True), from_int_validator] - ) - - return core_schema.union_schema( - [ - # Case 1: The value is already our field element type - core_schema.is_instance_schema(cls), - # Case 2: The value is a standard int and needs to be parsed and validated - python_schema, - ], - # For serialization (e.g., to JSON), convert the instance back to a plain int. - serialization=core_schema.plain_serializer_function_ser_schema( - lambda x: x.value, return_schema=core_schema.int_schema() - ), - ) - @classmethod def is_fixed_size(cls) -> bool: """Fp elements are fixed-size (4 bytes).""" From 51adfec1691d89f360e08f55a6ee81eded8d85e5 Mon Sep 17 00:00:00 2001 From: unnawut Date: Fri, 5 Dec 2025 22:11:54 +0700 Subject: [PATCH 3/6] fix: remove no longer needed imports --- src/lean_spec/subspecs/koalabear/field.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/src/lean_spec/subspecs/koalabear/field.py b/src/lean_spec/subspecs/koalabear/field.py index cbb2b49c4..0c5ea9b0f 100644 --- a/src/lean_spec/subspecs/koalabear/field.py +++ b/src/lean_spec/subspecs/koalabear/field.py @@ -1,9 +1,6 @@ """Core definition of the KoalaBear prime field Fp.""" -from typing import IO, Any, Self - -from pydantic.annotated_handlers import GetCoreSchemaHandler -from pydantic_core import core_schema +from typing import IO, Self from lean_spec.types import SSZType From 55a50bd742ee68db575a9ce69eaeb4787657a729 Mon Sep 17 00:00:00 2001 From: unnawut Date: Fri, 5 Dec 2025 23:29:22 +0700 Subject: [PATCH 4/6] fix: align naming --- src/lean_spec/types/collections.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/lean_spec/types/collections.py b/src/lean_spec/types/collections.py index 7f0f0bd9a..95a6d5d34 100644 --- a/src/lean_spec/types/collections.py +++ b/src/lean_spec/types/collections.py @@ -46,7 +46,7 @@ class Uint16Vector2(SSZVector): """The immutable data stored in the vector.""" @field_serializer("data", when_used="json") - def _serialize_vector_data(self, value: Tuple[SSZType, ...]) -> list[Any]: + 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 From 728f736abc70b7513e7cfd127bed0711514e63a8 Mon Sep 17 00:00:00 2001 From: unnawut Date: Fri, 5 Dec 2025 23:32:02 +0700 Subject: [PATCH 5/6] test: Fp list and vector --- tests/lean_spec/types/test_collections.py | 28 ++++++++++++++++++++++- 1 file changed, 27 insertions(+), 1 deletion(-) diff --git a/tests/lean_spec/types/test_collections.py b/tests/lean_spec/types/test_collections.py index 2fbc6c74c..2c1916b28 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,12 @@ 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), + "0a000000140000001e0000002800000032000000" + "3c0000004600000050000000", + ), ], ) def test_fixed_size_element_vector_serialization( @@ -372,7 +393,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 +418,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( From 2d0f421ca96a9b1161c177cd3a8ba1d8b709337c Mon Sep 17 00:00:00 2001 From: unnawut Date: Fri, 5 Dec 2025 23:33:26 +0700 Subject: [PATCH 6/6] fix: linting --- tests/lean_spec/types/test_collections.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/tests/lean_spec/types/test_collections.py b/tests/lean_spec/types/test_collections.py index 2c1916b28..c81a9cc2f 100644 --- a/tests/lean_spec/types/test_collections.py +++ b/tests/lean_spec/types/test_collections.py @@ -353,8 +353,7 @@ class TestSSZVectorSerialization: ( FpVector8, (10, 20, 30, 40, 50, 60, 70, 80), - "0a000000140000001e0000002800000032000000" - "3c0000004600000050000000", + "0a000000140000001e00000028000000320000003c0000004600000050000000", ), ], )