From 54e1d1d3d60ac03691cc7c629fcf34e41be9e933 Mon Sep 17 00:00:00 2001 From: Thomas Coratger Date: Wed, 7 Jan 2026 21:00:14 +0100 Subject: [PATCH] types: rm redundant ELEMENT_TYPE assignments --- .../subspecs/containers/block/types.py | 2 - .../subspecs/containers/state/types.py | 3 - src/lean_spec/subspecs/xmss/types.py | 5 -- src/lean_spec/types/collections.py | 56 ++++++++++----- tests/lean_spec/subspecs/ssz/test_hash.py | 18 ++--- tests/lean_spec/types/test_collections.py | 69 +++++++------------ tests/lean_spec/types/test_union.py | 6 +- 7 files changed, 71 insertions(+), 88 deletions(-) diff --git a/src/lean_spec/subspecs/containers/block/types.py b/src/lean_spec/subspecs/containers/block/types.py index b9aaa672..cd513e38 100644 --- a/src/lean_spec/subspecs/containers/block/types.py +++ b/src/lean_spec/subspecs/containers/block/types.py @@ -20,7 +20,6 @@ class AggregatedAttestations(SSZList[AggregatedAttestation]): """List of aggregated attestations included in a block.""" - ELEMENT_TYPE = AggregatedAttestation LIMIT = int(VALIDATOR_REGISTRY_LIMIT) @@ -35,5 +34,4 @@ class AttestationSignatures(SSZList[AggregatedSignatureProof]): - proof bytes from leanVM signature aggregation. """ - ELEMENT_TYPE = AggregatedSignatureProof LIMIT = int(VALIDATOR_REGISTRY_LIMIT) diff --git a/src/lean_spec/subspecs/containers/state/types.py b/src/lean_spec/subspecs/containers/state/types.py index 3bdf406d..effcfa89 100644 --- a/src/lean_spec/subspecs/containers/state/types.py +++ b/src/lean_spec/subspecs/containers/state/types.py @@ -12,14 +12,12 @@ class HistoricalBlockHashes(SSZList[Bytes32]): """List of historical block root hashes up to historical_roots_limit.""" - ELEMENT_TYPE = Bytes32 LIMIT = int(DEVNET_CONFIG.historical_roots_limit) class JustificationRoots(SSZList[Bytes32]): """List of justified block roots up to historical_roots_limit.""" - ELEMENT_TYPE = Bytes32 LIMIT = int(DEVNET_CONFIG.historical_roots_limit) @@ -38,5 +36,4 @@ class JustificationValidators(BaseBitlist): class Validators(SSZList[Validator]): """Validator registry tracked in the state.""" - ELEMENT_TYPE = Validator LIMIT = int(DEVNET_CONFIG.validator_registry_limit) diff --git a/src/lean_spec/subspecs/xmss/types.py b/src/lean_spec/subspecs/xmss/types.py index 076bf258..e2c4d78a 100644 --- a/src/lean_spec/subspecs/xmss/types.py +++ b/src/lean_spec/subspecs/xmss/types.py @@ -54,7 +54,6 @@ class HashDigestVector(SSZVector[Fp]): as SSZ can pack these back-to-back without per-element offsets. """ - ELEMENT_TYPE = Fp LENGTH = HASH_DIGEST_LENGTH @@ -67,7 +66,6 @@ class HashDigestList(SSZList[HashDigestVector]): This type is used to represent collections of hash digests in the XMSS scheme. """ - ELEMENT_TYPE = HashDigestVector LIMIT = NODE_LIST_LIMIT @@ -80,7 +78,6 @@ class Parameter(SSZVector[Fp]): certain cross-key attacks. It is public knowledge. """ - ELEMENT_TYPE = Fp LENGTH = TARGET_CONFIG.PARAMETER_LEN @@ -95,7 +92,6 @@ class Randomness(SSZVector[Fp]): SSZ notation: `Vector[Fp, RAND_LEN_FE]` """ - ELEMENT_TYPE = Fp LENGTH = TARGET_CONFIG.RAND_LEN_FE @@ -152,5 +148,4 @@ class HashTreeLayers(SSZList[HashTreeLayer]): - Maximum: `LOG_LIFETIME + 1` layers """ - ELEMENT_TYPE = HashTreeLayer LIMIT = LAYERS_LIMIT diff --git a/src/lean_spec/types/collections.py b/src/lean_spec/types/collections.py index d9d43a46..79dac736 100644 --- a/src/lean_spec/types/collections.py +++ b/src/lean_spec/types/collections.py @@ -38,7 +38,6 @@ Example: class Uint64Vector4(SSZVector[Uint64]): - ELEMENT_TYPE = Uint64 LENGTH = 4 vec = Uint64Vector4(data=[...]) @@ -54,12 +53,12 @@ class SSZVector(SSZModel, Generic[T]): The length is fixed at the type level and cannot change at runtime. Subclasses must define: - ELEMENT_TYPE: The SSZ type of each element LENGTH: The exact number of elements + The ELEMENT_TYPE is automatically inferred from the generic parameter. + Example: class Uint16Vector2(SSZVector[Uint16]): - ELEMENT_TYPE = Uint16 LENGTH = 2 vec = Uint16Vector2(data=[Uint16(1), Uint16(2)]) @@ -72,17 +71,29 @@ class Uint16Vector2(SSZVector[Uint16]): """ ELEMENT_TYPE: ClassVar[Type[SSZType]] - """The SSZ type of elements in this vector.""" + """The SSZ type of elements in this vector (auto-inferred from generic parameter).""" LENGTH: ClassVar[int] """The exact number of elements (fixed at the type level).""" data: Sequence[T] = Field(default_factory=tuple) - """ - The immutable sequence of elements. - Accepts lists or tuples on input; stored as a tuple after validation. - """ + def __init_subclass__(cls, **kwargs: Any) -> None: + """Automatically set ELEMENT_TYPE from the generic parameter.""" + super().__init_subclass__(**kwargs) + + # Skip if ELEMENT_TYPE is explicitly defined in this class + if "ELEMENT_TYPE" in cls.__dict__: + return + + # Extract type argument from Pydantic's generic metadata + for base in cls.__bases__: + metadata = getattr(base, "__pydantic_generic_metadata__", None) + if metadata and metadata.get("origin") is SSZVector: + args = metadata.get("args", ()) + if args: + cls.ELEMENT_TYPE = args[0] + return @field_serializer("data", when_used="json") def _serialize_data(self, value: Sequence[T]) -> list[Any]: @@ -243,12 +254,12 @@ class SSZList(SSZModel, Generic[T]): Unlike Vector, the length can vary at runtime. Subclasses must define: - ELEMENT_TYPE: The SSZ type of each element LIMIT: The maximum number of elements allowed + The ELEMENT_TYPE is automatically inferred from the generic parameter. + Example: class Uint64List32(SSZList[Uint64]): - ELEMENT_TYPE = Uint64 LIMIT = 32 my_list = Uint64List32(data=[Uint64(1), Uint64(2)]) @@ -262,17 +273,30 @@ class Uint64List32(SSZList[Uint64]): """ ELEMENT_TYPE: ClassVar[Type[SSZType]] - """The SSZ type of elements in this list.""" + """The SSZ type of elements in this list (auto-inferred from generic parameter).""" LIMIT: ClassVar[int] """The maximum number of elements allowed.""" data: Sequence[T] = Field(default_factory=tuple) - """ - The immutable sequence of elements. - - Accepts lists or tuples on input; stored as a tuple after validation. - """ + """The immutable sequence of elements.""" + + def __init_subclass__(cls, **kwargs: Any) -> None: + """Automatically set ELEMENT_TYPE from the generic parameter.""" + super().__init_subclass__(**kwargs) + + # Skip if ELEMENT_TYPE is explicitly defined in this class + if "ELEMENT_TYPE" in cls.__dict__: + return + + # Extract type argument from Pydantic's generic metadata + for base in cls.__bases__: + metadata = getattr(base, "__pydantic_generic_metadata__", None) + if metadata and metadata.get("origin") is SSZList: + args = metadata.get("args", ()) + if args: + cls.ELEMENT_TYPE = args[0] + return @field_serializer("data", when_used="json") def _serialize_data(self, value: Sequence[T]) -> list[Any]: diff --git a/tests/lean_spec/subspecs/ssz/test_hash.py b/tests/lean_spec/subspecs/ssz/test_hash.py index 692ded7c..2770c428 100644 --- a/tests/lean_spec/subspecs/ssz/test_hash.py +++ b/tests/lean_spec/subspecs/ssz/test_hash.py @@ -28,28 +28,23 @@ # Concrete SSZList classes for tests -class Uint16List32(SSZList): - ELEMENT_TYPE = Uint16 +class Uint16List32(SSZList[Uint16]): LIMIT = 32 -class Uint16List128(SSZList): - ELEMENT_TYPE = Uint16 +class Uint16List128(SSZList[Uint16]): LIMIT = 128 -class Uint16List1024(SSZList): - ELEMENT_TYPE = Uint16 +class Uint16List1024(SSZList[Uint16]): LIMIT = 1024 -class Uint32List128(SSZList): - ELEMENT_TYPE = Uint32 +class Uint32List128(SSZList[Uint32]): LIMIT = 128 -class Uint256List32(SSZList): - ELEMENT_TYPE = Uint256 +class Uint256List32(SSZList[Uint256]): LIMIT = 32 @@ -460,10 +455,9 @@ class Bitlist512(BaseBitlist): # Define explicit SSZVector types for testing our new approach -class Uint16Vector2(SSZVector): +class Uint16Vector2(SSZVector[Uint16]): """A vector of exactly 2 Uint16 values.""" - ELEMENT_TYPE = Uint16 LENGTH = 2 diff --git a/tests/lean_spec/types/test_collections.py b/tests/lean_spec/types/test_collections.py index 6529fb9c..33aaa394 100644 --- a/tests/lean_spec/types/test_collections.py +++ b/tests/lean_spec/types/test_collections.py @@ -18,10 +18,9 @@ # Define some List types that are needed for Container definitions -class Uint16List4(SSZList): +class Uint16List4(SSZList[Uint16]): """A list with up to 4 Uint16 values.""" - ELEMENT_TYPE = Uint16 LIMIT = 4 @@ -40,160 +39,138 @@ class VariableContainer(Container): # Define explicit SSZVector types for testing -class Uint16Vector2(SSZVector): +class Uint16Vector2(SSZVector[Uint16]): """A vector of exactly 2 Uint16 values.""" - ELEMENT_TYPE = Uint16 LENGTH = 2 -class Uint8Vector4(SSZVector): +class Uint8Vector4(SSZVector[Uint8]): """A vector of exactly 4 Uint8 values.""" - ELEMENT_TYPE = Uint8 LENGTH = 4 -class Uint8Vector48(SSZVector): +class Uint8Vector48(SSZVector[Uint8]): """A vector of exactly 48 Uint8 values.""" - ELEMENT_TYPE = Uint8 LENGTH = 48 -class Uint8Vector96(SSZVector): +class Uint8Vector96(SSZVector[Uint8]): """A vector of exactly 96 Uint8 values.""" - ELEMENT_TYPE = Uint8 LENGTH = 96 -class FixedContainerVector2(SSZVector): +class FixedContainerVector2(SSZVector[FixedContainer]): """A vector of exactly 2 FixedContainer values.""" - ELEMENT_TYPE = FixedContainer LENGTH = 2 -class VariableContainerVector2(SSZVector): +class VariableContainerVector2(SSZVector[VariableContainer]): """A vector of exactly 2 VariableContainer values.""" - ELEMENT_TYPE = VariableContainer LENGTH = 2 # Define explicit List types for testing -class Uint16List32(SSZList): +class Uint16List32(SSZList[Uint16]): """A list with up to 32 Uint16 values.""" - ELEMENT_TYPE = Uint16 LIMIT = 32 -class Uint8List10(SSZList): +class Uint8List10(SSZList[Uint8]): """A list with up to 10 Uint8 values.""" - ELEMENT_TYPE = Uint8 LIMIT = 10 -class Uint32List128(SSZList): +class Uint32List128(SSZList[Uint32]): """A list with up to 128 Uint32 values.""" - ELEMENT_TYPE = Uint32 LIMIT = 128 -class Uint256List32(SSZList): +class Uint256List32(SSZList[Uint256]): """A list with up to 32 Uint256 values.""" - ELEMENT_TYPE = Uint256 LIMIT = 32 -class Uint256List128(SSZList): +class Uint256List128(SSZList[Uint256]): """A list with up to 128 Uint256 values.""" - ELEMENT_TYPE = Uint256 LIMIT = 128 -class VariableContainerList2(SSZList): +class VariableContainerList2(SSZList[VariableContainer]): """A list with up to 2 VariableContainer values.""" - ELEMENT_TYPE = VariableContainer LIMIT = 2 # Additional SSZVector classes for tests -class Uint8Vector32(SSZVector): +class Uint8Vector32(SSZVector[Uint8]): """A vector of exactly 32 Uint8 values.""" - ELEMENT_TYPE = Uint8 LENGTH = 32 -class Uint16Vector32(SSZVector): +class Uint16Vector32(SSZVector[Uint16]): """A vector of exactly 32 Uint16 values.""" - ELEMENT_TYPE = Uint16 LENGTH = 32 -class Uint8Vector64(SSZVector): +class Uint8Vector64(SSZVector[Uint8]): """A vector of exactly 64 Uint8 values.""" - ELEMENT_TYPE = Uint8 LENGTH = 64 -class Uint8Vector2(SSZVector): +class Uint8Vector2(SSZVector[Uint8]): """A vector of exactly 2 Uint8 values.""" - ELEMENT_TYPE = Uint8 LENGTH = 2 -class FpVector8(SSZVector): +class FpVector8(SSZVector[Fp]): """A vector of exactly 8 Fp values.""" - ELEMENT_TYPE = Fp LENGTH = 8 # Additional List classes for tests -class Uint8List32(SSZList): +class Uint8List32(SSZList[Uint8]): """A list with up to 32 Uint8 values.""" - ELEMENT_TYPE = Uint8 LIMIT = 32 -class Uint8List64(SSZList): +class Uint8List64(SSZList[Uint8]): """A list with up to 64 Uint8 values.""" - ELEMENT_TYPE = Uint8 LIMIT = 64 -class Uint8List4(SSZList): +class Uint8List4(SSZList[Uint8]): """A list with up to 4 Uint8 values.""" - ELEMENT_TYPE = Uint8 LIMIT = 4 -class BooleanList4(SSZList): +class BooleanList4(SSZList[Boolean]): """A list with up to 4 Boolean values.""" - ELEMENT_TYPE = Boolean LIMIT = 4 -class FpList8(SSZList): +class FpList8(SSZList[Fp]): """A list with up to 8 Fp values.""" - ELEMENT_TYPE = Fp LIMIT = 8 diff --git a/tests/lean_spec/types/test_union.py b/tests/lean_spec/types/test_union.py index 59b010dd..297783dc 100644 --- a/tests/lean_spec/types/test_union.py +++ b/tests/lean_spec/types/test_union.py @@ -15,17 +15,15 @@ from lean_spec.types.union import SSZUnion -class Uint8Vector3(SSZVector): +class Uint8Vector3(SSZVector[Uint8]): """A vector of exactly 3 Uint8 values.""" - ELEMENT_TYPE = Uint8 LENGTH = 3 -class Uint16List8(SSZList): +class Uint16List8(SSZList[Uint16]): """A list with up to 8 Uint16 values.""" - ELEMENT_TYPE = Uint16 LIMIT = 8