feat: Add account filtering in aws configure sso

Signed-off-by: frauniki <[email protected]>

Author: frauniki

awscli/customizations/configure/sso.py

@@ -439,10 +439,12 @@ def _handle_single_account(self, accounts):
     def _handle_multiple_accounts(self, accounts):
         available_accounts_msg = (
             'There are {} AWS accounts available to you.\n'
+            'Use arrow keys to navigate, type to filter, and press Enter to select.\n'
         )
         uni_print(available_accounts_msg.format(len(accounts)))
         selected_account = self._selector(
-            accounts, display_format=display_account
+            accounts, display_format=display_account, enable_filter=True,
+            no_results_message='No matching accounts found'
         )
         sso_account_id = selected_account['accountId']
         return sso_account_id

awscli/customizations/wizard/ui/selectmenu.py

@@ -23,7 +23,9 @@
 from prompt_toolkit.utils import get_cwidth
 
 
-def select_menu(items, display_format=None, max_height=10):
+def select_menu(
+    items, display_format=None, max_height=10, enable_filter=False
+):
     """Presents a list of options and allows the user to select one.
 
     This presents a static list of options and prompts the user to select one.
@@ -42,6 +44,9 @@ def select_menu(items, display_format=None, max_height=10):
     :type max_height: int
     :param max_height: The max number of items to show in the list at a time.
 
+    :type enable_filter: bool
+    :param enable_filter: Enable keyboard filtering of items.
+
     :returns: The selected element from the items list.
     """
     app_bindings = KeyBindings()
@@ -51,8 +56,19 @@ def exit_app(event):
         event.app.exit(exception=KeyboardInterrupt, style='class:aborting')
 
     min_height = min(max_height, len(items))
