types: small refactor (leanEthereum#65)

Author: tcoratger

src/lean_spec/types/bitfields.py

@@ -54,16 +54,12 @@ def _validate_vector_data(cls, v: Any) -> Tuple[Boolean, ...]:
         if not isinstance(v, (list, tuple)):
             v = tuple(v)
 
-        # Convert each bit to Boolean
-        typed_values = tuple(Boolean(item) if not isinstance(item, Boolean) else item for item in v)
-
-        if len(typed_values) != cls.LENGTH:
+        if len(v) != cls.LENGTH:
             raise ValueError(
-                f"{cls.__name__} requires exactly {cls.LENGTH} bits, "
-                f"but {len(typed_values)} were provided."
+                f"{cls.__name__} requires exactly {cls.LENGTH} bits, but {len(v)} were provided."
             )
 
-        return typed_values
+        return tuple(Boolean(item) for item in v)
 
     @classmethod
     def is_fixed_size(cls) -> bool:
@@ -104,9 +100,7 @@ def encode_bytes(self) -> bytes:
         byte_array = bytearray(byte_len)
         for i, bit in enumerate(self.data):
             if bit:
-                byte_index = i // 8
-                bit_index_in_byte = i % 8
-                byte_array[byte_index] |= 1 << bit_index_in_byte
+                byte_array[i // 8] |= 1 << (i % 8)
         return bytes(byte_array)
 
     @classmethod
@@ -120,8 +114,8 @@ def decode_bytes(cls, data: bytes) -> Self:
         if len(data) != expected_len:
             raise ValueError(f"{cls.__name__} expected {expected_len} bytes, got {len(data)}")
 
-        bits_generator = (Boolean((data[i // 8] >> (i % 8)) & 1) for i in range(cls.LENGTH))
-        return cls(data=tuple(bits_generator))
+        bits = tuple(Boolean((data[i // 8] >> (i % 8)) & 1) for i in range(cls.LENGTH))
+        return cls(data=bits)
 
 
 class BaseBitlist(SSZModel):
@@ -158,18 +152,10 @@ def _validate_list_data(cls, v: Any) -> Tuple[Boolean, ...]:
                 f"{cls.__name__} cannot contain more than {cls.LIMIT} bits, got {len(elements)}"
             )
 
-        # Convert and validate each bit
-        typed_values = []
-        for i, element in enumerate(elements):
-            if isinstance(element, Boolean):
-                typed_values.append(element)
-            else:
-                try:
-                    typed_values.append(Boolean(element))
-                except Exception as e:
-                    raise ValueError(f"Bit {i} cannot be converted to Boolean: {e}") from e
-
-        return tuple(typed_values)
+        try:
+            return tuple(Boolean(element) for element in elements)
+        except Exception as e:
+            raise ValueError(f"Cannot convert elements to Boolean: {e}") from e
 
     def __getitem__(self, key: int | slice) -> Boolean | tuple[Boolean, ...]:
         """Get a bit by index or slice."""
@@ -219,19 +205,15 @@ def encode_bytes(self) -> bytes:
         # Pack data bits.
         for i, bit in enumerate(self.data):
             if bit:
-                byte_index = i // 8
-                bit_index_in_byte = i % 8
-                byte_array[byte_index] |= 1 << bit_index_in_byte
+                byte_array[i // 8] |= 1 << (i % 8)
 
         # Place delimiter bit (1) immediately after the last data bit.
         if num_bits % 8 == 0:
             # Delimiter starts a new byte.
             return bytes(byte_array) + b"\x01"
         else:
             # Delimiter lives in the last byte at position num_bits % 8.
-            delimiter_byte_index = num_bits // 8
-            delimiter_bit_index = num_bits % 8
-            byte_array[delimiter_byte_index] |= 1 << delimiter_bit_index
+            byte_array[num_bits // 8] |= 1 << (num_bits % 8)
             return bytes(byte_array)
 
     @classmethod
@@ -250,11 +232,9 @@ def decode_bytes(cls, data: bytes) -> Self:
         for byte_idx in range(len(data) - 1, -1, -1):
             byte_val = data[byte_idx]
             if byte_val != 0:
-                # Find the rightmost 1 bit in this byte.
-                for bit_idx in range(7, -1, -1):
-                    if (byte_val >> bit_idx) & 1:
-                        delimiter_pos = byte_idx * 8 + bit_idx
-                        break
+                # Find the highest set bit in this byte using bit_length
+                bit_idx = byte_val.bit_length() - 1
+                delimiter_pos = byte_idx * 8 + bit_idx
                 break
 
         if delimiter_pos is None:
@@ -267,14 +247,5 @@ def decode_bytes(cls, data: bytes) -> Self:
                 f"{cls.__name__} decoded length {num_data_bits} exceeds limit {cls.LIMIT}"
             )
 
-        bits = []
-        for i in range(num_data_bits):
-            byte_index = i // 8
-            bit_index_in_byte = i % 8
-            if byte_index < len(data):
-                bit_value = bool((data[byte_index] >> bit_index_in_byte) & 1)
-            else:
-                bit_value = False
-            bits.append(bit_value)
-
+        bits = [bool((data[i // 8] >> (i % 8)) & 1) for i in range(num_data_bits)]
         return cls(data=bits)

src/lean_spec/types/boolean.py

@@ -37,11 +37,10 @@ def __new__(cls, value: bool | int) -> Self:
         if not isinstance(value, int):
             raise TypeError(f"Expected bool or int, got {type(value).__name__}")
 
-        int_value = int(value)
-        if int_value not in (0, 1):
-            raise ValueError(f"Boolean value must be 0 or 1, not {int_value}")
+        if value not in (0, 1):
+            raise ValueError(f"Boolean value must be 0 or 1, not {value}")
 
-        return super().__new__(cls, int_value)
+        return super().__new__(cls, value)
 
     @classmethod
     def __get_pydantic_core_schema__(
@@ -142,7 +141,7 @@ def __and__(self, other: Any) -> Self:
         """Handle the bitwise AND operator (`&`) strictly."""
         if not isinstance(other, type(self)):
             self._raise_type_error(other, "&")
-        return type(self)(super().__and__(other))
+        return type(self)(int(self) & int(other))
 
     def __rand__(self, other: Any) -> Self:
         """Handle the reverse bitwise AND operator (`&`) strictly."""
@@ -152,7 +151,7 @@ def __or__(self, other: Any) -> Self:
         """Handle the bitwise OR operator (`|`) strictly."""
         if not isinstance(other, type(self)):
             self._raise_type_error(other, "|")
-        return type(self)(super().__or__(other))
+        return type(self)(int(self) | int(other))
 
     def __ror__(self, other: Any) -> Self:
         """Handle the reverse bitwise OR operator (`|`) strictly."""
@@ -162,7 +161,7 @@ def __xor__(self, other: Any) -> Self:
         """Handle the bitwise XOR operator (`^`) strictly."""
         if not isinstance(other, type(self)):
             self._raise_type_error(other, "^")
-        return type(self)(super().__xor__(other))
+        return type(self)(int(self) ^ int(other))
 
     def __rxor__(self, other: Any) -> Self:
         """Handle the reverse bitwise XOR operator (`^`) strictly."""
@@ -176,9 +175,7 @@ def __eq__(self, other: object) -> bool:
 
         It returns `False` for all other types.
         """
-        if isinstance(other, int):
-            return int(self) == int(other)
-        return False
+        return isinstance(other, int) and int(self) == int(other)
 
     def __ne__(self, other: object) -> bool:
         """

src/lean_spec/types/byte_arrays.py

@@ -34,9 +34,8 @@ def _coerce_to_bytes(value: Any) -> bytes:
     if isinstance(value, (bytes, bytearray)):
         return bytes(value)
     if isinstance(value, str):
-        s = value[2:] if value.startswith("0x") else value
         # bytes.fromhex handles empty string and validates hex characters
-        return bytes.fromhex(s)
+        return bytes.fromhex(value.removeprefix("0x"))
     if isinstance(value, Iterable):
         # bytes(bytearray(iterable)) enforces each element is an int in 0..255
         return bytes(bytearray(value))
@@ -185,9 +184,7 @@ def __hash__(self) -> int:
 
     def hex(self, sep: str | bytes | None = None, bytes_per_sep: SupportsIndex = 1) -> str:
         """Return the hexadecimal string representation of the underlying bytes."""
-        if sep is None:
-            return bytes(self).hex()
-        return bytes(self).hex(sep, bytes_per_sep)
+        return bytes(self).hex() if sep is None else bytes(self).hex(sep, bytes_per_sep)
 
 
 class Bytes1(BaseBytes):
@@ -288,11 +285,11 @@ def deserialize(cls, stream: IO[bytes], scope: int) -> Self:
         """
         if scope < 0:
             raise ValueError("Invalid scope for ByteList: negative")
+        if scope > cls.LIMIT:
+            raise ValueError(f"ByteList[{cls.LIMIT}] scope {scope} exceeds limit")
         data = stream.read(scope)
         if len(data) != scope:
             raise IOError("Stream ended prematurely while decoding ByteList")
-        if len(data) > cls.LIMIT:
-            raise ValueError(f"ByteList[{cls.LIMIT}] decoded length {len(data)} exceeds limit")
         return cls(data=data)
 
     def encode_bytes(self) -> bytes:
@@ -316,8 +313,6 @@ def __bytes__(self) -> bytes:
 
     def __add__(self, other: Any) -> bytes:
         """Return the concatenation of the byte list and the argument."""
-        if isinstance(other, (bytes, bytearray)):
-            return self.data + bytes(other)
         return self.data + bytes(other)
 
     def __radd__(self, other: Any) -> bytes:

src/lean_spec/types/collections.py

@@ -84,10 +84,7 @@ def serialize(self, stream: IO[bytes]) -> int:
         """Serialize the vector to a binary stream."""
         # If elements are fixed-size, serialize them back-to-back.
         if self.is_fixed_size():
-            total_bytes_written = 0
-            for element in self.data:
-                total_bytes_written += element.serialize(stream)
-            return total_bytes_written
+            return sum(element.serialize(stream) for element in self.data)
         # If elements are variable-size, serialize their offsets, then their data.
         else:
             # Use a temporary in-memory stream to hold the serialized variable data.
@@ -230,10 +227,7 @@ def serialize(self, stream: IO[bytes]) -> int:
         # Lists are always variable-size, so we serialize offsets + data
         if self.ELEMENT_TYPE.is_fixed_size():
             # Fixed-size elements: serialize them back-to-back
-            total_bytes_written = 0
-            for element in self.data:
-                total_bytes_written += element.serialize(stream)
-            return total_bytes_written
+            return sum(element.serialize(stream) for element in self.data)
         else:
             # Variable-size elements: serialize offsets, then data
             variable_data_stream = io.BytesIO()
@@ -268,10 +262,10 @@ def deserialize(cls, stream: IO[bytes], scope: int) -> Self:
             return cls(data=elements)
         else:
             # Variable-size elements: read offsets first, then data
-            if scope < OFFSET_BYTE_LENGTH:
+            if scope == 0:
                 # Empty list case
-                if scope == 0:
-                    return cls(data=[])
+                return cls(data=[])
+            if scope < OFFSET_BYTE_LENGTH:
                 raise ValueError(f"Invalid scope for variable-size list: {scope}")
 
             # Read the first offset to determine the number of elements.

src/lean_spec/types/container.py

@@ -34,14 +34,12 @@ def _get_ssz_field_type(annotation: Any) -> Type[SSZType]:
         TypeError: If the annotation is not a valid SSZType class.
     """
     # Check if it's a class and is a subclass of SSZType
-    if inspect.isclass(annotation) and issubclass(annotation, SSZType):
-        return annotation
-
-    # If we get here, the annotation is not a valid SSZType
-    raise TypeError(
-        f"Field annotation {annotation} is not a valid SSZType class. "
-        f"Container fields must be concrete SSZType subclasses."
-    )
+    if not (inspect.isclass(annotation) and issubclass(annotation, SSZType)):
+        raise TypeError(
+            f"Field annotation {annotation} is not a valid SSZType class. "
+            f"Container fields must be concrete SSZType subclasses."
+        )
+    return annotation
 
 
 class Container(SSZModel):
@@ -134,7 +132,7 @@ def serialize(self, stream: IO[bytes]) -> int:
         variable_data = []
 
         # Process each field in definition order
-        for field_name, _field_info in type(self).model_fields.items():
+        for field_name in type(self).model_fields:
             # Get the field value and its type
             value = getattr(self, field_name)
             # Use the actual runtime type of the value, which should be an SSZType
@@ -160,7 +158,7 @@ def serialize(self, stream: IO[bytes]) -> int:
             if part:  # Fixed-size field data
                 stream.write(part)
             else:  # Variable-size field offset
-                Uint32(offset).serialize(stream)
+                stream.write(Uint32(offset).encode_bytes())
                 offset += len(variable_data[var_index])
                 var_index += 1