1import ast
2import collections.abc as collections_abc
3import inspect
4import keyword
5import re
6from enum import Enum
7from typing import TYPE_CHECKING, Any, Callable, Dict, Iterable, List, Optional
8
9from prompt_toolkit.completion import (
10    CompleteEvent,
11    Completer,
12    Completion,
13    PathCompleter,
14)
15from prompt_toolkit.contrib.completers.system import SystemCompleter
16from prompt_toolkit.contrib.regular_languages.compiler import compile as compile_grammar
17from prompt_toolkit.contrib.regular_languages.completion import GrammarCompleter
18from prompt_toolkit.document import Document
19from prompt_toolkit.formatted_text import fragment_list_to_text, to_formatted_text
20
21from ptpython.utils import get_jedi_script_from_document
22
23if TYPE_CHECKING:
24    from prompt_toolkit.contrib.regular_languages.compiler import _CompiledGrammar
25
26__all__ = ["PythonCompleter", "CompletePrivateAttributes", "HidePrivateCompleter"]
27
28
29class CompletePrivateAttributes(Enum):
30    """
31    Should we display private attributes in the completion pop-up?
32    """
33
34    NEVER = "NEVER"
35    IF_NO_PUBLIC = "IF_NO_PUBLIC"
36    ALWAYS = "ALWAYS"
37
38
39class PythonCompleter(Completer):
40    """
41    Completer for Python code.
42    """
43
44    def __init__(
45        self,
46        get_globals: Callable[[], dict],
47        get_locals: Callable[[], dict],
48        enable_dictionary_completion: Callable[[], bool],
49    ) -> None:
50        super().__init__()
51
52        self.get_globals = get_globals
53        self.get_locals = get_locals
54        self.enable_dictionary_completion = enable_dictionary_completion
55
56        self._system_completer = SystemCompleter()
57        self._jedi_completer = JediCompleter(get_globals, get_locals)
58        self._dictionary_completer = DictionaryCompleter(get_globals, get_locals)
59
60        self._path_completer_cache: Optional[GrammarCompleter] = None
61        self._path_completer_grammar_cache: Optional["_CompiledGrammar"] = None
62
63    @property
64    def _path_completer(self) -> GrammarCompleter:
65        if self._path_completer_cache is None:
66            self._path_completer_cache = GrammarCompleter(
67                self._path_completer_grammar,
68                {
69                    "var1": PathCompleter(expanduser=True),
70                    "var2": PathCompleter(expanduser=True),
71                },
72            )
73        return self._path_completer_cache
74
75    @property
76    def _path_completer_grammar(self) -> "_CompiledGrammar":
77        """
78        Return the grammar for matching paths inside strings inside Python
79        code.
80        """
81        # We make this lazy, because it delays startup time a little bit.
82        # This way, the grammar is build during the first completion.
83        if self._path_completer_grammar_cache is None:
84            self._path_completer_grammar_cache = self._create_path_completer_grammar()
85        return self._path_completer_grammar_cache
86
87    def _create_path_completer_grammar(self) -> "_CompiledGrammar":
88        def unwrapper(text: str) -> str:
89            return re.sub(r"\\(.)", r"\1", text)
90
91        def single_quoted_wrapper(text: str) -> str:
92            return text.replace("\\", "\\\\").replace("'", "\\'")
93
94        def double_quoted_wrapper(text: str) -> str:
95            return text.replace("\\", "\\\\").replace('"', '\\"')
96
97        grammar = r"""
98                # Text before the current string.
99                (
100                    [^'"#]                                  |  # Not quoted characters.
101                    '''  ([^'\\]|'(?!')|''(?!')|\\.])*  ''' |  # Inside single quoted triple strings
102                    "" " ([^"\\]|"(?!")|""(?!^)|\\.])* "" " |  # Inside double quoted triple strings
103
104                    \#[^\n]*(\n|$)           |  # Comment.
105                    "(?!"") ([^"\\]|\\.)*"   |  # Inside double quoted strings.
106                    '(?!'') ([^'\\]|\\.)*'      # Inside single quoted strings.
107
108                        # Warning: The negative lookahead in the above two
109                        #          statements is important. If we drop that,
110                        #          then the regex will try to interpret every
111                        #          triple quoted string also as a single quoted
112                        #          string, making this exponentially expensive to
113                        #          execute!
114                )*
115                # The current string that we're completing.
116                (
117                    ' (?P<var1>([^\n'\\]|\\.)*) |  # Inside a single quoted string.
118                    " (?P<var2>([^\n"\\]|\\.)*)    # Inside a double quoted string.
119                )
120        """
121
122        return compile_grammar(
123            grammar,
124            escape_funcs={"var1": single_quoted_wrapper, "var2": double_quoted_wrapper},
125            unescape_funcs={"var1": unwrapper, "var2": unwrapper},
126        )
127
128    def _complete_path_while_typing(self, document: Document) -> bool:
129        char_before_cursor = document.char_before_cursor
130        return bool(
131            document.text
132            and (char_before_cursor.isalnum() or char_before_cursor in "/.~")
133        )
134
135    def _complete_python_while_typing(self, document: Document) -> bool:
136        """
137        When `complete_while_typing` is set, only return completions when this
138        returns `True`.
139        """
140        text = document.text_before_cursor  # .rstrip()
141        char_before_cursor = text[-1:]
142        return bool(
143            text and (char_before_cursor.isalnum() or char_before_cursor in "_.([,")
144        )
145
146    def get_completions(
147        self, document: Document, complete_event: CompleteEvent
148    ) -> Iterable[Completion]:
149        """
150        Get Python completions.
151        """
152        # If the input starts with an exclamation mark. Use the system completer.
153        if document.text.lstrip().startswith("!"):
154            yield from self._system_completer.get_completions(
155                Document(
156                    text=document.text[1:], cursor_position=document.cursor_position - 1
157                ),
158                complete_event,
159            )
160            return
161
162        # Do dictionary key completions.
163        if complete_event.completion_requested or self._complete_python_while_typing(
164            document
165        ):
166            if self.enable_dictionary_completion():
167                has_dict_completions = False
168                for c in self._dictionary_completer.get_completions(
169                    document, complete_event
170                ):
171                    if c.text not in "[.":
172                        # If we get the [ or . completion, still include the other
173                        # completions.
174                        has_dict_completions = True
175                    yield c
176                if has_dict_completions:
177                    return
178
179        # Do Path completions (if there were no dictionary completions).
180        if complete_event.completion_requested or self._complete_path_while_typing(
181            document
182        ):
183            yield from self._path_completer.get_completions(document, complete_event)
184
185        # Do Jedi completions.
186        if complete_event.completion_requested or self._complete_python_while_typing(
187            document
188        ):
189            # If we are inside a string, Don't do Jedi completion.
190            if not self._path_completer_grammar.match(document.text_before_cursor):
191
192                # Do Jedi Python completions.
193                yield from self._jedi_completer.get_completions(
194                    document, complete_event
195                )
196
197
198class JediCompleter(Completer):
199    """
200    Autocompleter that uses the Jedi library.
201    """
202
203    def __init__(self, get_globals, get_locals) -> None:
204        super().__init__()
205
206        self.get_globals = get_globals
207        self.get_locals = get_locals
208
209    def get_completions(
210        self, document: Document, complete_event: CompleteEvent
211    ) -> Iterable[Completion]:
212        script = get_jedi_script_from_document(
213            document, self.get_locals(), self.get_globals()
214        )
215
216        if script:
217            try:
218                jedi_completions = script.complete(
219                    column=document.cursor_position_col,
220                    line=document.cursor_position_row + 1,
221                )
222            except TypeError:
223                # Issue #9: bad syntax causes completions() to fail in jedi.
224                # https://github.com/jonathanslenders/python-prompt-toolkit/issues/9
225                pass
226            except UnicodeDecodeError:
227                # Issue #43: UnicodeDecodeError on OpenBSD
228                # https://github.com/jonathanslenders/python-prompt-toolkit/issues/43
229                pass
230            except AttributeError:
231                # Jedi issue #513: https://github.com/davidhalter/jedi/issues/513
232                pass
233            except ValueError:
234                # Jedi issue: "ValueError: invalid \x escape"
235                pass
236            except KeyError:
237                # Jedi issue: "KeyError: u'a_lambda'."
238                # https://github.com/jonathanslenders/ptpython/issues/89
239                pass
240            except IOError:
241                # Jedi issue: "IOError: No such file or directory."
242                # https://github.com/jonathanslenders/ptpython/issues/71
243                pass
244            except AssertionError:
245                # In jedi.parser.__init__.py: 227, in remove_last_newline,
246                # the assertion "newline.value.endswith('\n')" can fail.
247                pass
248            except SystemError:
249                # In jedi.api.helpers.py: 144, in get_stack_at_position
250                # raise SystemError("This really shouldn't happen. There's a bug in Jedi.")
251                pass
252            except NotImplementedError:
253                # See: https://github.com/jonathanslenders/ptpython/issues/223
254                pass
255            except Exception:
256                # Supress all other Jedi exceptions.
257                pass
258            else:
259                # Move function parameters to the top.
260                jedi_completions = sorted(
261                    jedi_completions,
262                    key=lambda jc: (
263                        # Params first.
264                        jc.type != "param",
265                        # Private at the end.
266                        jc.name.startswith("_"),
267                        # Then sort by name.
268                        jc.name_with_symbols.lower(),
269                    ),
270                )
271
272                for jc in jedi_completions:
273                    if jc.type == "function":
274                        suffix = "()"
275                    else:
276                        suffix = ""
277
278                    if jc.type == "param":
279                        suffix = "..."
280
281                    yield Completion(
282                        jc.name_with_symbols,
283                        len(jc.complete) - len(jc.name_with_symbols),
284                        display=jc.name_with_symbols + suffix,
285                        display_meta=jc.type,
286                        style=_get_style_for_jedi_completion(jc),
287                    )
288
289
290class DictionaryCompleter(Completer):
291    """
292    Experimental completer for Python dictionary keys.
293
294    Warning: This does an `eval` and `repr` on some Python expressions before
295             the cursor, which is potentially dangerous. It doesn't match on
296             function calls, so it only triggers attribute access.
297    """
298
299    def __init__(self, get_globals, get_locals):
300        super().__init__()
301
302        self.get_globals = get_globals
303        self.get_locals = get_locals
304
305        # Pattern for expressions that are "safe" to eval for auto-completion.
306        # These are expressions that contain only attribute and index lookups.
307        varname = r"[a-zA-Z_][a-zA-Z0-9_]*"
308
309        expression = rf"""
310            # Any expression safe enough to eval while typing.
311            # No operators, except dot, and only other dict lookups.
312            # Technically, this can be unsafe of course, if bad code runs
313            # in `__getattr__` or ``__getitem__``.
314            (
315                # Variable name
316                {varname}
317
318                \s*
319
320                (?:
321                    # Attribute access.
322                    \s* \. \s* {varname} \s*
323
324                    |
325
326                    # Item lookup.
327                    # (We match the square brackets. The key can be anything.
328                    # We don't care about matching quotes here in the regex.
329                    # Nested square brackets are not supported.)
330                    \s* \[ [^\[\]]+ \] \s*
331                )*
332            )
333        """
334
335        # Pattern for recognizing for-loops, so that we can provide
336        # autocompletion on the iterator of the for-loop. (According to the
337        # first item of the collection we're iterating over.)
338        self.for_loop_pattern = re.compile(
339            rf"""
340                for \s+ ([a-zA-Z0-9_]+) \s+ in \s+ {expression} \s* :
341            """,
342            re.VERBOSE,
343        )
344
345        # Pattern for matching a simple expression (for completing [ or .
346        # operators).
347        self.expression_pattern = re.compile(
348            rf"""
349                {expression}
350                $
351            """,
352            re.VERBOSE,
353        )
354
355        # Pattern for matching item lookups.
356        self.item_lookup_pattern = re.compile(
357            rf"""
358                {expression}
359
360                # Dict loopup to complete (square bracket open + start of
361                # string).
362                \[
363                \s* ([^\[\]]*)$
364            """,
365            re.VERBOSE,
366        )
367
368        # Pattern for matching attribute lookups.
369        self.attribute_lookup_pattern = re.compile(
370            rf"""
371                {expression}
372
373                # Attribute loopup to complete (dot + varname).
374                \.
375                \s* ([a-zA-Z0-9_]*)$
376            """,
377            re.VERBOSE,
378        )
379
380    def _lookup(self, expression: str, temp_locals: Dict[str, Any]) -> object:
381        """
382        Do lookup of `object_var` in the context.
383        `temp_locals` is a dictionary, used for the locals.
384        """
385        try:
386            return eval(expression.strip(), self.get_globals(), temp_locals)
387        except BaseException:
388            return None  # Many exception, like NameError can be thrown here.
389
390    def get_completions(
391        self, document: Document, complete_event: CompleteEvent
392    ) -> Iterable[Completion]:
393
394        # First, find all for-loops, and assign the first item of the
395        # collections they're iterating to the iterator variable, so that we
396        # can provide code completion on the iterators.
397        temp_locals = self.get_locals().copy()
398
399        for match in self.for_loop_pattern.finditer(document.text_before_cursor):
400            varname, expression = match.groups()
401            expression_val = self._lookup(expression, temp_locals)
402
403            # We do this only for lists and tuples. Calling `next()` on any
404            # collection would create undesired side effects.
405            if isinstance(expression_val, (list, tuple)) and expression_val:
406                temp_locals[varname] = expression_val[0]
407
408        # Get all completions.
409        yield from self._get_expression_completions(
410            document, complete_event, temp_locals
411        )
412        yield from self._get_item_lookup_completions(
413            document, complete_event, temp_locals
414        )
415        yield from self._get_attribute_completions(
416            document, complete_event, temp_locals
417        )
418
419    def _do_repr(self, obj: object) -> str:
420        try:
421            return str(repr(obj))
422        except BaseException:
423            raise ReprFailedError
424
425    def eval_expression(self, document: Document, locals: Dict[str, Any]) -> object:
426        """
427        Evaluate
428        """
429        match = self.expression_pattern.search(document.text_before_cursor)
430        if match is not None:
431            object_var = match.groups()[0]
432            return self._lookup(object_var, locals)
433
434        return None
435
436    def _get_expression_completions(
437        self,
438        document: Document,
439        complete_event: CompleteEvent,
440        temp_locals: Dict[str, Any],
441    ) -> Iterable[Completion]:
442        """
443        Complete the [ or . operator after an object.
444        """
445        result = self.eval_expression(document, temp_locals)
446
447        if result is not None:
448
449            if isinstance(
450                result,
451                (list, tuple, dict, collections_abc.Mapping, collections_abc.Sequence),
452            ):
453                yield Completion("[", 0)
454
455            else:
456                # Note: Don't call `if result` here. That can fail for types
457                #       that have custom truthness checks.
458                yield Completion(".", 0)
459
460    def _get_item_lookup_completions(
461        self,
462        document: Document,
463        complete_event: CompleteEvent,
464        temp_locals: Dict[str, Any],
465    ) -> Iterable[Completion]:
466        """
467        Complete dictionary keys.
468        """
469
470        def abbr_meta(text: str) -> str:
471            "Abbreviate meta text, make sure it fits on one line."
472            # Take first line, if multiple lines.
473            if len(text) > 20:
474                text = text[:20] + "..."
475            if "\n" in text:
476                text = text.split("\n", 1)[0] + "..."
477            return text
478
479        match = self.item_lookup_pattern.search(document.text_before_cursor)
480        if match is not None:
481            object_var, key = match.groups()
482
483            # Do lookup of `object_var` in the context.
484            result = self._lookup(object_var, temp_locals)
485
486            # If this object is a dictionary, complete the keys.
487            if isinstance(result, (dict, collections_abc.Mapping)):
488                # Try to evaluate the key.
489                key_obj = key
490                for k in [key, key + '"', key + "'"]:
491                    try:
492                        key_obj = ast.literal_eval(k)
493                    except (SyntaxError, ValueError):
494                        continue
495                    else:
496                        break
497
498                for k in result:
499                    if str(k).startswith(str(key_obj)):
500                        try:
501                            k_repr = self._do_repr(k)
502                            yield Completion(
503                                k_repr + "]",
504                                -len(key),
505                                display=f"[{k_repr}]",
506                                display_meta=abbr_meta(self._do_repr(result[k])),
507                            )
508                        except KeyError:
509                            # `result[k]` lookup failed. Trying to complete
510                            # broken object.
511                            pass
512                        except ReprFailedError:
513                            pass
514
515            # Complete list/tuple index keys.
516            elif isinstance(result, (list, tuple, collections_abc.Sequence)):
517                if not key or key.isdigit():
518                    for k in range(min(len(result), 1000)):
519                        if str(k).startswith(key):
520                            try:
521                                k_repr = self._do_repr(k)
522                                yield Completion(
523                                    k_repr + "]",
524                                    -len(key),
525                                    display=f"[{k_repr}]",
526                                    display_meta=abbr_meta(self._do_repr(result[k])),
527                                )
528                            except KeyError:
529                                # `result[k]` lookup failed. Trying to complete
530                                # broken object.
531                                pass
532                            except ReprFailedError:
533                                pass
534
535    def _get_attribute_completions(
536        self,
537        document: Document,
538        complete_event: CompleteEvent,
539        temp_locals: Dict[str, Any],
540    ) -> Iterable[Completion]:
541        """
542        Complete attribute names.
543        """
544        match = self.attribute_lookup_pattern.search(document.text_before_cursor)
545        if match is not None:
546            object_var, attr_name = match.groups()
547
548            # Do lookup of `object_var` in the context.
549            result = self._lookup(object_var, temp_locals)
550
551            names = self._sort_attribute_names(dir(result))
552
553            def get_suffix(name: str) -> str:
554                try:
555                    obj = getattr(result, name, None)
556                    if inspect.isfunction(obj) or inspect.ismethod(obj):
557                        return "()"
558                    if isinstance(obj, dict):
559                        return "{}"
560                    if isinstance(obj, (list, tuple)):
561                        return "[]"
562                except:
563                    pass
564                return ""
565
566            for name in names:
567                if name.startswith(attr_name):
568                    suffix = get_suffix(name)
569                    yield Completion(name, -len(attr_name), display=name + suffix)
570
571    def _sort_attribute_names(self, names: List[str]) -> List[str]:
572        """
573        Sort attribute names alphabetically, but move the double underscore and
574        underscore names to the end.
575        """
576
577        def sort_key(name: str):
578            if name.startswith("__"):
579                return (2, name)  # Double underscore comes latest.
580            if name.startswith("_"):
581                return (1, name)  # Single underscore before that.
582            return (0, name)  # Other names first.
583
584        return sorted(names, key=sort_key)
585
586
587class HidePrivateCompleter(Completer):
588    """
589    Wrapper around completer that hides private fields, deponding on whether or
590    not public fields are shown.
591
592    (The reason this is implemented as a `Completer` wrapper is because this
593    way it works also with `FuzzyCompleter`.)
594    """
595
596    def __init__(
597        self,
598        completer: Completer,
599        complete_private_attributes: Callable[[], CompletePrivateAttributes],
600    ) -> None:
601        self.completer = completer
602        self.complete_private_attributes = complete_private_attributes
603
604    def get_completions(
605        self, document: Document, complete_event: CompleteEvent
606    ) -> Iterable[Completion]:
607
608        completions = list(self.completer.get_completions(document, complete_event))
609        complete_private_attributes = self.complete_private_attributes()
610        hide_private = False
611
612        def is_private(completion: Completion) -> bool:
613            text = fragment_list_to_text(to_formatted_text(completion.display))
614            return text.startswith("_")
615
616        if complete_private_attributes == CompletePrivateAttributes.NEVER:
617            hide_private = True
618
619        elif complete_private_attributes == CompletePrivateAttributes.IF_NO_PUBLIC:
620            hide_private = any(not is_private(completion) for completion in completions)
621
622        if hide_private:
623            completions = [
624                completion for completion in completions if not is_private(completion)
625            ]
626
627        return completions
628
629
630class ReprFailedError(Exception):
631    "Raised when the repr() call in `DictionaryCompleter` fails."
632
633
634try:
635    import builtins
636
637    _builtin_names = dir(builtins)
638except ImportError:  # Python 2.
639    _builtin_names = []
640
641
642def _get_style_for_jedi_completion(jedi_completion) -> str:
643    """
644    Return completion style to use for this name.
645    """
646    name = jedi_completion.name_with_symbols
647
648    if jedi_completion.type == "param":
649        return "class:completion.param"
650
651    if name in _builtin_names:
652        return "class:completion.builtin"
653
654    if keyword.iskeyword(name):
655        return "class:completion.keyword"
656
657    return ""
658