From af0e2384952d1f9bd88e6b9ceaab6ca8e90ed035 Mon Sep 17 00:00:00 2001 From: dancergraham Date: Mon, 25 Nov 2024 23:02:00 +0100 Subject: [PATCH] refactor: for readability --- waffle/management/commands/waffle_flag.py | 116 +++++++++++----------- 1 file changed, 59 insertions(+), 57 deletions(-) diff --git a/waffle/management/commands/waffle_flag.py b/waffle/management/commands/waffle_flag.py index 9b393ac8..3fca539a 100644 --- a/waffle/management/commands/waffle_flag.py +++ b/waffle/management/commands/waffle_flag.py @@ -104,21 +104,7 @@ def handle(self, *args: Any, **options: Any) -> None: if options['list_flags']: self.stdout.write('Flags:') for flag in get_waffle_flag_model().objects.iterator(): - self.stdout.write(f'NAME: {flag.name}') - self.stdout.write(f'SUPERUSERS: {flag.superusers}') - self.stdout.write(f'EVERYONE: {flag.everyone}') - self.stdout.write(f'AUTHENTICATED: {flag.authenticated}') - self.stdout.write(f'PERCENT: {flag.percent}') - self.stdout.write(f'TESTING: {flag.testing}') - self.stdout.write(f'ROLLOUT: {flag.rollout}') - self.stdout.write(f'STAFF: {flag.staff}') - self.stdout.write('GROUPS: {}'.format(list( - flag.groups.values_list('name', flat=True))) - ) - self.stdout.write('USERS: {}'.format(list( - flag.users.values_list(UserModel.USERNAME_FIELD, flat=True))) - ) - self.stdout.write('') + self.log_flag_to_stdout(flag) return flag_name = options['name'] @@ -136,47 +122,63 @@ def handle(self, *args: Any, **options: Any) -> None: except get_waffle_flag_model().DoesNotExist: raise CommandError('This flag does not exist.') - # Loop through all options, setting Flag attributes that - # match (ie. don't want to try setting flag.verbosity) - for option in options: - # Group isn't an attribute on the Flag, but a related Many to Many - # field, so we handle it a bit differently by looking up groups and - # adding each group to the flag individually - if option == 'group': - group_hash = {} - for group in options['group']: - try: - group_instance = Group.objects.get(name=group) - group_hash[group_instance.name] = group_instance.id - except Group.DoesNotExist: - raise CommandError(f'Group {group} does not exist') - # If 'append' was not passed, we clear related groups - if not options['append']: - flag.groups.clear() - self.stdout.write('Setting group(s): %s' % ( - [name for name, _id in group_hash.items()]) - ) - for group_name, group_id in group_hash.items(): - flag.groups.add(group_id) - elif option == 'user': - user_hash = set() - for username in options['user']: - try: - user_instance = UserModel.objects.get( - Q(**{UserModel.USERNAME_FIELD: username}) - | Q(**{UserModel.EMAIL_FIELD: username}) - ) - user_hash.add(user_instance) - except UserModel.DoesNotExist: - raise CommandError(f'User {username} does not exist') - # If 'append' was not passed, we clear related users - if not options['append']: - flag.users.clear() - self.stdout.write(f'Setting user(s): {user_hash}') - # for user in user_hash: - flag.users.add(*[user.id for user in user_hash]) - elif hasattr(flag, option): - self.stdout.write(f'Setting {option}: {options[option]}') - setattr(flag, option, options[option]) + # Group isn't an attribute on the Flag, but a related Many-to-Many + # field, so we handle it a bit differently by looking up groups and + # adding each group to the flag individually + options_append = options.pop('append') + if groups := options.pop('group'): + group_hash = {} + for group in groups: + try: + group_instance = Group.objects.get(name=group) + group_hash[group_instance.name] = group_instance.id + except Group.DoesNotExist: + raise CommandError(f'Group {group} does not exist') + # If 'append' was not passed, we clear related groups + if not options_append: + flag.groups.clear() + self.stdout.write('Setting group(s): %s' % ( + [name for name, _id in group_hash.items()]) + ) + for group_id in group_hash.values(): + flag.groups.add(group_id) + if users := options.pop('user'): + user_hash = set() + for username in users: + try: + user_instance = UserModel.objects.get( + Q(**{UserModel.USERNAME_FIELD: username}) + | Q(**{UserModel.EMAIL_FIELD: username}) + ) + user_hash.add(user_instance) + except UserModel.DoesNotExist: + raise CommandError(f'User {username} does not exist') + # If 'append' was not passed, we clear related users + if not options_append: + flag.users.clear() + self.stdout.write(f'Setting user(s): {user_hash}') + # for user in user_hash: + flag.users.add(*[user.id for user in user_hash]) + for option_name, option in options.items(): + if hasattr(flag, option_name): + self.stdout.write(f'Setting {option_name}: {option}') + setattr(flag, option_name, option) flag.save() + + def log_flag_to_stdout(self, flag): + self.stdout.write(f'NAME: {flag.name}') + self.stdout.write(f'SUPERUSERS: {flag.superusers}') + self.stdout.write(f'EVERYONE: {flag.everyone}') + self.stdout.write(f'AUTHENTICATED: {flag.authenticated}') + self.stdout.write(f'PERCENT: {flag.percent}') + self.stdout.write(f'TESTING: {flag.testing}') + self.stdout.write(f'ROLLOUT: {flag.rollout}') + self.stdout.write(f'STAFF: {flag.staff}') + self.stdout.write('GROUPS: {}'.format(list( + flag.groups.values_list('name', flat=True))) + ) + self.stdout.write('USERS: {}'.format(list( + flag.users.values_list(UserModel.USERNAME_FIELD, flat=True))) + ) + self.stdout.write('')