diff --git a/synapse/streams/__init__.py b/synapse/streams/__init__.py
index 5e83dba2ed..806b671305 100644
--- a/synapse/streams/__init__.py
+++ b/synapse/streams/__init__.py
@@ -11,3 +11,25 @@
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
+
+from typing import Collection, Generic, List, Optional, Tuple, TypeVar
+
+from synapse.types import UserID
+
+# The key, this is either a stream token or int.
+K = TypeVar("K")
+# The return type.
+R = TypeVar("R")
+
+
+class EventSource(Generic[K, R]):
+ async def get_new_events(
+ self,
+ user: UserID,
+ from_key: K,
+ limit: Optional[int],
+ room_ids: Collection[str],
+ is_guest: bool,
+ explicit_room_id: Optional[str] = None,
+ ) -> Tuple[List[R], K]:
+ ...
diff --git a/synapse/streams/events.py b/synapse/streams/events.py
index 99b0aac2fb..21591d0bfd 100644
--- a/synapse/streams/events.py
+++ b/synapse/streams/events.py
@@ -12,29 +12,40 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from typing import Any, Dict
+from typing import TYPE_CHECKING, Iterator, Tuple
+
+import attr
from synapse.handlers.account_data import AccountDataEventSource
from synapse.handlers.presence import PresenceEventSource
from synapse.handlers.receipts import ReceiptEventSource
from synapse.handlers.room import RoomEventSource
from synapse.handlers.typing import TypingNotificationEventSource
+from synapse.streams import EventSource
from synapse.types import StreamToken
+if TYPE_CHECKING:
+ from synapse.server import HomeServer
-class EventSources:
- SOURCE_TYPES = {
- "room": RoomEventSource,
- "presence": PresenceEventSource,
- "typing": TypingNotificationEventSource,
- "receipt": ReceiptEventSource,
- "account_data": AccountDataEventSource,
- }
- def __init__(self, hs):
- self.sources: Dict[str, Any] = {
- name: cls(hs) for name, cls in EventSources.SOURCE_TYPES.items()
- }
+@attr.s(frozen=True, slots=True, auto_attribs=True)
+class _EventSourcesInner:
+ room: RoomEventSource
+ presence: PresenceEventSource
+ typing: TypingNotificationEventSource
+ receipt: ReceiptEventSource
+ account_data: AccountDataEventSource
+
+ def get_sources(self) -> Iterator[Tuple[str, EventSource]]:
+ for attribute in _EventSourcesInner.__attrs_attrs__: # type: ignore[attr-defined]
+ yield attribute.name, getattr(self, attribute.name)
+
+
+class EventSources:
+ def __init__(self, hs: "HomeServer"):
+ self.sources = _EventSourcesInner(
+ *(attribute.type(hs) for attribute in _EventSourcesInner.__attrs_attrs__) # type: ignore[attr-defined]
+ )
self.store = hs.get_datastore()
def get_current_token(self) -> StreamToken:
@@ -44,11 +55,11 @@ class EventSources:
groups_key = self.store.get_group_stream_token()
token = StreamToken(
- room_key=self.sources["room"].get_current_key(),
- presence_key=self.sources["presence"].get_current_key(),
- typing_key=self.sources["typing"].get_current_key(),
- receipt_key=self.sources["receipt"].get_current_key(),
- account_data_key=self.sources["account_data"].get_current_key(),
+ room_key=self.sources.room.get_current_key(),
+ presence_key=self.sources.presence.get_current_key(),
+ typing_key=self.sources.typing.get_current_key(),
+ receipt_key=self.sources.receipt.get_current_key(),
+ account_data_key=self.sources.account_data.get_current_key(),
push_rules_key=push_rules_key,
to_device_key=to_device_key,
device_list_key=device_list_key,
@@ -67,7 +78,7 @@ class EventSources:
The current token for pagination.
"""
token = StreamToken(
- room_key=self.sources["room"].get_current_key(),
+ room_key=self.sources.room.get_current_key(),
presence_key=0,
typing_key=0,
receipt_key=0,
|
