1""" 
    2Wrapper for the layout. 
    3""" 
    4 
    5from __future__ import annotations 
    6 
    7from typing import Generator, Iterable, Union 
    8 
    9from prompt_toolkit.buffer import Buffer 
    10 
    11from .containers import ( 
    12    AnyContainer, 
    13    ConditionalContainer, 
    14    Container, 
    15    Window, 
    16    to_container, 
    17) 
    18from .controls import BufferControl, SearchBufferControl, UIControl 
    19 
    20__all__ = [ 
    21    "Layout", 
    22    "InvalidLayoutError", 
    23    "walk", 
    24] 
    25 
    26FocusableElement = Union[str, Buffer, UIControl, AnyContainer] 
    27 
    28 
    29class Layout: 
    30    """ 
    31    The layout for a prompt_toolkit 
    32    :class:`~prompt_toolkit.application.Application`. 
    33    This also keeps track of which user control is focused. 
    34 
    35    :param container: The "root" container for the layout. 
    36    :param focused_element: element to be focused initially. (Can be anything 
    37        the `focus` function accepts.) 
    38    """ 
    39 
    40    def __init__( 
    41        self, 
    42        container: AnyContainer, 
    43        focused_element: FocusableElement | None = None, 
    44    ) -> None: 
    45        self.container = to_container(container) 
    46        self._stack: list[Window] = [] 
    47 
    48        # Map search BufferControl back to the original BufferControl. 
    49        # This is used to keep track of when exactly we are searching, and for 
    50        # applying the search. 
    51        # When a link exists in this dictionary, that means the search is 
    52        # currently active. 
    53        # Map: search_buffer_control -> original buffer control. 
    54        self.search_links: dict[SearchBufferControl, BufferControl] = {} 
    55 
    56        # Mapping that maps the children in the layout to their parent. 
    57        # This relationship is calculated dynamically, each time when the UI 
    58        # is rendered.  (UI elements have only references to their children.) 
    59        self._child_to_parent: dict[Container, Container] = {} 
    60 
    61        if focused_element is None: 
    62            try: 
    63                self._stack.append(next(self.find_all_windows())) 
    64            except StopIteration as e: 
    65                raise InvalidLayoutError( 
    66                    "Invalid layout. The layout does not contain any Window object." 
    67                ) from e 
    68        else: 
    69            self.focus(focused_element) 
    70 
    71        # List of visible windows. 
    72        self.visible_windows: list[Window] = []  # List of `Window` objects. 
    73 
    74    def __repr__(self) -> str: 
    75        return f"Layout({self.container!r}, current_window={self.current_window!r})" 
    76 
    77    def find_all_windows(self) -> Generator[Window, None, None]: 
    78        """ 
    79        Find all the :class:`.UIControl` objects in this layout. 
    80        """ 
    81        for item in self.walk(): 
    82            if isinstance(item, Window): 
    83                yield item 
    84 
    85    def find_all_controls(self) -> Iterable[UIControl]: 
    86        for container in self.find_all_windows(): 
    87            yield container.content 
    88 
    89    def focus(self, value: FocusableElement) -> None: 
    90        """ 
    91        Focus the given UI element. 
    92 
    93        `value` can be either: 
    94 
    95        - a :class:`.UIControl` 
    96        - a :class:`.Buffer` instance or the name of a :class:`.Buffer` 
    97        - a :class:`.Window` 
    98        - Any container object. In this case we will focus the :class:`.Window` 
    99          from this container that was focused most recent, or the very first 
    100          focusable :class:`.Window` of the container. 
    101        """ 
    102        # BufferControl by buffer name. 
    103        if isinstance(value, str): 
    104            for control in self.find_all_controls(): 
    105                if isinstance(control, BufferControl) and control.buffer.name == value: 
    106                    self.focus(control) 
    107                    return 
    108            raise ValueError(f"Couldn't find Buffer in the current layout: {value!r}.") 
    109 
    110        # BufferControl by buffer object. 
    111        elif isinstance(value, Buffer): 
    112            for control in self.find_all_controls(): 
    113                if isinstance(control, BufferControl) and control.buffer == value: 
    114                    self.focus(control) 
    115                    return 
    116            raise ValueError(f"Couldn't find Buffer in the current layout: {value!r}.") 
    117 
    118        # Focus UIControl. 
    119        elif isinstance(value, UIControl): 
    120            if value not in self.find_all_controls(): 
    121                raise ValueError( 
    122                    "Invalid value. Container does not appear in the layout." 
    123                ) 
    124            if not value.is_focusable(): 
    125                raise ValueError("Invalid value. UIControl is not focusable.") 
    126 
    127            self.current_control = value 
    128 
    129        # Otherwise, expecting any Container object. 
    130        else: 
    131            value = to_container(value) 
    132 
    133            if isinstance(value, Window): 
    134                # This is a `Window`: focus that. 
    135                if value not in self.find_all_windows(): 
    136                    raise ValueError( 
    137                        f"Invalid value. Window does not appear in the layout: {value!r}" 
    138                    ) 
    139 
    140                self.current_window = value 
    141            else: 
    142                # Focus a window in this container. 
    143                # If we have many windows as part of this container, and some 
    144                # of them have been focused before, take the last focused 
    145                # item. (This is very useful when the UI is composed of more 
    146                # complex sub components.) 
    147                windows = [] 
    148                for c in walk(value, skip_hidden=True): 
    149                    if isinstance(c, Window) and c.content.is_focusable(): 
    150                        windows.append(c) 
    151 
    152                # Take the first one that was focused before. 
    153                for w in reversed(self._stack): 
    154                    if w in windows: 
    155                        self.current_window = w 
    156                        return 
    157 
    158                # None was focused before: take the very first focusable window. 
    159                if windows: 
    160                    self.current_window = windows[0] 
    161                    return 
    162 
    163                raise ValueError( 
    164                    f"Invalid value. Container cannot be focused: {value!r}" 
    165                ) 
    166 
    167    def has_focus(self, value: FocusableElement) -> bool: 
    168        """ 
    169        Check whether the given control has the focus. 
    170        :param value: :class:`.UIControl` or :class:`.Window` instance. 
    171        """ 
    172        if isinstance(value, str): 
    173            if self.current_buffer is None: 
    174                return False 
    175            return self.current_buffer.name == value 
    176        if isinstance(value, Buffer): 
    177            return self.current_buffer == value 
    178        if isinstance(value, UIControl): 
    179            return self.current_control == value 
    180        else: 
    181            value = to_container(value) 
    182            if isinstance(value, Window): 
    183                return self.current_window == value 
    184            else: 
    185                # Check whether this "container" is focused. This is true if 
    186                # one of the elements inside is focused. 
    187                for element in walk(value): 
    188                    if element == self.current_window: 
    189                        return True 
    190                return False 
    191 
    192    @property 
    193    def current_control(self) -> UIControl: 
    194        """ 
    195        Get the :class:`.UIControl` to currently has the focus. 
    196        """ 
    197        return self._stack[-1].content 
    198 
    199    @current_control.setter 
    200    def current_control(self, control: UIControl) -> None: 
    201        """ 
    202        Set the :class:`.UIControl` to receive the focus. 
    203        """ 
    204        for window in self.find_all_windows(): 
    205            if window.content == control: 
    206                self.current_window = window 
    207                return 
    208 
    209        raise ValueError("Control not found in the user interface.") 
    210 
    211    @property 
    212    def current_window(self) -> Window: 
    213        "Return the :class:`.Window` object that is currently focused." 
    214        return self._stack[-1] 
    215 
    216    @current_window.setter 
    217    def current_window(self, value: Window) -> None: 
    218        "Set the :class:`.Window` object to be currently focused." 
    219        self._stack.append(value) 
    220 
    221    @property 
    222    def is_searching(self) -> bool: 
    223        "True if we are searching right now." 
    224        return self.current_control in self.search_links 
    225 
    226    @property 
    227    def search_target_buffer_control(self) -> BufferControl | None: 
    228        """ 
    229        Return the :class:`.BufferControl` in which we are searching or `None`. 
    230        """ 
    231        # Not every `UIControl` is a `BufferControl`. This only applies to 
    232        # `BufferControl`. 
    233        control = self.current_control 
    234 
    235        if isinstance(control, SearchBufferControl): 
    236            return self.search_links.get(control) 
    237        else: 
    238            return None 
    239 
    240    def get_focusable_windows(self) -> Iterable[Window]: 
    241        """ 
    242        Return all the :class:`.Window` objects which are focusable (in the 
    243        'modal' area). 
    244        """ 
    245        for w in self.walk_through_modal_area(): 
    246            if isinstance(w, Window) and w.content.is_focusable(): 
    247                yield w 
    248 
    249    def get_visible_focusable_windows(self) -> list[Window]: 
    250        """ 
    251        Return a list of :class:`.Window` objects that are focusable. 
    252        """ 
    253        # focusable windows are windows that are visible, but also part of the 
    254        # modal container. Make sure to keep the ordering. 
    255        visible_windows = self.visible_windows 
    256        return [w for w in self.get_focusable_windows() if w in visible_windows] 
    257 
    258    @property 
    259    def current_buffer(self) -> Buffer | None: 
    260        """ 
    261        The currently focused :class:`~.Buffer` or `None`. 
    262        """ 
    263        ui_control = self.current_control 
    264        if isinstance(ui_control, BufferControl): 
    265            return ui_control.buffer 
    266        return None 
    267 
    268    def get_buffer_by_name(self, buffer_name: str) -> Buffer | None: 
    269        """ 
    270        Look in the layout for a buffer with the given name. 
    271        Return `None` when nothing was found. 
    272        """ 
    273        for w in self.walk(): 
    274            if isinstance(w, Window) and isinstance(w.content, BufferControl): 
    275                if w.content.buffer.name == buffer_name: 
    276                    return w.content.buffer 
    277        return None 
    278 
    279    @property 
    280    def buffer_has_focus(self) -> bool: 
    281        """ 
    282        Return `True` if the currently focused control is a 
    283        :class:`.BufferControl`. (For instance, used to determine whether the 
    284        default key bindings should be active or not.) 
    285        """ 
    286        ui_control = self.current_control 
    287        return isinstance(ui_control, BufferControl) 
    288 
    289    @property 
    290    def previous_control(self) -> UIControl: 
    291        """ 
    292        Get the :class:`.UIControl` to previously had the focus. 
    293        """ 
    294        try: 
    295            return self._stack[-2].content 
    296        except IndexError: 
    297            return self._stack[-1].content 
    298 
    299    def focus_last(self) -> None: 
    300        """ 
    301        Give the focus to the last focused control. 
    302        """ 
    303        if len(self._stack) > 1: 
    304            self._stack = self._stack[:-1] 
    305 
    306    def focus_next(self) -> None: 
    307        """ 
    308        Focus the next visible/focusable Window. 
    309        """ 
    310        windows = self.get_visible_focusable_windows() 
    311 
    312        if len(windows) > 0: 
    313            try: 
    314                index = windows.index(self.current_window) 
    315            except ValueError: 
    316                index = 0 
    317            else: 
    318                index = (index + 1) % len(windows) 
    319 
    320            self.focus(windows[index]) 
    321 
    322    def focus_previous(self) -> None: 
    323        """ 
    324        Focus the previous visible/focusable Window. 
    325        """ 
    326        windows = self.get_visible_focusable_windows() 
    327 
    328        if len(windows) > 0: 
    329            try: 
    330                index = windows.index(self.current_window) 
    331            except ValueError: 
    332                index = 0 
    333            else: 
    334                index = (index - 1) % len(windows) 
    335 
    336            self.focus(windows[index]) 
    337 
    338    def walk(self) -> Iterable[Container]: 
    339        """ 
    340        Walk through all the layout nodes (and their children) and yield them. 
    341        """ 
    342        yield from walk(self.container) 
    343 
    344    def walk_through_modal_area(self) -> Iterable[Container]: 
    345        """ 
    346        Walk through all the containers which are in the current 'modal' part 
    347        of the layout. 
    348        """ 
    349        # Go up in the tree, and find the root. (it will be a part of the 
    350        # layout, if the focus is in a modal part.) 
    351        root: Container = self.current_window 
    352        while not root.is_modal() and root in self._child_to_parent: 
    353            root = self._child_to_parent[root] 
    354 
    355        yield from walk(root) 
    356 
    357    def update_parents_relations(self) -> None: 
    358        """ 
    359        Update child->parent relationships mapping. 
    360        """ 
    361        parents = {} 
    362 
    363        def walk(e: Container) -> None: 
    364            for c in e.get_children(): 
    365                parents[c] = e 
    366                walk(c) 
    367 
    368        walk(self.container) 
    369 
    370        self._child_to_parent = parents 
    371 
    372    def reset(self) -> None: 
    373        # Remove all search links when the UI starts. 
    374        # (Important, for instance when control-c is been pressed while 
    375        #  searching. The prompt cancels, but next `run()` call the search 
    376        #  links are still there.) 
    377        self.search_links.clear() 
    378 
    379        self.container.reset() 
    380 
    381    def get_parent(self, container: Container) -> Container | None: 
    382        """ 
    383        Return the parent container for the given container, or ``None``, if it 
    384        wasn't found. 
    385        """ 
    386        try: 
    387            return self._child_to_parent[container] 
    388        except KeyError: 
    389            return None 
    390 
    391 
    392class InvalidLayoutError(Exception): 
    393    pass 
    394 
    395 
    396def walk(container: Container, skip_hidden: bool = False) -> Iterable[Container]: 
    397    """ 
    398    Walk through layout, starting at this container. 
    399    """ 
    400    # When `skip_hidden` is set, don't go into disabled ConditionalContainer containers. 
    401    if ( 
    402        skip_hidden 
    403        and isinstance(container, ConditionalContainer) 
    404        and not container.filter() 
    405    ): 
    406        return 
    407 
    408    yield container 
    409 
    410    for c in container.get_children(): 
    411        # yield from walk(c) 
    412        yield from walk(c, skip_hidden=skip_hidden)