refactor: Extract get_stream_subscriptions_for_user().

zerver/lib/actions.py

@@ -37,6 +37,7 @@ from zerver.lib.send_email import send_email, FromAddress
 from zerver.lib.stream_subscription import (
     get_active_subscriptions_for_stream_id,
     get_active_subscriptions_for_stream_ids,
+    get_stream_subscriptions_for_user,
     num_subscribers_for_stream_id,
 )
 from zerver.lib.stream_topic import StreamTopicTarget
@@ -1876,9 +1877,7 @@ def internal_send_private_message(realm, sender, recipient_user, content):
 
 def pick_color(user_profile):
     # type: (UserProfile) -> Text
-    subs = Subscription.objects.filter(user_profile=user_profile,
+    subs = get_stream_subscriptions_for_user(user_profile).filter(active=True)
-                                       active=True,
-                                       recipient__type=Recipient.STREAM)
     return pick_color_helper(user_profile, subs)
 
 def pick_color_helper(user_profile, subs):
@@ -3486,10 +3485,7 @@ def decode_email_address(email):
 # subscriptions, so it's worth optimizing.
 def gather_subscriptions_helper(user_profile, include_subscribers=True):
     # type: (UserProfile, bool) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]], List[Dict[str, Any]]]
-    sub_dicts = Subscription.objects.filter(
+    sub_dicts = get_stream_subscriptions_for_user(user_profile).values(
-        user_profile    = user_profile,
-        recipient__type = Recipient.STREAM
-    ).values(
         "recipient_id", "in_home_view", "color", "desktop_notifications",
         "audible_notifications", "push_notifications", "active", "pin_to_top"
     ).order_by("recipient_id")
@@ -4021,9 +4017,9 @@ def do_get_streams(user_profile, include_public=True, include_subscribed=True,
     query = get_occupied_streams(user_profile.realm)
 
     if not include_all_active:
-        user_subs = Subscription.objects.select_related("recipient").filter(
+        user_subs = get_stream_subscriptions_for_user(user_profile).filter(
-            active=True, user_profile=user_profile,
+            active=True,
-            recipient__type=Recipient.STREAM)
+        ).select_related('recipient')
 
         if include_subscribed:
             recipient_check = Q(id__in=[sub.recipient.type_id for sub in user_subs])

zerver/lib/message.py

@@ -12,6 +12,9 @@ import zerver.lib.bugdown as bugdown
 from zerver.lib.cache import cache_with_key, to_dict_cache_key
 from zerver.lib.request import JsonableError
 from zerver.lib.str_utils import force_bytes, dict_with_str_keys
+from zerver.lib.stream_subscription import (
+    get_stream_subscriptions_for_user,
+)
 from zerver.lib.timestamp import datetime_to_timestamp
 from zerver.lib.topic_mutes import (
     build_topic_mute_checker,
@@ -593,9 +596,7 @@ def aggregate_message_dict(input_dict, lookup_fields, collect_senders):
 
 def get_inactive_recipient_ids(user_profile):
     # type: (UserProfile) -> List[int]
-    rows = Subscription.objects.filter(
+    rows = get_stream_subscriptions_for_user(user_profile).filter(
-        user_profile=user_profile,
-        recipient__type=Recipient.STREAM,
         active=False,
     ).values(
         'recipient_id'
@@ -607,9 +608,7 @@ def get_inactive_recipient_ids(user_profile):
 
 def get_muted_stream_ids(user_profile):
     # type: (UserProfile) -> List[int]
-    rows = Subscription.objects.filter(
+    rows = get_stream_subscriptions_for_user(user_profile).filter(
-        user_profile=user_profile,
-        recipient__type=Recipient.STREAM,
         active=True,
         in_home_view=False,
     ).values(

zerver/lib/stream_subscription.py

@@ -4,6 +4,7 @@ from django.db.models.query import QuerySet
 from zerver.models import (
     Recipient,
     Subscription,
+    UserProfile,
 )
 
 def get_active_subscriptions_for_stream_id(stream_id):
@@ -22,6 +23,13 @@ def get_active_subscriptions_for_stream_ids(stream_ids):
         active=True
     )
 
+def get_stream_subscriptions_for_user(user_profile):
+    # type: (UserProfile) -> QuerySet
+    return Subscription.objects.filter(
+        user_profile=user_profile,
+        recipient__type=Recipient.STREAM,
+    )
+
 def num_subscribers_for_stream_id(stream_id):
     # type: (int) -> int
     return get_active_subscriptions_for_stream_id(stream_id).filter(

zerver/lib/test_classes.py

@@ -28,6 +28,10 @@ from zerver.lib.actions import (
     check_send_stream_message,
 )
 
+from zerver.lib.stream_subscription import (
+    get_stream_subscriptions_for_user,
+)
+
 from zerver.lib.test_helpers import (
     instrument_url, find_key_by_email,
 )
@@ -375,10 +379,9 @@ class ZulipTestCase(TestCase):
         Helper function to get the stream names for a user
         """
         user_profile = get_user(email, realm)
-        subs = Subscription.objects.filter(
+        subs = get_stream_subscriptions_for_user(user_profile).filter(
-            user_profile=user_profile,
             active=True,
-            recipient__type=Recipient.STREAM)
+        )
         return [cast(Text, get_display_recipient(sub.recipient)) for sub in subs]
 
     def send_personal_message(self, from_email, to_email, content=u"test content"):