1# -*- coding: utf-8 -*-
2"""The prompt_toolkit based xonsh shell."""
3import sys
4import builtins
5
6from prompt_toolkit.key_binding.manager import KeyBindingManager
7from prompt_toolkit.auto_suggest import AutoSuggestFromHistory
8from prompt_toolkit.layout.lexers import PygmentsLexer
9from prompt_toolkit.shortcuts import print_tokens
10from prompt_toolkit.styles import PygmentsStyle, style_from_dict
11
12from xonsh.base_shell import BaseShell
13from xonsh.tools import print_exception, carriage_return, ansicolors_to_ptk1_names
14from xonsh.ptk.completer import PromptToolkitCompleter
15from xonsh.ptk.history import PromptToolkitHistory
16from xonsh.ptk.key_bindings import load_xonsh_bindings
17from xonsh.ptk.shortcuts import Prompter
18from xonsh.events import events
19from xonsh.shell import transform_command
20from xonsh.platform import HAS_PYGMENTS, ON_WINDOWS
21from xonsh.style_tools import (
22    partial_color_tokenize,
23    _TokenType,
24    DEFAULT_STYLE_DICT as _DEFAULT_STYLE_DICT,
25)
26from xonsh.lazyimps import pygments, pyghooks, winutils
27from xonsh.pygments_cache import get_all_styles
28from xonsh.lazyasd import LazyObject
29
30
31Token = _TokenType()
32
33events.transmogrify("on_ptk_create", "LoadEvent")
34events.doc(
35    "on_ptk_create",
36    """
37on_ptk_create(prompter: Prompter, history: PromptToolkitHistory, completer: PromptToolkitCompleter, bindings: KeyBindingManager) ->
38
39Fired after prompt toolkit has been initialized
40""",
41)
42
43# Convert new ansicolor names to names
44# understood by PTK1
45DEFAULT_STYLE_DICT = LazyObject(
46    lambda: ansicolors_to_ptk1_names(_DEFAULT_STYLE_DICT),
47    globals(),
48    "DEFAULT_STYLE_DICT",
49)
50
51
52class PromptToolkitShell(BaseShell):
53    """The xonsh shell."""
54
55    def __init__(self, **kwargs):
56        super().__init__(**kwargs)
57        if ON_WINDOWS:
58            winutils.enable_virtual_terminal_processing()
59        self._first_prompt = True
60        self.prompter = Prompter()
61        self.history = PromptToolkitHistory()
62        self.pt_completer = PromptToolkitCompleter(self.completer, self.ctx, self)
63        key_bindings_manager_args = {
64            "enable_auto_suggest_bindings": True,
65            "enable_search": True,
66            "enable_abort_and_exit_bindings": True,
67        }
68        self.key_bindings_manager = KeyBindingManager(**key_bindings_manager_args)
69        load_xonsh_bindings(self.key_bindings_manager)
70        # This assumes that PromptToolkitShell is a singleton
71        events.on_ptk_create.fire(
72            prompter=self.prompter,
73            history=self.history,
74            completer=self.pt_completer,
75            bindings=self.key_bindings_manager,
76        )
77
78    def singleline(
79        self,
80        store_in_history=True,
81        auto_suggest=None,
82        enable_history_search=True,
83        multiline=True,
84        **kwargs
85    ):
86        """Reads a single line of input from the shell. The store_in_history
87        kwarg flags whether the input should be stored in PTK's in-memory
88        history.
89        """
90        events.on_pre_prompt.fire()
91        env = builtins.__xonsh_env__
92        mouse_support = env.get("MOUSE_SUPPORT")
93        if store_in_history:
94            history = self.history
95        else:
96            history = None
97            enable_history_search = False
98        auto_suggest = auto_suggest if env.get("AUTO_SUGGEST") else None
99        completions_display = env.get("COMPLETIONS_DISPLAY")
100        multicolumn = completions_display == "multi"
101        complete_while_typing = env.get("UPDATE_COMPLETIONS_ON_KEYPRESS")
102        if complete_while_typing:
103            # PTK requires history search to be none when completing while typing
104            enable_history_search = False
105        if HAS_PYGMENTS:
106            self.styler.style_name = env.get("XONSH_COLOR_STYLE")
107        completer = None if completions_display == "none" else self.pt_completer
108        if not env.get("UPDATE_PROMPT_ON_KEYPRESS"):
109            prompt_tokens_cached = self.prompt_tokens(None)
110            get_prompt_tokens = lambda cli: prompt_tokens_cached
111            rprompt_tokens_cached = self.rprompt_tokens(None)
112            get_rprompt_tokens = lambda cli: rprompt_tokens_cached
113            bottom_toolbar_tokens_cached = self.bottom_toolbar_tokens(None)
114            get_bottom_toolbar_tokens = lambda cli: bottom_toolbar_tokens_cached
115        else:
116            get_prompt_tokens = self.prompt_tokens
117            get_rprompt_tokens = self.rprompt_tokens
118            get_bottom_toolbar_tokens = self.bottom_toolbar_tokens
119
120        with self.prompter:
121            prompt_args = {
122                "mouse_support": mouse_support,
123                "auto_suggest": auto_suggest,
124                "get_prompt_tokens": get_prompt_tokens,
125                "get_rprompt_tokens": get_rprompt_tokens,
126                "get_bottom_toolbar_tokens": get_bottom_toolbar_tokens,
127                "completer": completer,
128                "multiline": multiline,
129                "get_continuation_tokens": self.continuation_tokens,
130                "history": history,
131                "enable_history_search": enable_history_search,
132                "reserve_space_for_menu": 0,
133                "key_bindings_registry": self.key_bindings_manager.registry,
134                "display_completions_in_columns": multicolumn,
135                "complete_while_typing": complete_while_typing,
136            }
137            if builtins.__xonsh_env__.get("COLOR_INPUT"):
138                if HAS_PYGMENTS:
139                    prompt_args["lexer"] = PygmentsLexer(pyghooks.XonshLexer)
140                    prompt_args["style"] = PygmentsStyle(
141                        pyghooks.xonsh_style_proxy(self.styler)
142                    )
143                else:
144                    prompt_args["style"] = style_from_dict(DEFAULT_STYLE_DICT)
145            line = self.prompter.prompt(**prompt_args)
146            events.on_post_prompt.fire()
147        return line
148
149    def _push(self, line):
150        """Pushes a line onto the buffer and compiles the code in a way that
151        enables multiline input.
152        """
153        code = None
154        self.buffer.append(line)
155        if self.need_more_lines:
156            return None, code
157        src = "".join(self.buffer)
158        src = transform_command(src)
159        try:
160            code = self.execer.compile(src, mode="single", glbs=self.ctx, locs=None)
161            self.reset_buffer()
162        except Exception:  # pylint: disable=broad-except
163            self.reset_buffer()
164            print_exception()
165            return src, None
166        return src, code
167
168    def cmdloop(self, intro=None):
169        """Enters a loop that reads and execute input from user."""
170        if intro:
171            print(intro)
172        auto_suggest = AutoSuggestFromHistory()
173        self.push = self._push
174        while not builtins.__xonsh_exit__:
175            try:
176                line = self.singleline(auto_suggest=auto_suggest)
177                if not line:
178                    self.emptyline()
179                else:
180                    line = self.precmd(line)
181                    self.default(line)
182            except (KeyboardInterrupt, SystemExit):
183                self.reset_buffer()
184            except EOFError:
185                if builtins.__xonsh_env__.get("IGNOREEOF"):
186                    print('Use "exit" to leave the shell.', file=sys.stderr)
187                else:
188                    break
189
190    def prompt_tokens(self, cli):
191        """Returns a list of (token, str) tuples for the current prompt."""
192        p = builtins.__xonsh_env__.get("PROMPT")
193        try:
194            p = self.prompt_formatter(p)
195        except Exception:  # pylint: disable=broad-except
196            print_exception()
197        toks = partial_color_tokenize(p)
198        if self._first_prompt:
199            carriage_return()
200            self._first_prompt = False
201        self.settitle()
202        return toks
203
204    def rprompt_tokens(self, cli):
205        """Returns a list of (token, str) tuples for the current right
206        prompt.
207        """
208        p = builtins.__xonsh_env__.get("RIGHT_PROMPT")
209        # self.prompt_formatter does handle empty strings properly,
210        # but this avoids descending into it in the common case of
211        # $RIGHT_PROMPT == ''.
212        if isinstance(p, str) and len(p) == 0:
213            return []
214        try:
215            p = self.prompt_formatter(p)
216        except Exception:  # pylint: disable=broad-except
217            print_exception()
218        toks = partial_color_tokenize(p)
219        return toks
220
221    def bottom_toolbar_tokens(self, cli):
222        """Returns a list of (token, str) tuples for the current bottom
223        toolbar.
224        """
225        p = builtins.__xonsh_env__.get("BOTTOM_TOOLBAR")
226        # self.prompt_formatter does handle empty strings properly,
227        # but this avoids descending into it in the common case of
228        # $TOOLBAR == ''.
229        if isinstance(p, str) and len(p) == 0:
230            return []
231        try:
232            p = self.prompt_formatter(p)
233        except Exception:  # pylint: disable=broad-except
234            print_exception()
235        toks = partial_color_tokenize(p)
236        return toks
237
238    def continuation_tokens(self, cli, width):
239        """Displays dots in multiline prompt"""
240        width = width - 1
241        dots = builtins.__xonsh_env__.get("MULTILINE_PROMPT")
242        dots = dots() if callable(dots) else dots
243        if dots is None:
244            return [(Token, " " * (width + 1))]
245        basetoks = self.format_color(dots)
246        baselen = sum(len(t[1]) for t in basetoks)
247        if baselen == 0:
248            return [(Token, " " * (width + 1))]
249        toks = basetoks * (width // baselen)
250        n = width % baselen
251        count = 0
252        for tok in basetoks:
253            slen = len(tok[1])
254            newcount = slen + count
255            if slen == 0:
256                continue
257            elif newcount <= n:
258                toks.append(tok)
259            else:
260                toks.append((tok[0], tok[1][: n - count]))
261            count = newcount
262            if n <= count:
263                break
264        toks.append((Token, " "))  # final space
265        return toks
266
267    def format_color(self, string, hide=False, force_string=False, **kwargs):
268        """Formats a color string using Pygments. This, therefore, returns
269        a list of (Token, str) tuples. If force_string is set to true, though,
270        this will return a color formatted string.
271        """
272        tokens = partial_color_tokenize(string)
273        if force_string and HAS_PYGMENTS:
274            env = builtins.__xonsh_env__
275            self.styler.style_name = env.get("XONSH_COLOR_STYLE")
276            proxy_style = pyghooks.xonsh_style_proxy(self.styler)
277            formatter = pyghooks.XonshTerminal256Formatter(style=proxy_style)
278            s = pygments.format(tokens, formatter)
279            return s
280        elif force_string:
281            print("To force colorization of string, install Pygments")
282            return tokens
283        else:
284            return tokens
285
286    def print_color(self, string, end="\n", **kwargs):
287        """Prints a color string using prompt-toolkit color management."""
288        if isinstance(string, str):
289            tokens = partial_color_tokenize(string + end)
290        else:
291            # assume this is a list of (Token, str) tuples and just print
292            tokens = string
293        if HAS_PYGMENTS:
294            env = builtins.__xonsh_env__
295            self.styler.style_name = env.get("XONSH_COLOR_STYLE")
296            proxy_style = PygmentsStyle(pyghooks.xonsh_style_proxy(self.styler))
297        else:
298            proxy_style = style_from_dict(DEFAULT_STYLE_DICT)
299        print_tokens(tokens, style=proxy_style)
300
301    def color_style_names(self):
302        """Returns an iterable of all available style names."""
303        if not HAS_PYGMENTS:
304            return ["For other xonsh styles, please install pygments"]
305        return get_all_styles()
306
307    def color_style(self):
308        """Returns the current color map."""
309        if not HAS_PYGMENTS:
310            return DEFAULT_STYLE_DICT
311        env = builtins.__xonsh_env__
312        self.styler.style_name = env.get("XONSH_COLOR_STYLE")
313        return self.styler.styles
314
315    def restore_tty_sanity(self):
316        """An interface for resetting the TTY stdin mode. This is highly
317        dependent on the shell backend. Also it is mostly optional since
318        it only affects ^Z backgrounding behaviour.
319        """
320        # PTK does not seem to need any specialization here. However,
321        # if it does for some reason in the future...
322        # The following writes an ANSI escape sequence that sends the cursor
323        # to the end of the line. This has the effect of restoring ECHO mode.
324        # See http://unix.stackexchange.com/a/108014/129048 for more details.
325        # This line can also be replaced by os.system("stty sane"), as per
326        # http://stackoverflow.com/questions/19777129/interactive-python-interpreter-run-in-background#comment29421919_19778355
327        # However, it is important to note that not termios-based solution
328        # seems to work. My guess is that this is because termios restoration
329        # needs to be performed by the subprocess itself. This fix is important
330        # when subprocesses don't properly restore the terminal attributes,
331        # like Python in interactive mode. Also note that the sequences "\033M"
332        # and "\033E" seem to work too, but these are technically VT100 codes.
333        # I used the more primitive ANSI sequence to maximize compatibility.
334        # -scopatz 2017-01-28
335        #   if not ON_POSIX:
336        #       return
337        #   sys.stdout.write('\033[9999999C\n')
338