management: Use add_user_list_args in turn_off_digests.

zerver/management/commands/turn_off_digests.py

@@ -14,27 +14,33 @@ class Command(ZulipBaseCommand):
 
     def add_arguments(self, parser):
         # type: (CommandParser) -> None
-        parser.add_argument('-u', '--users',
-                            dest='users',
-                            type=str,
+        self.add_realm_args(parser)
+
+        self.add_user_list_args(parser,
                                 help='Turn off digests for this comma-separated '
                                      'list of email addresses.')
-        self.add_realm_args(parser)
+
+        parser.add_argument('-a', '--all-users',
+                            dest='all_users',
+                            action="store_true",
+                            default=False,
+                            help="Turn off digests for everyone in a realm. "
+                                 "Don't forget to specify the realm.")
 
     def handle(self, **options):
         # type: (**str) -> None
         realm = self.get_realm(options)
-        if realm is None and options["users"] is None:
+        user_profiles = self.get_users(options, realm)
+        all_users = options["all_users"]
+
+        # If all_users flag is passed user list should not be passed and vice versa.
+        # If all_users flag is passed it is manadatory to pass the realm.
+        if (bool(user_profiles) == all_users) or (all_users and not realm):
             self.print_help("./manage.py", "turn_off_digests")
             exit(1)
 
-        if realm and not options["users"]:
+        if all_users:
             user_profiles = UserProfile.objects.filter(realm=realm)
-        else:
-            emails = set([email.strip() for email in options["users"].split(",")])
-            user_profiles = []
-            for email in emails:
-                user_profiles.append(self.get_user(email, realm))
 
         print("Turned off digest emails for:")
         for user_profile in user_profiles: