From a1ba3ca066f29625b4f7e4dda686b655e001bfee Mon Sep 17 00:00:00 2001
From: Anders Kaseorg <anders@zulip.com>
Date: Mon, 15 Feb 2021 15:44:06 -0800
Subject: [PATCH] import_util: Strengthen get_users type using a Protocol.

Signed-off-by: Anders Kaseorg <anders@zulip.com>
---
 zerver/data_import/import_util.py | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

diff --git a/zerver/data_import/import_util.py b/zerver/data_import/import_util.py
index 3f51b61636..be5898f623 100644
--- a/zerver/data_import/import_util.py
+++ b/zerver/data_import/import_util.py
@@ -9,6 +9,7 @@ from typing import AbstractSet, Any, Callable, Dict, Iterable, List, Optional, S
 import orjson
 import requests
 from django.forms.models import model_to_dict
+from typing_extensions import Protocol
 
 from zerver.data_import.sequencer import NEXT_ID
 from zerver.lib.actions import STREAM_ASSIGNMENT_COLORS as stream_colors
@@ -223,8 +224,13 @@ def build_public_stream_subscriptions(
     return subscriptions
 
 
+class GetUsers(Protocol):
+    def __call__(self, stream_id: int = ..., huddle_id: int = ...) -> Set[int]:
+        ...
+
+
 def build_stream_subscriptions(
-    get_users: Callable[..., Set[int]],
+    get_users: GetUsers,
     zerver_recipient: List[ZerverFieldsT],
     zerver_stream: List[ZerverFieldsT],
 ) -> List[ZerverFieldsT]:
@@ -253,7 +259,7 @@ def build_stream_subscriptions(
 
 
 def build_huddle_subscriptions(
-    get_users: Callable[..., Set[int]],
+    get_users: GetUsers,
     zerver_recipient: List[ZerverFieldsT],
     zerver_huddle: List[ZerverFieldsT],
 ) -> List[ZerverFieldsT]: