xref: /aosp_15_r20/external/pigweed/pw_console/py/pw_console/embed.py (revision 61c4878ac05f98d0ceed94b57d316916de578985)
1# Copyright 2021 The Pigweed Authors
2#
3# Licensed under the Apache License, Version 2.0 (the "License"); you may not
4# use this file except in compliance with the License. You may obtain a copy of
5# the License at
6#
7#     https://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
11# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
12# License for the specific language governing permissions and limitations under
13# the License.
14"""pw_console embed class."""
15
16import asyncio
17import logging
18from pathlib import Path
19from typing import Any, Iterable
20
21from prompt_toolkit.completion import WordCompleter
22
23from pw_console.console_app import ConsoleApp
24from pw_console.get_pw_console_app import PW_CONSOLE_APP_CONTEXTVAR
25from pw_console.plugin_mixin import PluginMixin
26from pw_console.python_logging import (
27    setup_python_logging as pw_console_setup_python_logging,
28)
29from pw_console.widgets import (
30    FloatingWindowPane,
31    WindowPane,
32    WindowPaneToolbar,
33)
34
35
36def _set_console_app_instance(plugin: Any, console_app: ConsoleApp) -> None:
37    if hasattr(plugin, 'pw_console_init'):
38        plugin.pw_console_init(console_app)
39    else:
40        plugin.application = console_app
41
42
43def create_word_completer(
44    word_meta_dict: dict[str, str], ignore_case=True
45) -> WordCompleter:
46    sentences: list[str] = list(word_meta_dict.keys())
47    return WordCompleter(
48        sentences,
49        meta_dict=word_meta_dict,
50        ignore_case=ignore_case,
51        # Whole input field should match
52        sentence=True,
53    )
54
55
56class PwConsoleEmbed:
57    """Embed class for customizing the console before startup."""
58
59    # pylint: disable=too-many-instance-attributes
60    def __init__(
61        self,
62        global_vars=None,
63        local_vars=None,
64        loggers: dict[str, Iterable[logging.Logger]] | Iterable | None = None,
65        test_mode=False,
66        repl_startup_message: str | None = None,
67        help_text: str | None = None,
68        app_title: str | None = None,
69        config_file_path: str | Path | None = None,
70    ) -> None:
71        """Call this to embed pw console at the call point within your program.
72
73        Example usage:
74
75        .. code-block:: python
76
77            import logging
78
79            from pw_console import PwConsoleEmbed
80
81            # Create the pw_console embed instance
82            console = PwConsoleEmbed(
83                global_vars=globals(),
84                local_vars=locals(),
85                loggers={
86                    'Host Logs': [
87                        logging.getLogger(__package__),
88                        logging.getLogger(__name__),
89                    ],
90                    'Device Logs': [
91                        logging.getLogger('usb_gadget'),
92                    ],
93                },
94                app_title='My Awesome Console',
95                config_file_path='/home/user/project/.pw_console.yaml',
96            )
97            # Optional: Add custom completions
98            console.add_sentence_completer(
99                {
100                    'some_function', 'Function',
101                    'some_variable', 'Variable',
102                }
103            )
104
105            # Setup Python loggers to output to a file instead of STDOUT.
106            console.setup_python_logging()
107
108            # Then run the console with:
109            console.embed()
110
111        Args:
112            global_vars: dictionary representing the desired global symbol
113                table. Similar to what is returned by `globals()`.
114            local_vars: dictionary representing the desired local symbol
115                table. Similar to what is returned by `locals()`.
116            loggers: dict with keys of log window titles and values of either:
117
118                    1. List of `logging.getLogger()
119                       <https://docs.python.org/3/library/logging.html#logging.getLogger>`_
120                       instances.
121                    2. A single pw_console.log_store.LogStore instance.
122
123            app_title: Custom title text displayed in the user interface.
124            repl_startup_message: Custom text shown by default in the repl
125                output pane.
126            help_text: Custom text shown at the top of the help window before
127                keyboard shortcuts.
128            config_file_path: Path to a pw_console yaml config file.
129        """
130
131        self.global_vars = global_vars
132        self.local_vars = local_vars
133        self.loggers = loggers
134        self.test_mode = test_mode
135        self.repl_startup_message = repl_startup_message
136        self.help_text = help_text
137        self.app_title = app_title
138        self.config_file_path = (
139            Path(config_file_path) if config_file_path else None
140        )
141
142        self.console_app: ConsoleApp | None = None
143        self.extra_completers: list = []
144
145        self.setup_python_logging_called = False
146        self.hidden_by_default_windows: list[str] = []
147        self.window_plugins: list[WindowPane] = []
148        self.floating_window_plugins: list[tuple[FloatingWindowPane, dict]] = []
149        self.top_toolbar_plugins: list[WindowPaneToolbar] = []
150        self.bottom_toolbar_plugins: list[WindowPaneToolbar] = []
151
152    def add_window_plugin(self, window_pane: WindowPane) -> None:
153        """Include a custom window pane plugin.
154
155        Args:
156            window_pane: Any instance of the WindowPane class.
157        """
158        self.window_plugins.append(window_pane)
159
160    def add_floating_window_plugin(
161        self, window_pane: FloatingWindowPane, **float_args
162    ) -> None:
163        """Include a custom floating window pane plugin.
164
165        This adds a FloatingWindowPane class to the pw_console UI. The first
166        argument should be the window to add and the remaining keyword arguments
167        are passed to the prompt_toolkit Float() class. This allows positioning
168        of the floating window. By default the floating window will be
169        centered. To anchor the window to a side or corner of the screen set the
170        ``left``, ``right``, ``top``, or ``bottom`` keyword args.
171
172        For example:
173
174        .. code-block:: python
175
176           from pw_console import PwConsoleEmbed
177
178           console = PwConsoleEmbed(...)
179           my_plugin = MyPlugin()
180           # Anchor this floating window 2 rows away from the top and 4 columns
181           # away from the left edge of the screen.
182           console.add_floating_window_plugin(my_plugin, top=2, left=4)
183
184        See all possible keyword args in the prompt_toolkit documentation:
185        https://python-prompt-toolkit.readthedocs.io/en/stable/pages/reference.html#prompt_toolkit.layout.Float
186
187        Args:
188            window_pane: Any instance of the FloatingWindowPane class.
189            left: Distance to the left edge of the screen
190            right: Distance to the right edge of the screen
191            top: Distance to the top edge of the screen
192            bottom: Distance to the bottom edge of the screen
193        """
194        self.floating_window_plugins.append((window_pane, float_args))
195
196    def add_top_toolbar(self, toolbar: WindowPaneToolbar) -> None:
197        """Include a toolbar plugin to display on the top of the screen.
198
199        Top toolbars appear above all window panes and just below the main menu
200        bar. They span the full width of the screen.
201
202        Args:
203            toolbar: Instance of the WindowPaneToolbar class.
204        """
205        self.top_toolbar_plugins.append(toolbar)
206
207    def add_bottom_toolbar(self, toolbar: WindowPaneToolbar) -> None:
208        """Include a toolbar plugin to display at the bottom of the screen.
209
210        Bottom toolbars appear below all window panes and span the full width of
211        the screen.
212
213        Args:
214            toolbar: Instance of the WindowPaneToolbar class.
215        """
216        self.bottom_toolbar_plugins.append(toolbar)
217
218    def add_sentence_completer(
219        self, word_meta_dict: dict[str, str], ignore_case=True
220    ) -> None:
221        """Include a custom completer that matches on the entire repl input.
222
223        Args:
224            word_meta_dict: dictionary representing the sentence completions
225                and descriptions. Keys are completion text, values are
226                descriptions.
227        """
228
229        # Don't modify completion if empty.
230        if len(word_meta_dict) == 0:
231            return
232
233        word_completer = create_word_completer(word_meta_dict, ignore_case)
234        self.extra_completers.append(word_completer)
235
236    def _setup_log_panes(self) -> None:
237        """Add loggers to ConsoleApp log pane(s)."""
238        if not self.loggers:
239            return
240
241        assert isinstance(self.console_app, ConsoleApp)
242
243        if isinstance(self.loggers, list):
244            self.console_app.add_log_handler('Logs', self.loggers)
245
246        elif isinstance(self.loggers, dict):
247            for window_title, logger_instances in self.loggers.items():
248                window_pane = self.console_app.add_log_handler(
249                    window_title, logger_instances
250                )
251
252                if (
253                    window_pane
254                    and window_pane.pane_title()
255                    in self.hidden_by_default_windows
256                ):
257                    window_pane.show_pane = False
258
259    def setup_python_logging(
260        self,
261        last_resort_filename: str | None = None,
262        loggers_with_no_propagation: Iterable[logging.Logger] | None = None,
263    ) -> None:
264        """Setup friendly logging for full-screen prompt_toolkit applications.
265
266        This function sets up Python log handlers to be friendly for full-screen
267        prompt_toolkit applications. That is, logging to terminal STDOUT and
268        STDERR is disabled so the terminal user interface can be drawn.
269
270        Specifically, all Python STDOUT and STDERR log handlers are
271        disabled. It also sets `log propagation to True
272        <https://docs.python.org/3/library/logging.html#logging.Logger.propagate>`_.
273        to ensure that all log messages are sent to the root logger.
274
275        Args:
276            last_resort_filename: If specified use this file as a fallback for
277                unhandled Python logging messages. Normally Python will output
278                any log messages with no handlers to STDERR as a fallback. If
279                None, a temp file will be created instead. See Python
280                documentation on `logging.lastResort
281                <https://docs.python.org/3/library/logging.html#logging.lastResort>`_
282                for more info.
283            loggers_with_no_propagation: List of logger instances to skip
284               setting ``propagate = True``. This is useful if you would like
285               log messages from a particular source to not appear in the root
286               logger.
287        """
288        self.setup_python_logging_called = True
289        pw_console_setup_python_logging(
290            last_resort_filename, loggers_with_no_propagation
291        )
292
293    def hide_windows(self, *window_titles) -> None:
294        """Hide window panes specified by title on console startup."""
295        for window_title in window_titles:
296            self.hidden_by_default_windows.append(window_title)
297
298    def embed(self, override_window_config: dict | None = None) -> None:
299        """Start the console."""
300
301        # Create the ConsoleApp instance.
302        self.console_app = ConsoleApp(
303            global_vars=self.global_vars,
304            local_vars=self.local_vars,
305            repl_startup_message=self.repl_startup_message,
306            help_text=self.help_text,
307            app_title=self.app_title,
308            extra_completers=self.extra_completers,
309            floating_window_plugins=self.floating_window_plugins,
310        )
311        PW_CONSOLE_APP_CONTEXTVAR.set(self.console_app)  # type: ignore
312        # Setup Python logging and log panes.
313        if not self.setup_python_logging_called:
314            self.setup_python_logging()
315        self._setup_log_panes()
316
317        # Add window pane plugins to the layout.
318        for window_pane in self.window_plugins:
319            _set_console_app_instance(window_pane, self.console_app)
320            # Hide window plugins if the title is hidden by default.
321            if window_pane.pane_title() in self.hidden_by_default_windows:
322                window_pane.show_pane = False
323            self.console_app.window_manager.add_pane(window_pane)
324
325        # Add toolbar plugins to the layout.
326        for toolbar in self.top_toolbar_plugins:
327            _set_console_app_instance(toolbar, self.console_app)
328            self.console_app.window_manager.add_top_toolbar(toolbar)
329        for toolbar in self.bottom_toolbar_plugins:
330            _set_console_app_instance(toolbar, self.console_app)
331            self.console_app.window_manager.add_bottom_toolbar(toolbar)
332
333        # Init floating window plugins.
334        for floating_window, _ in self.floating_window_plugins:
335            _set_console_app_instance(floating_window, self.console_app)
336
337        # Rebuild prompt_toolkit containers, menu items, and help content with
338        # any new plugins added above.
339        self.console_app.refresh_layout()
340
341        # Load external config if passed in.
342        if self.config_file_path:
343            self.console_app.load_config(self.config_file_path)
344
345        if override_window_config:
346            self.console_app.prefs.set_windows(override_window_config)
347        self.console_app.apply_window_config()
348
349        # Hide the repl pane if it's in the hidden windows list.
350        if 'Python Repl' in self.hidden_by_default_windows:
351            self.console_app.repl_pane.show_pane = False
352
353        # Start a thread for running user code.
354        self.console_app.start_user_code_thread()
355
356        # Startup any background threads and tasks required by plugins.
357        for window_pane in self.window_plugins:
358            if isinstance(window_pane, PluginMixin):
359                window_pane.plugin_start()
360        for toolbar in self.bottom_toolbar_plugins:
361            if isinstance(toolbar, PluginMixin):
362                toolbar.plugin_start()
363        for toolbar in self.top_toolbar_plugins:
364            if isinstance(toolbar, PluginMixin):
365                toolbar.plugin_start()
366
367        # Start the prompt_toolkit UI app.
368        asyncio.run(
369            self.console_app.run(test_mode=self.test_mode), debug=self.test_mode
370        )
371