diff --git a/aleph_message/models/__init__.py b/aleph_message/models/__init__.py index 31fd060..32fdba7 100644 --- a/aleph_message/models/__init__.py +++ b/aleph_message/models/__init__.py @@ -7,7 +7,6 @@ from pathlib import Path from typing import Any, Dict, List, Literal, Optional, Type, TypeVar, Union, cast -from aleph_message.models.execution.volume import EphemeralVolumeSize, PersistentVolumeSizeMib from pydantic import BaseModel, Field, field_validator, ConfigDict from typing_extensions import TypeAlias @@ -265,7 +264,10 @@ def convert_float_to_datetime(cls, v, values): assert isinstance(v, datetime.datetime) return v - model_config = ConfigDict(extra="forbid", exclude={"id_", "_id"}) + model_config = ConfigDict(extra="forbid") + + def custom_dump(self): + return self.model_dump(exclude={"id_", "_id"}) class PostMessage(BaseMessage): @@ -306,36 +308,38 @@ def cannot_be_forgotten(cls, v: Optional[List[str]], values) -> Optional[List[st class ProgramMessage(BaseMessage): type: Literal[MessageType.program] - content: Optional[ProgramContent] = None + content: ProgramContent forgotten_by: Optional[List[str]] = None @staticmethod - def normalize_content(content): + def normalize_content(content: Union[Dict[str, Any], Any]) -> Any: if not isinstance(content, dict): return content - normalized_content = {} + normalized_content: Dict[str, Any] = {} + for key, value in content.items(): - # Handle special cases such as ItemHash, Enum, list and dict if isinstance(value, ItemHash): - normalized_content[key] = str(value) # ItemHash to string + normalized_content[key] = str(value) elif isinstance(value, Enum): - normalized_content[key] = value.value # Enum to string + normalized_content[key] = value.value elif isinstance(value, list): - normalized_content[key] = [ProgramMessage.normalize_content(v) for v in value] + if key == "volumes" and all(isinstance(v, str) for v in value): + normalized_content[key] = value + else: + normalized_content[key] = [ + ProgramMessage.normalize_content(v) for v in value + ] elif isinstance(value, dict): - # Special case for 'size_mib' and 'data' - if key == 'size_mib': + if key == "size_mib": normalized_content[key] = list(value.values())[0] else: normalized_content[key] = ProgramMessage.normalize_content(value) else: - normalized_content[key] = value # Keep the value as is + normalized_content[key] = value return normalized_content - - @field_validator("content") def check_content(cls, v, values): item_type = values.data.get("item_type") @@ -343,7 +347,7 @@ def check_content(cls, v, values): item_content = json.loads(values.data.get("item_content")) # Normalizing content to fit the structure of item_content - normalized_content = cls.normalize_content(v.dict(exclude_none=True)) + normalized_content = cls.normalize_content(v.model_dump(exclude_none=True)) if normalized_content != item_content: # Print les différences @@ -354,9 +358,6 @@ def check_content(cls, v, values): return v - - - class InstanceMessage(BaseMessage): type: Literal[MessageType.instance] content: InstanceContent diff --git a/aleph_message/models/execution/instance.py b/aleph_message/models/execution/instance.py index 4c54b25..3dc66d3 100644 --- a/aleph_message/models/execution/instance.py +++ b/aleph_message/models/execution/instance.py @@ -8,6 +8,7 @@ from .abstract import BaseExecutableContent from .environment import InstanceEnvironment from .volume import ParentVolume, PersistentVolumeSizeMib, VolumePersistence +from .base import Payment class RootfsVolume(HashableModel): @@ -24,7 +25,7 @@ class RootfsVolume(HashableModel): size_mib: PersistentVolumeSizeMib forgotten_by: Optional[List[str]] = None - @field_validator('size_mib', mode="before") + @field_validator("size_mib", mode="before") def convert_size_mib(cls, v): if isinstance(v, int): return PersistentVolumeSizeMib(persistent_volume_size=v) @@ -35,7 +36,7 @@ class InstanceContent(BaseExecutableContent): """Message content for scheduling a VM instance on the network.""" metadata: Optional[dict] = None - payment: Optional[dict] = None + payment: Optional[Payment] = None environment: InstanceEnvironment = Field( description="Properties of the instance execution environment" ) diff --git a/aleph_message/models/execution/program.py b/aleph_message/models/execution/program.py index ee3237c..9bb2228 100644 --- a/aleph_message/models/execution/program.py +++ b/aleph_message/models/execution/program.py @@ -7,7 +7,7 @@ from ..abstract import HashableModel from ..item_hash import ItemHash from .abstract import BaseExecutableContent -from .base import Encoding, Interface, MachineType +from .base import Encoding, Interface, MachineType, Payment from .environment import FunctionTriggers @@ -72,4 +72,4 @@ class ProgramContent(BaseExecutableContent): metadata: Optional[dict] = None authorized_keys: Optional[List[str]] = None - payment: Optional[dict] = None + payment: Optional[Payment] = None diff --git a/aleph_message/models/execution/volume.py b/aleph_message/models/execution/volume.py index adbfe3c..7e215e1 100644 --- a/aleph_message/models/execution/volume.py +++ b/aleph_message/models/execution/volume.py @@ -22,7 +22,7 @@ def is_read_only(self): ... class ImmutableVolume(AbstractVolume): - ref: ItemHash = None + ref: Optional[ItemHash] = None use_latest: bool = True def is_read_only(self): @@ -30,9 +30,7 @@ def is_read_only(self): class EphemeralVolumeSize(BaseModel): - ephemeral_volume_size: int = Field(gt=0, - le=1000, #Limit to 1GiB - strict=True) + ephemeral_volume_size: int = Field(gt=-1, le=1000, strict=True) # Limit to 1GiB def __hash__(self): return hash(self.ephemeral_volume_size) @@ -40,9 +38,9 @@ def __hash__(self): class EphemeralVolume(AbstractVolume): ephemeral: Literal[True] = True - size_mib: EphemeralVolumeSize = 0 + size_mib: EphemeralVolumeSize = EphemeralVolumeSize(ephemeral_volume_size=0) - @field_validator('size_mib', mode="before") + @field_validator("size_mib", mode="before") def convert_size_mib(cls, v): if isinstance(v, int): return EphemeralVolumeSize(ephemeral_volume_size=v) @@ -67,9 +65,9 @@ class VolumePersistence(str, Enum): class PersistentVolumeSizeMib(BaseModel): - persistent_volume_size: int = Field(gt=0, - le=gigabyte_to_mebibyte(Gigabytes(100)), #Limit to 1GiB - strict=True) + persistent_volume_size: int = Field( + gt=-1, le=gigabyte_to_mebibyte(Gigabytes(100)), strict=True # Limit to 1GiB + ) def __hash__(self): return hash(self.persistent_volume_size) @@ -77,11 +75,13 @@ def __hash__(self): class PersistentVolume(AbstractVolume): parent: Optional[ParentVolume] = None - persistence: VolumePersistence = None + persistence: Optional[VolumePersistence] = None name: Optional[str] = None - size_mib: PersistentVolumeSizeMib = 0 + size_mib: Optional[PersistentVolumeSizeMib] = PersistentVolumeSizeMib( + persistent_volume_size=0 + ) - @field_validator('size_mib', mode="before") + @field_validator("size_mib", mode="before") def convert_size_mib(cls, v): if isinstance(v, int): return PersistentVolumeSizeMib(persistent_volume_size=v) diff --git a/aleph_message/models/item_hash.py b/aleph_message/models/item_hash.py index 7037a99..e4074c5 100644 --- a/aleph_message/models/item_hash.py +++ b/aleph_message/models/item_hash.py @@ -2,9 +2,10 @@ from functools import lru_cache from ..exceptions import UnknownHashError -from pydantic_core import CoreSchema, core_schema +from pydantic_core import core_schema from pydantic import GetCoreSchemaHandler + class ItemType(str, Enum): """Item storage options""" @@ -46,14 +47,15 @@ def __new__(cls, value: str): return obj @classmethod - def __get_pydantic_core_schema__(cls, source, handler: GetCoreSchemaHandler) -> core_schema.CoreSchema: + 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() + cls.validate, core_schema.str_schema() ) @classmethod diff --git a/aleph_message/tests/test_types.py b/aleph_message/tests/test_types.py index f0a2dbf..1e8c03f 100644 --- a/aleph_message/tests/test_types.py +++ b/aleph_message/tests/test_types.py @@ -1,7 +1,7 @@ import copy import pytest -from pydantic import BaseModel, ValidationError, field_validator +from pydantic import BaseModel, ValidationError from aleph_message.exceptions import UnknownHashError from aleph_message.models import ItemHash, ItemType