Skip to content

Commit

Permalink
Fix: Refactor ItemHash to align with Pydantic v2 validation schema.
Browse files Browse the repository at this point in the history
Replaced `__get_pydantic_core_schema__` with a more efficient schema
handling using `core_schema.str_schema()` and custom validation for ItemHash.
  • Loading branch information
Antonyjin committed Sep 27, 2024
1 parent 1cb37b6 commit 8d6ac80
Showing 1 changed file with 12 additions and 8 deletions.
20 changes: 12 additions & 8 deletions aleph_message/models/item_hash.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,8 @@
from functools import lru_cache

from ..exceptions import UnknownHashError

from pydantic_core import CoreSchema, core_schema
from pydantic import GetCoreSchemaHandler

class ItemType(str, Enum):
"""Item storage options"""
Expand Down Expand Up @@ -45,18 +46,21 @@ def __new__(cls, value: str):
return obj

@classmethod
def __get_validators__(cls):
# one or more validators may be yielded which will be called in the
# order to validate the input, each validator will receive as an input
# the value returned from the previous validator
yield cls.validate
def __get_pydantic_core_schema__(cls, source, handler: GetCoreSchemaHandler) -> core_schema.CoreSchema:
# This function validates the input after the initial type validation (as a string).
# The returned value from this function will be used as the final validated value.

# Return a string schema and add a post-validation function to convert to ItemHash
return core_schema.no_info_after_validator_function(
cls.validate,
core_schema.str_schema()
)

@classmethod
def validate(cls, v):
if not isinstance(v, str):
raise TypeError("Item hash must be a string")

return cls(v)
return cls(v) # Convert to ItemHash

def __repr__(self):
return f"<ItemHash value={super().__repr__()} item_type={self.item_type!r}>"

0 comments on commit 8d6ac80

Please sign in to comment.