Fix typing issues with mypy by using a generic type for AlephMessage

aleph_message/models/__init__.py

@@ -12,7 +12,7 @@
     Literal,
     Optional,
     Type,
-    Union,
+    Union, TypeVar, cast,
 )
 
 from pydantic import BaseModel, Extra, Field, validator
@@ -342,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,
@@ -391,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(

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
 
 
@@ -308,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)