diff --git a/aleph_message/models/__init__.py b/aleph_message/models/__init__.py index 9306012..8afec37 100644 --- a/aleph_message/models/__init__.py +++ b/aleph_message/models/__init__.py @@ -1,3 +1,4 @@ +import datetime import json from copy import copy from enum import Enum @@ -11,7 +12,7 @@ Literal, Optional, Type, - Union, + Union, TypeVar, cast, ) from pydantic import BaseModel, Extra, Field, validator @@ -203,7 +204,7 @@ class BaseMessage(BaseModel): size: Optional[int] = Field( default=None, description="Size of the content" ) # Almost always present - time: float = Field(description="Unix timestamp when the message was published") + time: datetime.datetime = Field(description="Unix timestamp or datetime when the message was published") item_type: ItemType = Field(description="Storage method used for the content") item_content: Optional[str] = Field( default=None, @@ -265,6 +266,13 @@ def check_confirmed(cls, v, values): raise ValueError("Message cannot be 'confirmed' without 'confirmations'") return v + @validator("time") + def convert_float_to_datetime(cls, v, values): + if isinstance(v, float): + v = datetime.datetime.fromtimestamp(v) + assert isinstance(v, datetime.datetime) + return v + class Config: extra = Extra.forbid exclude = {"id_", "_id"} @@ -334,14 +342,10 @@ class InstanceMessage(BaseMessage): ForgetMessage, ] -AlephMessageType: TypeAlias = Union[ - Type[PostMessage], - Type[AggregateMessage], - Type[StoreMessage], - Type[ProgramMessage], - Type[InstanceMessage], - Type[ForgetMessage], -] + +T = TypeVar('T', bound=AlephMessage) + +AlephMessageType: TypeAlias = Type[T] message_classes: List[AlephMessageType] = [ PostMessage, @@ -383,16 +387,16 @@ def add_item_content_and_hash(message_dict: Dict, inplace: bool = False): def create_new_message( message_dict: Dict, - factory: Optional[AlephMessageType] = None, -) -> AlephMessage: + factory: Optional[Type[T]] = None, +) -> T: """Create a new message from a dict. Computes the 'item_content' and 'item_hash' fields. """ message_content = add_item_content_and_hash(message_dict) if factory: - return factory.parse_obj(message_content) + return cast(T, factory.parse_obj(message_content)) else: - return parse_message(message_content) + return cast(T, parse_message(message_content)) def create_message_from_json( diff --git a/aleph_message/tests/test_models.py b/aleph_message/tests/test_models.py index ce170bf..48d587a 100644 --- a/aleph_message/tests/test_models.py +++ b/aleph_message/tests/test_models.py @@ -25,7 +25,7 @@ create_message_from_file, create_message_from_json, create_new_message, - parse_message, + parse_message, AlephMessage, ) from aleph_message.tests.download_messages import MESSAGES_STORAGE_PATH @@ -87,7 +87,7 @@ def test_messages_last_page(): if message_dict["item_hash"] in HASHES_TO_IGNORE: continue - message = parse_message(message_dict) + message: AlephMessage = parse_message(message_dict) assert message @@ -285,9 +285,14 @@ def test_create_new_message(): "signature": "0x123456789", # Signature validation requires using aleph-client } - new_message_1 = create_new_message(message_dict, factory=PostMessage) + new_message_1: PostMessage = create_new_message(message_dict, factory=PostMessage) assert new_message_1 assert new_message_1.type == MessageType.post + # Check that the time was converted to a datetime + assert new_message_1.time.isoformat() == '2021-07-07T10:04:47.017000+00:00' + + # The time field can be either a float or a datetime as string + message_dict["time"] = '2021-07-07T10:04:47.017000+00:00' new_message_2 = create_message_from_json( json.dumps(message_dict), factory=PostMessage ) @@ -303,7 +308,7 @@ def test_messages_from_disk(): data_dict = json.load(page_fd) for message_dict in data_dict["messages"]: try: - message = parse_message(message_dict) + message: AlephMessage = parse_message(message_dict) assert message except ValidationError as e: console.print("-" * 79) diff --git a/requirements.txt b/requirements.txt deleted file mode 100644 index d5a2c3a..0000000 --- a/requirements.txt +++ /dev/null @@ -1,18 +0,0 @@ -attrs==22.2.0 -certifi==2022.12.7 -charset-normalizer==3.0.1 -exceptiongroup==1.1.0 -idna==3.4 -iniconfig==2.0.0 -markdown-it-py==2.2.0 -mdurl==0.1.2 -packaging==23.0 -pluggy==1.0.0 -pydantic==1.10.5 -Pygments==2.14.0 -pytest==7.2.1 -requests==2.28.2 -rich==13.3.1 -tomli==2.0.1 -typing-extensions==4.5.0 -urllib3==1.26.14