diff --git a/zerver/management/commands/turn_off_digests.py b/zerver/management/commands/turn_off_digests.py index d87f1f008b..4f1688552b 100644 --- a/zerver/management/commands/turn_off_digests.py +++ b/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, - help='Turn off digests for this comma-separated ' - 'list of email addresses.') self.add_realm_args(parser) + self.add_user_list_args(parser, + help='Turn off digests for this comma-separated ' + 'list of email addresses.') + + 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: