From 9e766611a1046a9d7fe2fecf22fd61355a303d2e Mon Sep 17 00:00:00 2001
From: Sumner Evans <sumner@beeper.com>
Date: Tue, 12 Apr 2022 07:43:35 -0600
Subject: [PATCH] batch send: use classes without unnecessary/unknown fields

The batch send endpoint now accepts events and state events that don't
have room_id, event_id, etc. All batch send classes allow for the
specification of the timestamp.
---
 mautrix/appservice/api/intent.py | 10 +++++-----
 mautrix/types/__init__.py        |  4 ++++
 mautrix/types/event/__init__.py  |  1 +
 mautrix/types/event/batch.py     | 32 ++++++++++++++++++++++++++++++++
 mautrix/types/event/message.py   |  4 ++--
 5 files changed, 44 insertions(+), 7 deletions(-)
 create mode 100644 mautrix/types/event/batch.py

diff --git a/mautrix/appservice/api/intent.py b/mautrix/appservice/api/intent.py
index f9c991b7..be8d87ff 100644
--- a/mautrix/appservice/api/intent.py
+++ b/mautrix/appservice/api/intent.py
@@ -21,7 +21,9 @@
 from mautrix.types import (
     JSON,
     BatchID,
+    BatchSendEvent,
     BatchSendResponse,
+    BatchSendStateEvent,
     ContentURI,
     EventContent,
     EventID,
@@ -30,7 +32,6 @@
     JoinRulesStateEventContent,
     Member,
     Membership,
-    MessageEvent,
     PowerLevelStateEventContent,
     PresenceState,
     RoomAvatarStateEventContent,
@@ -38,7 +39,6 @@
     RoomNameStateEventContent,
     RoomPinnedEventsStateEventContent,
     RoomTopicStateEventContent,
-    StateEvent,
     StateEventContent,
     UserID,
 )
@@ -435,8 +435,8 @@ async def batch_send(
         prev_event_id: EventID,
         *,
         batch_id: BatchID | None = None,
-        events: Iterable[MessageEvent],
-        state_events_at_start: Iterable[StateEvent] = None,
+        events: Iterable[BatchSendEvent],
+        state_events_at_start: Iterable[BatchSendStateEvent] = (),
     ) -> BatchSendResponse:
         """
         Send a batch of historical events into a room. See `MSC2716`_ for more info.
@@ -459,7 +459,7 @@ async def batch_send(
             All the event IDs generated, plus a batch ID that can be passed back to this method.
         """
         path = Path.unstable["org.matrix.msc2716"].rooms[room_id].batch_send
-        query = {"prev_event_id": prev_event_id}
+        query: JSON = {"prev_event_id": prev_event_id}
         if batch_id:
             query["batch_id"] = batch_id
         resp = await self.api.request(
diff --git a/mautrix/types/__init__.py b/mautrix/types/__init__.py
index a53e965c..6040c525 100644
--- a/mautrix/types/__init__.py
+++ b/mautrix/types/__init__.py
@@ -34,6 +34,8 @@
     BaseMessageEventContentFuncs,
     BaseRoomEvent,
     BaseUnsigned,
+    BatchSendEvent,
+    BatchSendStateEvent,
     CallAnswerEventContent,
     CallCandidate,
     CallCandidatesEventContent,
@@ -214,6 +216,8 @@
     "BaseMessageEventContentFuncs",
     "BaseRoomEvent",
     "BaseUnsigned",
+    "BatchSendEvent",
+    "BatchSendStateEvent",
     "CallAnswerEventContent",
     "CallCandidate",
     "CallCandidatesEventContent",
diff --git a/mautrix/types/event/__init__.py b/mautrix/types/event/__init__.py
index a81358b0..2d7e889c 100644
--- a/mautrix/types/event/__init__.py
+++ b/mautrix/types/event/__init__.py
@@ -10,6 +10,7 @@
     RoomTagInfo,
 )
 from .base import BaseEvent, BaseRoomEvent, BaseUnsigned, GenericEvent
+from .batch import BatchSendEvent, BatchSendStateEvent
 from .encrypted import (
     EncryptedEvent,
     EncryptedEventContent,
diff --git a/mautrix/types/event/batch.py b/mautrix/types/event/batch.py
new file mode 100644
index 00000000..df45b742
--- /dev/null
+++ b/mautrix/types/event/batch.py
@@ -0,0 +1,32 @@
+# Copyright (c) 2022 Tulir Asokan, Sumner Evans
+#
+# This Source Code Form is subject to the terms of the Mozilla Public
+# License, v. 2.0. If a copy of the MPL was not distributed with this
+# file, You can obtain one at http://mozilla.org/MPL/2.0/.
+from typing import Any
+
+from attr import dataclass
+import attr
+
+from ..primitive import UserID
+from ..util import SerializableAttrs
+from .base import BaseEvent
+
+
+@dataclass
+class BatchSendEvent(BaseEvent, SerializableAttrs):
+    """Base event class for events sent via a batch send request."""
+
+    sender: UserID
+    timestamp: int = attr.ib(metadata={"json": "origin_server_ts"})
+    content: Any
+
+
+@dataclass
+class BatchSendStateEvent(BatchSendEvent, SerializableAttrs):
+    """
+    State events to be used as initial state events on batch send events. These never need to be
+    deserialized.
+    """
+
+    state_key: str
diff --git a/mautrix/types/event/message.py b/mautrix/types/event/message.py
index fd95a24a..45c34ce6 100644
--- a/mautrix/types/event/message.py
+++ b/mautrix/types/event/message.py
@@ -3,7 +3,7 @@
 # This Source Code Form is subject to the terms of the Mozilla Public
 # License, v. 2.0. If a copy of the MPL was not distributed with this
 # file, You can obtain one at http://mozilla.org/MPL/2.0/.
-from typing import Any, Dict, List, Optional, Pattern, Union
+from typing import Dict, List, Optional, Pattern, Union
 from html import escape
 import re
 
@@ -11,7 +11,7 @@
 import attr
 
 from ..primitive import JSON, ContentURI, EventID
-from ..util import ExtensibleEnum, Obj, Serializable, SerializableAttrs, deserializer, field
+from ..util import ExtensibleEnum, Obj, SerializableAttrs, deserializer, field
 from .base import BaseRoomEvent, BaseUnsigned
 
 # region Message types