1"""
2
3Adaptor for using the input system of `prompt_toolkit` with the IPython
4backend.
5
6This gives a powerful interactive shell that has a nice user interface, but
7also the power of for instance all the %-magic functions that IPython has to
8offer.
9
10"""
11from warnings import warn
12
13from IPython import utils as ipy_utils
14from IPython.core.inputsplitter import IPythonInputSplitter
15from IPython.terminal.embed import InteractiveShellEmbed as _InteractiveShellEmbed
16from IPython.terminal.ipapp import load_default_config
17from prompt_toolkit.completion import (
18    Completer,
19    Completion,
20    PathCompleter,
21    WordCompleter,
22)
23from prompt_toolkit.contrib.completers import SystemCompleter
24from prompt_toolkit.contrib.regular_languages.compiler import compile
25from prompt_toolkit.contrib.regular_languages.completion import GrammarCompleter
26from prompt_toolkit.contrib.regular_languages.lexer import GrammarLexer
27from prompt_toolkit.document import Document
28from prompt_toolkit.formatted_text import PygmentsTokens
29from prompt_toolkit.lexers import PygmentsLexer, SimpleLexer
30from prompt_toolkit.styles import Style
31from pygments.lexers import BashLexer, PythonLexer
32
33from ptpython.prompt_style import PromptStyle
34
35from .python_input import PythonCompleter, PythonInput, PythonValidator
36from .style import default_ui_style
37
38__all__ = ["embed"]
39
40
41class IPythonPrompt(PromptStyle):
42    """
43    Style for IPython >5.0, use the prompt_toolkit tokens directly.
44    """
45
46    def __init__(self, prompts):
47        self.prompts = prompts
48
49    def in_prompt(self):
50        return PygmentsTokens(self.prompts.in_prompt_tokens())
51
52    def in2_prompt(self, width):
53        return PygmentsTokens(self.prompts.continuation_prompt_tokens())
54
55    def out_prompt(self):
56        return []
57
58
59class IPythonValidator(PythonValidator):
60    def __init__(self, *args, **kwargs):
61        super(IPythonValidator, self).__init__(*args, **kwargs)
62        self.isp = IPythonInputSplitter()
63
64    def validate(self, document):
65        document = Document(text=self.isp.transform_cell(document.text))
66        super(IPythonValidator, self).validate(document)
67
68
69def create_ipython_grammar():
70    """
71    Return compiled IPython grammar.
72    """
73    return compile(
74        r"""
75        \s*
76        (
77            (?P<percent>%)(
78                (?P<magic>pycat|run|loadpy|load)  \s+ (?P<py_filename>[^\s]+)  |
79                (?P<magic>cat)                    \s+ (?P<filename>[^\s]+)     |
80                (?P<magic>pushd|cd|ls)            \s+ (?P<directory>[^\s]+)    |
81                (?P<magic>pdb)                    \s+ (?P<pdb_arg>[^\s]+)      |
82                (?P<magic>autocall)               \s+ (?P<autocall_arg>[^\s]+) |
83                (?P<magic>time|timeit|prun)       \s+ (?P<python>.+)           |
84                (?P<magic>psource|pfile|pinfo|pinfo2) \s+ (?P<python>.+)       |
85                (?P<magic>system)                 \s+ (?P<system>.+)           |
86                (?P<magic>unalias)                \s+ (?P<alias_name>.+)       |
87                (?P<magic>[^\s]+)   .* |
88            ) .*            |
89            !(?P<system>.+) |
90            (?![%!]) (?P<python>.+)
91        )
92        \s*
93    """
94    )
95
96
97def create_completer(
98    get_globals,
99    get_locals,
100    magics_manager,
101    alias_manager,
102    get_enable_dictionary_completion,
103):
104    g = create_ipython_grammar()
105
106    return GrammarCompleter(
107        g,
108        {
109            "python": PythonCompleter(
110                get_globals, get_locals, get_enable_dictionary_completion
111            ),
112            "magic": MagicsCompleter(magics_manager),
113            "alias_name": AliasCompleter(alias_manager),
114            "pdb_arg": WordCompleter(["on", "off"], ignore_case=True),
115            "autocall_arg": WordCompleter(["0", "1", "2"], ignore_case=True),
116            "py_filename": PathCompleter(
117                only_directories=False, file_filter=lambda name: name.endswith(".py")
118            ),
119            "filename": PathCompleter(only_directories=False),
120            "directory": PathCompleter(only_directories=True),
121            "system": SystemCompleter(),
122        },
123    )
124
125
126def create_lexer():
127    g = create_ipython_grammar()
128
129    return GrammarLexer(
130        g,
131        lexers={
132            "percent": SimpleLexer("class:pygments.operator"),
133            "magic": SimpleLexer("class:pygments.keyword"),
134            "filename": SimpleLexer("class:pygments.name"),
135            "python": PygmentsLexer(PythonLexer),
136            "system": PygmentsLexer(BashLexer),
137        },
138    )
139
140
141class MagicsCompleter(Completer):
142    def __init__(self, magics_manager):
143        self.magics_manager = magics_manager
144
145    def get_completions(self, document, complete_event):
146        text = document.text_before_cursor.lstrip()
147
148        for m in sorted(self.magics_manager.magics["line"]):
149            if m.startswith(text):
150                yield Completion("%s" % m, -len(text))
151
152
153class AliasCompleter(Completer):
154    def __init__(self, alias_manager):
155        self.alias_manager = alias_manager
156
157    def get_completions(self, document, complete_event):
158        text = document.text_before_cursor.lstrip()
159        # aliases = [a for a, _ in self.alias_manager.aliases]
160        aliases = self.alias_manager.aliases
161
162        for a, cmd in sorted(aliases, key=lambda a: a[0]):
163            if a.startswith(text):
164                yield Completion("%s" % a, -len(text), display_meta=cmd)
165
166
167class IPythonInput(PythonInput):
168    """
169    Override our `PythonCommandLineInterface` to add IPython specific stuff.
170    """
171
172    def __init__(self, ipython_shell, *a, **kw):
173        kw["_completer"] = create_completer(
174            kw["get_globals"],
175            kw["get_globals"],
176            ipython_shell.magics_manager,
177            ipython_shell.alias_manager,
178            lambda: self.enable_dictionary_completion,
179        )
180        kw["_lexer"] = create_lexer()
181        kw["_validator"] = IPythonValidator(get_compiler_flags=self.get_compiler_flags)
182
183        super().__init__(*a, **kw)
184        self.ipython_shell = ipython_shell
185
186        self.all_prompt_styles["ipython"] = IPythonPrompt(ipython_shell.prompts)
187        self.prompt_style = "ipython"
188
189        # UI style for IPython. Add tokens that are used by IPython>5.0
190        style_dict = {}
191        style_dict.update(default_ui_style)
192        style_dict.update(
193            {
194                "pygments.prompt": "#009900",
195                "pygments.prompt-num": "#00ff00 bold",
196                "pygments.out-prompt": "#990000",
197                "pygments.out-prompt-num": "#ff0000 bold",
198            }
199        )
200
201        self.ui_styles = {"default": Style.from_dict(style_dict)}
202        self.use_ui_colorscheme("default")
203
204
205class InteractiveShellEmbed(_InteractiveShellEmbed):
206    """
207    Override the `InteractiveShellEmbed` from IPython, to replace the front-end
208    with our input shell.
209
210    :param configure: Callable for configuring the repl.
211    """
212
213    def __init__(self, *a, **kw):
214        vi_mode = kw.pop("vi_mode", False)
215        history_filename = kw.pop("history_filename", None)
216        configure = kw.pop("configure", None)
217        title = kw.pop("title", None)
218
219        # Don't ask IPython to confirm for exit. We have our own exit prompt.
220        self.confirm_exit = False
221
222        super().__init__(*a, **kw)
223
224        def get_globals():
225            return self.user_ns
226
227        python_input = IPythonInput(
228            self,
229            get_globals=get_globals,
230            vi_mode=vi_mode,
231            history_filename=history_filename,
232        )
233
234        if title:
235            python_input.terminal_title = title
236
237        if configure:
238            configure(python_input)
239            python_input.prompt_style = "ipython"  # Don't take from config.
240
241        self.python_input = python_input
242
243    def prompt_for_code(self):
244        try:
245            return self.python_input.app.run()
246        except KeyboardInterrupt:
247            self.python_input.default_buffer.document = Document()
248            return ""
249
250
251def initialize_extensions(shell, extensions):
252    """
253    Partial copy of `InteractiveShellApp.init_extensions` from IPython.
254    """
255    try:
256        iter(extensions)
257    except TypeError:
258        pass  # no extensions found
259    else:
260        for ext in extensions:
261            try:
262                shell.extension_manager.load_extension(ext)
263            except:
264                warn(
265                    "Error in loading extension: %s" % ext
266                    + "\nCheck your config files in %s"
267                    % ipy_utils.path.get_ipython_dir()
268                )
269                shell.showtraceback()
270
271
272def embed(**kwargs):
273    """
274    Copied from `IPython/terminal/embed.py`, but using our `InteractiveShellEmbed` instead.
275    """
276    config = kwargs.get("config")
277    header = kwargs.pop("header", "")
278    compile_flags = kwargs.pop("compile_flags", None)
279    if config is None:
280        config = load_default_config()
281        config.InteractiveShellEmbed = config.TerminalInteractiveShell
282        kwargs["config"] = config
283    shell = InteractiveShellEmbed.instance(**kwargs)
284    initialize_extensions(shell, config["InteractiveShellApp"]["extensions"])
285    run_startup_scripts(shell)
286    shell(header=header, stack_depth=2, compile_flags=compile_flags)
287
288
289def run_startup_scripts(shell):
290    """
291    Contributed by linyuxu:
292    https://github.com/prompt-toolkit/ptpython/issues/126#issue-161242480
293    """
294    import glob
295    import os
296
297    startup_dir = shell.profile_dir.startup_dir
298    startup_files = []
299    startup_files += glob.glob(os.path.join(startup_dir, "*.py"))
300    startup_files += glob.glob(os.path.join(startup_dir, "*.ipy"))
301    for file in startup_files:
302        shell.run_cell(open(file).read())
303