1from __future__ import annotations 
    2 
    3from typing import Callable, Iterable, Sequence 
    4 
    5from prompt_toolkit.application.current import get_app 
    6from prompt_toolkit.filters import Condition 
    7from prompt_toolkit.formatted_text.base import OneStyleAndTextTuple, StyleAndTextTuples 
    8from prompt_toolkit.key_binding.key_bindings import KeyBindings, KeyBindingsBase 
    9from prompt_toolkit.key_binding.key_processor import KeyPressEvent 
    10from prompt_toolkit.keys import Keys 
    11from prompt_toolkit.layout.containers import ( 
    12    AnyContainer, 
    13    ConditionalContainer, 
    14    Container, 
    15    Float, 
    16    FloatContainer, 
    17    HSplit, 
    18    Window, 
    19) 
    20from prompt_toolkit.layout.controls import FormattedTextControl 
    21from prompt_toolkit.mouse_events import MouseEvent, MouseEventType 
    22from prompt_toolkit.utils import get_cwidth 
    23from prompt_toolkit.widgets import Shadow 
    24 
    25from .base import Border 
    26 
    27__all__ = [ 
    28    "MenuContainer", 
    29    "MenuItem", 
    30] 
    31 
    32E = KeyPressEvent 
    33 
    34 
    35class MenuContainer: 
    36    """ 
    37    :param floats: List of extra Float objects to display. 
    38    :param menu_items: List of `MenuItem` objects. 
    39    """ 
    40 
    41    def __init__( 
    42        self, 
    43        body: AnyContainer, 
    44        menu_items: list[MenuItem], 
    45        floats: list[Float] | None = None, 
    46        key_bindings: KeyBindingsBase | None = None, 
    47    ) -> None: 
    48        self.body = body 
    49        self.menu_items = menu_items 
    50        self.selected_menu = [0] 
    51 
    52        # Key bindings. 
    53        kb = KeyBindings() 
    54 
    55        @Condition 
    56        def in_main_menu() -> bool: 
    57            return len(self.selected_menu) == 1 
    58 
    59        @Condition 
    60        def in_sub_menu() -> bool: 
    61            return len(self.selected_menu) > 1 
    62 
    63        # Navigation through the main menu. 
    64 
    65        @kb.add("left", filter=in_main_menu) 
    66        def _left(event: E) -> None: 
    67            self.selected_menu[0] = max(0, self.selected_menu[0] - 1) 
    68 
    69        @kb.add("right", filter=in_main_menu) 
    70        def _right(event: E) -> None: 
    71            self.selected_menu[0] = min( 
    72                len(self.menu_items) - 1, self.selected_menu[0] + 1 
    73            ) 
    74 
    75        @kb.add("down", filter=in_main_menu) 
    76        def _down(event: E) -> None: 
    77            self.selected_menu.append(0) 
    78 
    79        @kb.add("c-c", filter=in_main_menu) 
    80        @kb.add("c-g", filter=in_main_menu) 
    81        def _cancel(event: E) -> None: 
    82            "Leave menu." 
    83            event.app.layout.focus_last() 
    84 
    85        # Sub menu navigation. 
    86 
    87        @kb.add("left", filter=in_sub_menu) 
    88        @kb.add("c-g", filter=in_sub_menu) 
    89        @kb.add("c-c", filter=in_sub_menu) 
    90        def _back(event: E) -> None: 
    91            "Go back to parent menu." 
    92            if len(self.selected_menu) > 1: 
    93                self.selected_menu.pop() 
    94 
    95        @kb.add("right", filter=in_sub_menu) 
    96        def _submenu(event: E) -> None: 
    97            "go into sub menu." 
    98            if self._get_menu(len(self.selected_menu) - 1).children: 
    99                self.selected_menu.append(0) 
    100 
    101            # If This item does not have a sub menu. Go up in the parent menu. 
    102            elif ( 
    103                len(self.selected_menu) == 2 
    104                and self.selected_menu[0] < len(self.menu_items) - 1 
    105            ): 
    106                self.selected_menu = [ 
    107                    min(len(self.menu_items) - 1, self.selected_menu[0] + 1) 
    108                ] 
    109                if self.menu_items[self.selected_menu[0]].children: 
    110                    self.selected_menu.append(0) 
    111 
    112        @kb.add("up", filter=in_sub_menu) 
    113        def _up_in_submenu(event: E) -> None: 
    114            "Select previous (enabled) menu item or return to main menu." 
    115            # Look for previous enabled items in this sub menu. 
    116            menu = self._get_menu(len(self.selected_menu) - 2) 
    117            index = self.selected_menu[-1] 
    118 
    119            previous_indexes = [ 
    120                i 
    121                for i, item in enumerate(menu.children) 
    122                if i < index and not item.disabled 
    123            ] 
    124 
    125            if previous_indexes: 
    126                self.selected_menu[-1] = previous_indexes[-1] 
    127            elif len(self.selected_menu) == 2: 
    128                # Return to main menu. 
    129                self.selected_menu.pop() 
    130 
    131        @kb.add("down", filter=in_sub_menu) 
    132        def _down_in_submenu(event: E) -> None: 
    133            "Select next (enabled) menu item." 
    134            menu = self._get_menu(len(self.selected_menu) - 2) 
    135            index = self.selected_menu[-1] 
    136 
    137            next_indexes = [ 
    138                i 
    139                for i, item in enumerate(menu.children) 
    140                if i > index and not item.disabled 
    141            ] 
    142 
    143            if next_indexes: 
    144                self.selected_menu[-1] = next_indexes[0] 
    145 
    146        @kb.add("enter") 
    147        def _click(event: E) -> None: 
    148            "Click the selected menu item." 
    149            item = self._get_menu(len(self.selected_menu) - 1) 
    150            if item.handler: 
    151                event.app.layout.focus_last() 
    152                item.handler() 
    153 
    154        # Controls. 
    155        self.control = FormattedTextControl( 
    156            self._get_menu_fragments, key_bindings=kb, focusable=True, show_cursor=False 
    157        ) 
    158 
    159        self.window = Window(height=1, content=self.control, style="class:menu-bar") 
    160 
    161        submenu = self._submenu(0) 
    162        submenu2 = self._submenu(1) 
    163        submenu3 = self._submenu(2) 
    164 
    165        @Condition 
    166        def has_focus() -> bool: 
    167            return get_app().layout.current_window == self.window 
    168 
    169        self.container = FloatContainer( 
    170            content=HSplit( 
    171                [ 
    172                    # The titlebar. 
    173                    self.window, 
    174                    # The 'body', like defined above. 
    175                    body, 
    176                ] 
    177            ), 
    178            floats=[ 
    179                Float( 
    180                    xcursor=True, 
    181                    ycursor=True, 
    182                    content=ConditionalContainer( 
    183                        content=Shadow(body=submenu), filter=has_focus 
    184                    ), 
    185                ), 
    186                Float( 
    187                    attach_to_window=submenu, 
    188                    xcursor=True, 
    189                    ycursor=True, 
    190                    allow_cover_cursor=True, 
    191                    content=ConditionalContainer( 
    192                        content=Shadow(body=submenu2), 
    193                        filter=has_focus 
    194                        & Condition(lambda: len(self.selected_menu) >= 1), 
    195                    ), 
    196                ), 
    197                Float( 
    198                    attach_to_window=submenu2, 
    199                    xcursor=True, 
    200                    ycursor=True, 
    201                    allow_cover_cursor=True, 
    202                    content=ConditionalContainer( 
    203                        content=Shadow(body=submenu3), 
    204                        filter=has_focus 
    205                        & Condition(lambda: len(self.selected_menu) >= 2), 
    206                    ), 
    207                ), 
    208                # -- 
    209            ] 
    210            + (floats or []), 
    211            key_bindings=key_bindings, 
    212        ) 
    213 
    214    def _get_menu(self, level: int) -> MenuItem: 
    215        menu = self.menu_items[self.selected_menu[0]] 
    216 
    217        for i, index in enumerate(self.selected_menu[1:]): 
    218            if i < level: 
    219                try: 
    220                    menu = menu.children[index] 
    221                except IndexError: 
    222                    return MenuItem("debug") 
    223 
    224        return menu 
    225 
    226    def _get_menu_fragments(self) -> StyleAndTextTuples: 
    227        focused = get_app().layout.has_focus(self.window) 
    228 
    229        # This is called during the rendering. When we discover that this 
    230        # widget doesn't have the focus anymore. Reset menu state. 
    231        if not focused: 
    232            self.selected_menu = [0] 
    233 
    234        # Generate text fragments for the main menu. 
    235        def one_item(i: int, item: MenuItem) -> Iterable[OneStyleAndTextTuple]: 
    236            def mouse_handler(mouse_event: MouseEvent) -> None: 
    237                hover = mouse_event.event_type == MouseEventType.MOUSE_MOVE 
    238                if ( 
    239                    mouse_event.event_type == MouseEventType.MOUSE_DOWN 
    240                    or hover 
    241                    and focused 
    242                ): 
    243                    # Toggle focus. 
    244                    app = get_app() 
    245                    if not hover: 
    246                        if app.layout.has_focus(self.window): 
    247                            if self.selected_menu == [i]: 
    248                                app.layout.focus_last() 
    249                        else: 
    250                            app.layout.focus(self.window) 
    251                    self.selected_menu = [i] 
    252 
    253            yield ("class:menu-bar", " ", mouse_handler) 
    254            if i == self.selected_menu[0] and focused: 
    255                yield ("[SetMenuPosition]", "", mouse_handler) 
    256                style = "class:menu-bar.selected-item" 
    257            else: 
    258                style = "class:menu-bar" 
    259            yield style, item.text, mouse_handler 
    260 
    261        result: StyleAndTextTuples = [] 
    262        for i, item in enumerate(self.menu_items): 
    263            result.extend(one_item(i, item)) 
    264 
    265        return result 
    266 
    267    def _submenu(self, level: int = 0) -> Window: 
    268        def get_text_fragments() -> StyleAndTextTuples: 
    269            result: StyleAndTextTuples = [] 
    270            if level < len(self.selected_menu): 
    271                menu = self._get_menu(level) 
    272                if menu.children: 
    273                    result.append(("class:menu", Border.TOP_LEFT)) 
    274                    result.append(("class:menu", Border.HORIZONTAL * (menu.width + 4))) 
    275                    result.append(("class:menu", Border.TOP_RIGHT)) 
    276                    result.append(("", "\n")) 
    277                    try: 
    278                        selected_item = self.selected_menu[level + 1] 
    279                    except IndexError: 
    280                        selected_item = -1 
    281 
    282                    def one_item( 
    283                        i: int, item: MenuItem 
    284                    ) -> Iterable[OneStyleAndTextTuple]: 
    285                        def mouse_handler(mouse_event: MouseEvent) -> None: 
    286                            if item.disabled: 
    287                                # The arrow keys can't interact with menu items that are disabled. 
    288                                # The mouse shouldn't be able to either. 
    289                                return 
    290                            hover = mouse_event.event_type == MouseEventType.MOUSE_MOVE 
    291                            if ( 
    292                                mouse_event.event_type == MouseEventType.MOUSE_UP 
    293                                or hover 
    294                            ): 
    295                                app = get_app() 
    296                                if not hover and item.handler: 
    297                                    app.layout.focus_last() 
    298                                    item.handler() 
    299                                else: 
    300                                    self.selected_menu = self.selected_menu[ 
    301                                        : level + 1 
    302                                    ] + [i] 
    303 
    304                        if i == selected_item: 
    305                            yield ("[SetCursorPosition]", "") 
    306                            style = "class:menu-bar.selected-item" 
    307                        else: 
    308                            style = "" 
    309 
    310                        yield ("class:menu", Border.VERTICAL) 
    311                        if item.text == "-": 
    312                            yield ( 
    313                                style + "class:menu-border", 
    314                                f"{Border.HORIZONTAL * (menu.width + 3)}", 
    315                                mouse_handler, 
    316                            ) 
    317                        else: 
    318                            yield ( 
    319                                style, 
    320                                f" {item.text}".ljust(menu.width + 3), 
    321                                mouse_handler, 
    322                            ) 
    323 
    324                        if item.children: 
    325                            yield (style, ">", mouse_handler) 
    326                        else: 
    327                            yield (style, " ", mouse_handler) 
    328 
    329                        if i == selected_item: 
    330                            yield ("[SetMenuPosition]", "") 
    331                        yield ("class:menu", Border.VERTICAL) 
    332 
    333                        yield ("", "\n") 
    334 
    335                    for i, item in enumerate(menu.children): 
    336                        result.extend(one_item(i, item)) 
    337 
    338                    result.append(("class:menu", Border.BOTTOM_LEFT)) 
    339                    result.append(("class:menu", Border.HORIZONTAL * (menu.width + 4))) 
    340                    result.append(("class:menu", Border.BOTTOM_RIGHT)) 
    341            return result 
    342 
    343        return Window(FormattedTextControl(get_text_fragments), style="class:menu") 
    344 
    345    @property 
    346    def floats(self) -> list[Float] | None: 
    347        return self.container.floats 
    348 
    349    def __pt_container__(self) -> Container: 
    350        return self.container 
    351 
    352 
    353class MenuItem: 
    354    def __init__( 
    355        self, 
    356        text: str = "", 
    357        handler: Callable[[], None] | None = None, 
    358        children: list[MenuItem] | None = None, 
    359        shortcut: Sequence[Keys | str] | None = None, 
    360        disabled: bool = False, 
    361    ) -> None: 
    362        self.text = text 
    363        self.handler = handler 
    364        self.children = children or [] 
    365        self.shortcut = shortcut 
    366        self.disabled = disabled 
    367        self.selected_item = 0 
    368 
    369    @property 
    370    def width(self) -> int: 
    371        if self.children: 
    372            return max(get_cwidth(c.text) for c in self.children) 
    373        else: 
    374            return 0