+    if enable_filter:
+        # Add 1 to height for filter line
+        min_height = min(max_height + 1, len(items) + 1)
+        menu_control = FilterableSelectionMenuControl(
+            items, display_format=display_format
+        )
+    else:
+        menu_control = SelectionMenuControl(
+            items, display_format=display_format
+        )
+
     menu_window = Window(
-        SelectionMenuControl(items, display_format=display_format),
+        menu_control,
         always_hide_cursor=False,
         height=Dimension(min=min_height, max=min_height),
         scroll_offsets=ScrollOffsets(),
@@ -122,6 +138,8 @@ def is_focusable(self):
 
     def preferred_width(self, max_width):
         items = self._get_items()
+        if not items:
+            return self.MIN_WIDTH
         if self._display_format:
             items = (self._display_format(i) for i in items)
         max_item_width = max(get_cwidth(i) for i in items)
@@ -188,6 +206,157 @@ def app_result(event):
         return kb
 
 
+class FilterableSelectionMenuControl(SelectionMenuControl):
+    """Menu that supports keyboard filtering of items"""
+
+    def __init__(self, items, display_format=None, cursor='>', no_results_message=None):
+        super().__init__(items, display_format=display_format, cursor=cursor)
+        self._filter_text = ''
+        self._filtered_items = items if items else []
+        self._all_items = items if items else []
+        self._filter_enabled = True
+        self._no_results_message = no_results_message or 'No matching items found'
+
+    def _get_items(self):
+        if callable(self._all_items):
+            self._all_items = self._all_items()
+        return self._filtered_items
+
+    def preferred_width(self, max_width):
+        # Ensure minimum width for search display
+        min_search_width = max(20, len("Search: " + self._filter_text) + 5)
+
+        # Get width from filtered items
+        items = self._filtered_items
+        if not items:
+            # Width for no results message
+            no_results_width = get_cwidth(self._no_results_message) + 4
+            return max(no_results_width, min_search_width)
+
+        if self._display_format:
+            items_display = [self._display_format(i) for i in items]
+        else:
+            items_display = [str(i) for i in items]
+
+        if items_display:
+            max_item_width = max(get_cwidth(i) for i in items_display)
+            max_item_width += self._format_overhead
+        else:
+            max_item_width = self.MIN_WIDTH
+
+        max_item_width = max(max_item_width, min_search_width)
+
+        if max_item_width < self.MIN_WIDTH:
+            max_item_width = self.MIN_WIDTH
+        return min(max_width, max_item_width)
+
+    def _update_filtered_items(self):
+        """Update the filtered items based on the current filter text"""
+        if not self._filter_text:
+            self._filtered_items = self._all_items
+        else:
+            filter_lower = self._filter_text.lower()
+            self._filtered_items = [
+                item
+                for item in self._all_items
+                if filter_lower
+                in (
+                    self._display_format(item)
+                    if self._display_format
+                    else str(item)
+                ).lower()
+            ]
+
+        # Reset selection if it's out of bounds
+        if self._selection >= len(self._filtered_items):
+            self._selection = 0
+
+    def preferred_height(self, width, max_height, wrap_lines, get_line_prefix):
+        # Add 1 extra line for the filter display
+        return min(max_height, len(self._get_items()) + 1)
+
+    def create_content(self, width, height):
+        def get_line(i):
+            # First line shows the filter
+            if i == 0:
+                filter_display = (
+                    f"Search: {self._filter_text}_"
+                    if self._filter_enabled
+                    else f"Search: {self._filter_text}"
+                )
+                return [('class:filter', filter_display)]
+
+            # Show "No results" message if filtered items is empty
+            if not self._filtered_items:
+                if i == 1:
+                    return [
+                        ('class:no-results', f'  {self._no_results_message}')
+                    ]
+                return [('', '')]
+
+            # Adjust for the filter line
+            item_index = i - 1
+            if item_index >= len(self._filtered_items):
+                return [('', '')]
+
+            item = self._filtered_items[item_index]
+            is_selected = item_index == self._selection
+            return self._menu_item_fragment(item, is_selected, width)
+
+        # Ensure at least 2 lines (search + no results or items)
+        line_count = max(2, len(self._filtered_items) + 1)
+        cursor_y = self._selection + 1 if self._filtered_items else 0
+
+        return UIContent(
+            get_line=get_line,
+            cursor_position=Point(x=0, y=cursor_y),
+            line_count=line_count,
+        )
+
+    def get_key_bindings(self):
+        kb = KeyBindings()
+
+        @kb.add('up')
+        def move_up(event):
+            if len(self._filtered_items) > 0:
+                self._move_cursor(-1)
+
+        @kb.add('down')
+        def move_down(event):
+            if len(self._filtered_items) > 0:
+                self._move_cursor(1)
+
+        @kb.add('enter')
+        def app_result(event):
+            if len(self._filtered_items) > 0:
+                result = self._filtered_items[self._selection]
+                event.app.exit(result=result)
+
+        @kb.add('backspace')
+        def delete_char(event):
+            if self._filter_text:
+                self._filter_text = self._filter_text[:-1]
+                self._update_filtered_items()
+
+        @kb.add('c-u')
+        def clear_filter(event):
+            self._filter_text = ''
+            self._update_filtered_items()
+
+        # Add support for typing any character
+        from string import printable
+
+        for char in printable:
+            if char not in ('\n', '\r', '\t'):
+
+                @kb.add(char)
+                def add_char(event, c=char):
+                    self._filter_text += c
+                    self._update_filtered_items()
+
+        return kb
+
+
 class CollapsableSelectionMenuControl(SelectionMenuControl):
     """Menu that collapses to text with selection when loses focus"""
 

tests/unit/customizations/configure/test_sso.py

@@ -747,6 +747,8 @@ class PTKStubber:
     _ALLOWED_SELECT_MENU_KWARGS = {
         "display_format",
         "max_height",
+        "enable_filter",
+        "no_results_message",
     }
 
     def __init__(self, user_inputs=None):
@@ -763,9 +765,9 @@ def prompt(self, message, **kwargs):
                 f'Received prompt with no stubbed answer: "{message}"'
             )
         prompt = self._expected_inputs.pop(0)
