refactor: Upgrade to Pydantic v2

Author: msbrogli

hathor/conf/settings.py

@@ -16,7 +16,7 @@
 from enum import StrEnum, auto, unique
 from math import log
 from pathlib import Path
-from typing import Any, NamedTuple, Optional, Union
+from typing import NamedTuple, Optional, Union
 
 import pydantic
 
@@ -561,37 +561,19 @@ def parse_hex_str(hex_str: Union[str, bytes]) -> bytes:
     return hex_str
 
 
-def _validate_consensus_algorithm(consensus_algorithm: ConsensusSettings, values: dict[str, Any]) -> ConsensusSettings:
-    """Validate that if Proof-of-Authority is enabled, block rewards must not be set."""
-    if consensus_algorithm.is_pow():
-        return consensus_algorithm
-
-    assert consensus_algorithm.is_poa()
-    blocks_per_halving = values.get('BLOCKS_PER_HALVING')
-    initial_token_units_per_block = values.get('INITIAL_TOKEN_UNITS_PER_BLOCK')
-    minimum_token_units_per_block = values.get('MINIMUM_TOKEN_UNITS_PER_BLOCK')
-    assert initial_token_units_per_block is not None, 'INITIAL_TOKEN_UNITS_PER_BLOCK must be set'
-    assert minimum_token_units_per_block is not None, 'MINIMUM_TOKEN_UNITS_PER_BLOCK must be set'
-
-    if blocks_per_halving is not None or initial_token_units_per_block != 0 or minimum_token_units_per_block != 0:
-        raise ValueError('PoA networks do not support block rewards')
+def _validate_consensus_algorithm(consensus_algorithm: ConsensusSettings) -> ConsensusSettings:
+    """Validate that if Proof-of-Authority is enabled, block rewards must not be set.
 
+    Note: Full validation with other fields is done in the named_tuple validator.
+    """
     return consensus_algorithm
 
 
-def _validate_tokens(genesis_tokens: int, values: dict[str, Any]) -> int:
-    """Validate genesis tokens."""
-    genesis_token_units = values.get('GENESIS_TOKEN_UNITS')
-    decimal_places = values.get('DECIMAL_PLACES')
-    assert genesis_token_units is not None, 'GENESIS_TOKEN_UNITS must be set'
-    assert decimal_places is not None, 'DECIMAL_PLACES must be set'
-
-    if genesis_tokens != genesis_token_units * (10 ** decimal_places):
-        raise ValueError(
-            f'invalid tokens: GENESIS_TOKENS={genesis_tokens}, GENESIS_TOKEN_UNITS={genesis_token_units}, '
-            f'DECIMAL_PLACES={decimal_places}',
-        )
+def _validate_tokens(genesis_tokens: int) -> int:
+    """Validate genesis tokens.
 
+    Note: Full validation with other fields is done in the named_tuple validator.
+    """
     return genesis_tokens
 
 
