Skip to content
Back to top

# Source code for xonsh.shells.ptk_shell.updator

    """Has classes that help updating Prompt sections using Threads."""

    import concurrent.futures
    import threading
    import typing as tp

    from prompt_toolkit import PromptSession
    from prompt_toolkit.formatted_text import PygmentsTokens

    from xonsh.built_ins import XSH
    from xonsh.prompt.base import ParsedTokens
    from xonsh.style_tools import partial_color_tokenize, style_as_faded

    class Executor:
        """Caches thread results across prompts."""

        def __init__(self):
            self.thread_pool = concurrent.futures.ThreadPoolExecutor(
                max_workers=XSH.env["ASYNC_PROMPT_THREAD_WORKERS"]
            )

            # the attribute, .cache is cleared between calls.
            # This caches results from callback alone by field name.
            self.thread_results = {}

        def submit(self, func: tp.Callable, field: str):
            future = self.thread_pool.submit(self._run_func, func, field)
            place_holder = "{" + field + "}"

            return (
                future,
                (
                    self.thread_results[field]
                    if field in self.thread_results
                    else place_holder
                ),
                place_holder,
            )

        def _run_func(self, func, field):
            """Run the callback and store the result."""
            result = func()
            self.thread_results[field] = (
                result if result is None else style_as_faded(result)
            )
            return result

    class AsyncPrompt:
        """Represent an asynchronous prompt."""

        def __init__(self, name: str, session: PromptSession, executor: Executor):
            """

            Parameters
            ----------
            name: str
                what prompt to update. One of ['message', 'rprompt', 'bottom_toolbar']
            session: PromptSession
                current ptk session
            """

            self.name = name

            # list of tokens in that prompt. It could either be resolved or not resolved.
            self.tokens: ParsedTokens | None = None
            self.timer = None
            self.session = session
            self.executor = executor

            # (Key: the future object) that is created for the (value: index/field_name) in the tokens list
            self.futures: dict[
                concurrent.futures.Future,
                tuple[str, int | None, str | None, str | None],
            ] = {}

        def start_update(self, on_complete):
            """Listen on futures and update the prompt as each one completed.

            Timer is used to avoid clogging multiple calls at the same time.

            Parameters
            -----------
            on_complete:
                callback to notify after all the futures are completed
            """
            if not self.tokens:
                print(f"Warn: AsyncPrompt is created without tokens - {self.name}")
                return
            # Snapshot to avoid RuntimeError if stop() clears self.futures
            # from the main thread while we iterate here.
            futures = dict(self.futures)
            for fut in concurrent.futures.as_completed(futures):
                try:
                    val = fut.result()
                except concurrent.futures.CancelledError:
                    continue
                except Exception:
                    from xonsh.tools import print_exception

                    print_exception("Exception in async prompt field callback; ignored.")
                    continue

                placeholder, idx, spec, conv = futures[fut]
                # example: placeholder="{field}", idx=10, spec="env: {}"

                if isinstance(idx, int):
                    self.tokens.update(idx, val, spec, conv)
                else:  # when the function is called outside shell.
                    for idx, ptok in enumerate(self.tokens.tokens):
                        if placeholder in ptok.value:
                            val = ptok.value.replace(placeholder, val)
                            self.tokens.update(idx, val, spec, conv)

                # calling invalidate in less period is inefficient
                self.invalidate()

            on_complete(self.name)

        def invalidate(self):
            """Create a timer to update the prompt. The timing can be configured through env variables.
            threading.Timer is used to stop calling invalidate frequently.
            """
            from xonsh.shells.ptk_shell import tokenize_ansi

            if self.timer:
                self.timer.cancel()

            def _invalidate():
                new_prompt = self.tokens.process()
                formatted_tokens = tokenize_ansi(
                    PygmentsTokens(partial_color_tokenize(new_prompt))
                )
                app = self.session.app

                def _update_in_main():
                    setattr(self.session, self.name, formatted_tokens)
                    app.invalidate()

                if app.loop is not None:
                    app.loop.call_soon_threadsafe(_update_in_main)
                else:
                    _update_in_main()

            self.timer = threading.Timer(XSH.env["ASYNC_INVALIDATE_INTERVAL"], _invalidate)
            self.timer.start()

        def stop(self):
            """Stop any running threads"""
            if self.timer:
                self.timer.cancel()
            for fut in self.futures:
                fut.cancel()
            self.futures.clear()

        def submit_section(
            self,
            func: tp.Callable,
            field: str,
            idx: int | None = None,
            spec: str | None = None,
            conv=None,
        ):
            future, intermediate_value, placeholder = self.executor.submit(func, field)
            self.futures[future] = (placeholder, idx, spec, conv)
            return intermediate_value

    class PromptUpdator:
        """Handle updating multiple AsyncPrompt instances prompt/rprompt/bottom_toolbar"""

        def __init__(self, shell):
            from xonsh.shells.ptk_shell import PromptToolkitShell

            self.prompts: dict[str, AsyncPrompt] = {}
            self.shell: PromptToolkitShell = shell
            self.executor = Executor()
            self.futures = {}
            self.attrs_loaded = None

        def add(self, prompt_name: str | None) -> AsyncPrompt | None:
            # clear out old futures from the same prompt
            if prompt_name is None:
                return None

            self.stop(prompt_name)

            self.prompts[prompt_name] = AsyncPrompt(
                prompt_name, self.shell.prompter, self.executor
            )
            return self.prompts[prompt_name]

        def add_attrs(self):
            attrs = list(self.shell.get_lazy_ptk_kwargs())
            app = self.shell.prompter.app

            def _update_in_main():
                for attr, val in attrs:
                    setattr(self.shell.prompter, attr, val)
                app.invalidate()

            if app.loop is not None:
                app.loop.call_soon_threadsafe(_update_in_main)
            else:
                _update_in_main()

        def start(self):
            """after ptk prompt is created, update it in background."""
            if not self.attrs_loaded:
                self.attrs_loaded = self.executor.thread_pool.submit(self.add_attrs)

            prompts = list(self.prompts)  # removal safe
            for pt_name in prompts:
                if pt_name not in self.prompts:
                    continue
                prompt = self.prompts[pt_name]
                future = self.executor.thread_pool.submit(
                    prompt.start_update, self.on_complete
                )
                self.futures[pt_name] = future

        def stop(self, prompt_name: str):
            if prompt_name in self.prompts:
                self.prompts[prompt_name].stop()
            if prompt_name in self.futures:
                self.futures[prompt_name].cancel()

        def on_complete(self, prompt_name):
            self.prompts.pop(prompt_name, None)
            self.futures.pop(prompt_name, None)