-        assert isinstance(
-            prompt, Prompt
-        ), f'Did not receive user input of type Prompt for: "{message}"'
+        assert isinstance(prompt, Prompt), (
+            f'Did not receive user input of type Prompt for: "{message}"'
+        )
         if prompt.expected_message is not None:
             assert message == prompt.expected_message, (
                 f"Prompt does not match expected "
@@ -792,9 +794,9 @@ def select_menu(self, items, **kwargs):
                 f'Received select_menu with no stubbed answer: "{items}"'
             )
         select_menu = self._expected_inputs.pop(0)
-        assert isinstance(
-            select_menu, SelectMenu
-        ), f'Did not receive user input of type SelectMenu for: "{items}"'
+        assert isinstance(select_menu, SelectMenu), (
+            f'Did not receive user input of type SelectMenu for: "{items}"'
+        )
         if select_menu.expected_choices is not None:
             assert items == select_menu.expected_choices, (
                 f"Choices does not match expected select_menu choices "
@@ -807,9 +809,9 @@ def _initialize_expected_inputs_if_needed(self):
             self._expected_inputs = self.user_inputs.get_expected_inputs()
 
     def _validate_kwargs(self, provided_kwargs, allowed_kwargs):
-        assert set(provided_kwargs).issubset(
-            allowed_kwargs
-        ), "Arguments provided does not matched allowed keyword arguments"
+        assert set(provided_kwargs).issubset(allowed_kwargs), (
+            "Arguments provided does not matched allowed keyword arguments"
+        )
 
 
 def write_aws_config(aws_config, lines):
@@ -1484,6 +1486,66 @@ def test_configure_sso_suggests_values_from_sessions(
         sso_cmd = sso_cmd_factory(session=session)
         assert sso_cmd(args, parsed_globals) == 0
 
+    def test_multiple_accounts_uses_filterable_menu(
+        self,
+        sso_cmd,
+        ptk_stubber,
+        aws_config,
+        stub_sso_list_roles,
+        stub_sso_list_accounts,
+        mock_do_sso_login,
+        botocore_session,
+        args,
+        parsed_globals,
+        configure_sso_legacy_inputs,
+        capsys,
+    ):
+        """Test that multiple accounts selection shows filter instructions"""
+        inputs = configure_sso_legacy_inputs
+        selected_account_id = inputs.account_id_select.answer["accountId"]
+        ptk_stubber.user_inputs = inputs
+
+        stub_sso_list_accounts(inputs.account_id_select.expected_choices)
+        stub_sso_list_roles(
+            inputs.role_name_select.expected_choices,
+            expected_account_id=selected_account_id,
+        )
+
+        sso_cmd(args, parsed_globals)
+
+        # Verify the filter instructions are shown for multiple accounts
+        stdout = capsys.readouterr().out
+        assert "Use arrow keys to navigate, type to filter" in stdout
+        assert "There are 2 AWS accounts available to you" in stdout
+
+    def test_single_account_does_not_use_filterable_menu(
+        self,
+        sso_cmd,
+        ptk_stubber,
+        aws_config,
+        stub_simple_single_item_sso_responses,
+        mock_do_sso_login,
+        botocore_session,
+        args,
+        parsed_globals,
+        configure_sso_legacy_inputs,
+        account_id,
+        role_name,
+        capsys,
+    ):
+        """Test that single account does not show filter instructions"""
+        inputs = configure_sso_legacy_inputs
+        inputs.skip_account_and_role_selection()
+        ptk_stubber.user_inputs = inputs
+        stub_simple_single_item_sso_responses(account_id, role_name)
+
+        sso_cmd(args, parsed_globals)
+
+        # Verify the filter instructions are NOT shown for single account
+        stdout = capsys.readouterr().out
+        assert "Use arrow keys to navigate, type to filter" not in stdout
+        assert "The only AWS account available to you is" in stdout
+
     def test_single_account_single_role_device_code_fallback(
         self,
         sso_cmd,