Use UpdateFunction more

synapse/replication/tcp/streams/_base.py

@@ -429,9 +429,9 @@ class TypingStream(Stream):
         if hs.get_instance_name() in hs.config.worker.writers.typing:
             # On the writer, query the typing handler
             typing_writer_handler = hs.get_typing_writer_handler()
-            update_function: Callable[
-                [str, int, int, int], Awaitable[tuple[list[tuple[int, Any]], int, bool]]
-            ] = typing_writer_handler.get_all_typing_updates
+            update_function: UpdateFunction = (
+                typing_writer_handler.get_all_typing_updates
+            )
             self.current_token_function = typing_writer_handler.get_current_token
         else:
             # Query the typing writer process

synapse/replication/tcp/streams/federation.py

@@ -18,13 +18,14 @@
 # [This file includes modifications made by New Vector Limited]
 #
 #
-from typing import TYPE_CHECKING, Any, Awaitable, Callable
+from typing import TYPE_CHECKING
 
 import attr
 
 from synapse.replication.tcp.streams._base import (
     Stream,
     Token,
+    UpdateFunction,
     current_token_without_instance,
     make_http_update_function,
 )
@@ -57,9 +58,7 @@ class FederationStream(Stream):
             self.current_token_func = current_token_without_instance(
                 federation_sender.get_current_token
             )
-            update_function: Callable[
-                [str, int, int, int], Awaitable[tuple[list[tuple[int, Any]], int, bool]]
-            ] = federation_sender.get_replication_rows
+            update_function: UpdateFunction = federation_sender.get_replication_rows
 
         elif hs.should_send_federation():
             # federation sender: Query master process