@@ -607,43 +589,41 @@ def _validate_token_deposit_percentage(token_deposit_percentage: float) -> float
 
 
 _VALIDATORS = dict(
-    _parse_hex_str=pydantic.validator(
+    _parse_hex_str=pydantic.field_validator(
         'P2PKH_VERSION_BYTE',
         'MULTISIG_VERSION_BYTE',
         'GENESIS_OUTPUT_SCRIPT',
         'GENESIS_BLOCK_HASH',
         'GENESIS_TX1_HASH',
         'GENESIS_TX2_HASH',
-        pre=True,
-        allow_reuse=True
+        mode='before',
     )(parse_hex_str),
-    _parse_soft_voided_tx_id=pydantic.validator(
+    _parse_soft_voided_tx_id=pydantic.field_validator(
         'SOFT_VOIDED_TX_IDS',
-        pre=True,
-        allow_reuse=True,
-        each_item=True
-    )(parse_hex_str),
-    _parse_skipped_verification_tx_id=pydantic.validator(
+        mode='before',
+    )(lambda v: [parse_hex_str(x) for x in v] if isinstance(v, list) else v),
+    _parse_skipped_verification_tx_id=pydantic.field_validator(
         'SKIP_VERIFICATION',
-        pre=True,
-        allow_reuse=True,
-        each_item=True
-    )(parse_hex_str),
-    _parse_checkpoints=pydantic.validator(
+        mode='before',
+    )(lambda v: [parse_hex_str(x) for x in v] if isinstance(v, list) else v),
+    _parse_checkpoints=pydantic.field_validator(
         'CHECKPOINTS',
-        pre=True
+        mode='before',
     )(_parse_checkpoints),
-    _parse_blueprints=pydantic.validator(
+    _parse_blueprints=pydantic.field_validator(
         'BLUEPRINTS',
-        pre=True
+        mode='before',
     )(_parse_blueprints),
-    _validate_consensus_algorithm=pydantic.validator(
-        'CONSENSUS_ALGORITHM'
+    _validate_consensus_algorithm=pydantic.field_validator(
+        'CONSENSUS_ALGORITHM',
+        mode='after',
     )(_validate_consensus_algorithm),
-    _validate_tokens=pydantic.validator(
-        'GENESIS_TOKENS'
+    _validate_tokens=pydantic.field_validator(
+        'GENESIS_TOKENS',
+        mode='after',
     )(_validate_tokens),
-    _validate_token_deposit_percentage=pydantic.validator(
-        'TOKEN_DEPOSIT_PERCENTAGE'
+    _validate_token_deposit_percentage=pydantic.field_validator(
+        'TOKEN_DEPOSIT_PERCENTAGE',
+        mode='after',
     )(_validate_token_deposit_percentage),
 )

hathor/consensus/consensus_settings.py

@@ -17,14 +17,14 @@
 import hashlib
 from abc import ABC, abstractmethod
 from enum import Enum, unique
-from typing import TYPE_CHECKING, Annotated, Any, Literal, TypeAlias
+from typing import TYPE_CHECKING, Annotated, Any, Literal, TypeAlias, Union
 
-from pydantic import Field, NonNegativeInt, PrivateAttr, validator
+from pydantic import Discriminator, NonNegativeInt, PrivateAttr, Tag, field_validator
 from typing_extensions import override
 
 from hathor.transaction import TxVersion
 from hathor.util import json_dumpb
-from hathor.utils.pydantic import BaseModel
+from hathor.utils.pydantic import BaseModel, Hex
 
 if TYPE_CHECKING:
     from hathor.conf.settings import HathorSettings
@@ -97,34 +97,19 @@ def get_peer_hello_hash(self) -> str | None:
 
 
 class PoaSignerSettings(BaseModel):
-    public_key: bytes
+    public_key: Hex[bytes]
     start_height: NonNegativeInt = 0
     end_height: NonNegativeInt | None = None
 
-    @validator('public_key', pre=True)
-    def _parse_hex_str(cls, hex_str: str | bytes) -> bytes:
-        from hathor.conf.settings import parse_hex_str
-        return parse_hex_str(hex_str)
-
-    @validator('end_height')
-    def _validate_end_height(cls, end_height: int | None, values: dict[str, Any]) -> int | None:
-        start_height = values.get('start_height')
-        assert start_height is not None, 'start_height must be set'
-
-        if end_height is None:
-            return None
-
-        if end_height <= start_height:
-            raise ValueError(f'end_height ({end_height}) must be greater than start_height ({start_height})')
-
-        return end_height
+    def __init__(self, **data: Any) -> None:
+        super().__init__(**data)
+        # Validate end_height > start_height after initialization
+        if self.end_height is not None and self.end_height <= self.start_height:
+            raise ValueError(f'end_height ({self.end_height}) must be greater than start_height ({self.start_height})')
 
     def to_json_dict(self) -> dict[str, Any]:
         """Return this signer settings instance as a json dict."""
-        json_dict = self.dict()
-        # TODO: We can use a custom serializer to convert bytes to hex when we update to Pydantic V2.
-        json_dict['public_key'] = self.public_key.hex()
-        return json_dict
+        return self.model_dump()
 
 
 class PoaSettings(_BaseConsensusSettings):
@@ -133,7 +118,8 @@ class PoaSettings(_BaseConsensusSettings):
     # A list of Proof-of-Authority signer public keys that have permission to produce blocks.
     signers: tuple[PoaSignerSettings, ...]
 
-    @validator('signers')
+    @field_validator('signers', mode='after')
+    @classmethod
     def _validate_signers(cls, signers: tuple[PoaSignerSettings, ...]) -> tuple[PoaSignerSettings, ...]:
         if len(signers) == 0:
             raise ValueError('At least one signer must be provided in PoA networks')
@@ -165,4 +151,10 @@ def _calculate_peer_hello_hash(self) -> str | None:
         return hashlib.sha256(data).digest().hex()
 
 
-ConsensusSettings: TypeAlias = Annotated[PowSettings | PoaSettings, Field(discriminator='type')]
+ConsensusSettings: TypeAlias = Annotated[
+    Union[
+        Annotated[PowSettings, Tag(ConsensusType.PROOF_OF_WORK)],
+        Annotated[PoaSettings, Tag(ConsensusType.PROOF_OF_AUTHORITY)],
+    ],
+    Discriminator('type')
+]

hathor/consensus/poa/poa_signer.py

@@ -19,7 +19,7 @@
 
 from cryptography.hazmat.primitives import hashes
 from cryptography.hazmat.primitives.asymmetric import ec
-from pydantic import Field, validator
+from pydantic import ConfigDict, Field, field_validator, model_validator
 
 from hathor.consensus import poa
 from hathor.crypto.util import (
@@ -33,46 +33,62 @@
     from hathor.transaction.poa import PoaBlock
 
 
-class PoaSignerFile(BaseModel, arbitrary_types_allowed=True):
+class PoaSignerFile(BaseModel):
     """Class that represents a Proof-of-Authority signer configuration file."""
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
     private_key: ec.EllipticCurvePrivateKeyWithSerialization = Field(alias='private_key_hex')
     public_key: ec.EllipticCurvePublicKey = Field(alias='public_key_hex')
     address: str
+    _public_key_hex: bytes = bytes()  # Store original hex for validation
 
-    @validator('private_key', pre=True)
+    @field_validator('private_key', mode='before')
+    @classmethod
     def _parse_private_key(cls, private_key_hex: str) -> ec.EllipticCurvePrivateKeyWithSerialization:
         """Parse a private key hex into a private key instance."""
         private_key_bytes = bytes.fromhex(private_key_hex)
         return get_private_key_from_bytes(private_key_bytes)
 
-    @validator('public_key', pre=True)
-    def _validate_public_key_first_bytes(
-        cls,
-        public_key_hex: str,
-        values: dict[str, Any]
-    ) -> ec.EllipticCurvePublicKey:
-        """Parse a public key hex into a public key instance, and validate that it corresponds to the private key."""
-        private_key = values.get('private_key')
-        assert isinstance(private_key, ec.EllipticCurvePrivateKey), 'private_key must be set'
-
-        public_key_bytes = bytes.fromhex(public_key_hex)
-        actual_public_key = private_key.public_key()
-
-        if public_key_bytes != get_public_key_bytes_compressed(actual_public_key):
+    @model_validator(mode='before')
+    @classmethod
+    def _store_public_key_hex(cls, data: dict[str, Any]) -> dict[str, Any]:
+        """Store the public key hex for validation and convert to actual key."""
+        public_key_hex = data.get('public_key_hex')
+        if public_key_hex is not None:
+            data['_public_key_hex_input'] = bytes.fromhex(public_key_hex)
+        return data
+
+    @model_validator(mode='after')
+    def _validate_keys_and_address(self) -> 'PoaSignerFile':
+        """Validate that public key and address correspond to the private key."""
+        actual_public_key = self.private_key.public_key()
+        actual_public_key_bytes = get_public_key_bytes_compressed(actual_public_key)
+
+        # Validate the provided public key matches the one derived from private key
+        provided_public_key_bytes = get_public_key_bytes_compressed(self.public_key)
+        if provided_public_key_bytes != actual_public_key_bytes:
             raise ValueError('invalid public key')
 
-        return actual_public_key
+        if self.address != get_address_b58_from_public_key(actual_public_key):
+            raise ValueError('invalid address')
 
-    @validator('address')
-    def _validate_address(cls, address: str, values: dict[str, Any]) -> str:
-        """Validate that the provided address corresponds to the provided private key."""
-        private_key = values.get('private_key')
-        assert isinstance(private_key, ec.EllipticCurvePrivateKey), 'private_key must be set'
+        return self
 
-        if address != get_address_b58_from_public_key(private_key.public_key()):
-            raise ValueError('invalid address')
+    @field_validator('public_key', mode='before')
+    @classmethod
+    def _parse_public_key(cls, public_key_hex: str | ec.EllipticCurvePublicKey) -> ec.EllipticCurvePublicKey:
+        """Parse public key hex to public key object."""
+        if isinstance(public_key_hex, ec.EllipticCurvePublicKey):
+            return public_key_hex
+        # The public key is provided as compressed bytes
+        public_key_bytes = bytes.fromhex(public_key_hex)
+        # For compressed public keys, we need to use a different approach
+        from cryptography.hazmat.primitives.asymmetric import ec as ec_module
 
-        return address
+        # Load compressed public key
+        return ec_module.EllipticCurvePublicKey.from_encoded_point(
+            ec_module.SECP256K1(), public_key_bytes
+        )
 
     def get_signer(self) -> PoaSigner:
         """Get a PoaSigner for this file."""

hathor/event/model/base_event.py

@@ -12,17 +12,19 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Any, Optional
+from typing import Optional
 
-from pydantic import NonNegativeInt, validator
+from pydantic import ConfigDict, NonNegativeInt, model_validator
 
-from hathor.event.model.event_data import BaseEventData, EventData
+from hathor.event.model.event_data import EventData
 from hathor.event.model.event_type import EventType
 from hathor.pubsub import EventArguments
 from hathor.utils.pydantic import BaseModel
 
 
-class BaseEvent(BaseModel, use_enum_values=True):
+class BaseEvent(BaseModel):
+    model_config = ConfigDict(use_enum_values=True)
+
     # Event unique id, determines event order
     id: NonNegativeInt
     # Timestamp in which the event was emitted, this follows the unix_timestamp format, it's only informative, events
@@ -57,12 +59,12 @@ def from_event_arguments(
             group_id=group_id,
         )
 
-    @validator('data')
-    def data_type_must_match_event_type(cls, v: BaseEventData, values: dict[str, Any]) -> BaseEventData:
-        event_type = EventType(values['type'])
+    @model_validator(mode='after')
+    def data_type_must_match_event_type(self) -> 'BaseEvent':
+        event_type = EventType(self.type)
         expected_data_type = event_type.data_type()
 
-        if type(v) is not expected_data_type:
+        if type(self.data) is not expected_data_type:
             raise ValueError('event data type does not match event type')
 
-        return v
+        return self