Subject: bug#68559: [PATCH] Improve Python shell completion

From: Liu Hui <[email-masked]>
Date: Fri, 26 Jan 2024 18:12:19 +0800

kobarity <[email-masked]> 于2024年1月25日周四 23:38写道：
>
> Liu Hui wrote:
> > kobarity <[email-masked]> 于2024年1月23日周二 22:15写道：
> > > Thank you.  I tried the new patch and confirmed that the Python Shell
> > > buffer allows keyword argument completion on multi-line function
> > > calls.  However, my expectation is that it can be done in Python
> > > buffer as well (after calling `python-shell-send-buffer').  Is this
> > > also possible?
> > Thank you for the feedback. I have updated the patch to add support in
> > Python buffer.
>
> Thank you.  It's working very nice.
>
> I apologize for reporting in the piecemeal fashion, but I have noticed
> the following differences from the current completion.

No worries. I appreciate your valuable feedback for improving the
quality of the patch.

> 1. Module names cannot be completed in Python buffer.

Can you elaborate? The completion of module names never worked for me
except using IPython with non-native completion mechanism.

Because module names cannot be completed using vanilla python
interpreter in the terminal, they are also not completed even with the
patch, unless jedi is available.

> 2. Completion is not working in comments.
>
> I'm not sure if the completion should work in comments, but some
> people might want it because the current python-mode allows it in both
> Python buffer and Python Shell buffer.

Thanks for pointing it out. The attached patch should now restore
previous basic completion behavior in comments.
From 770b306fc9ac2406517536e5276649578a18955a Mon Sep 17 00:00:00 2001
From: Liu Hui <[email-masked]>
Date: Thu, 18 Jan 2024 12:00:00 +0800
Subject: [PATCH] Improve Python shell completion (bug#68559)

* lisp/progmodes/python.el (python-shell-completion-setup-code): Fix
the completion code of IPython.  Change the return value to JSON
string and ...
(python-shell-completion-get-completions): ... simplify parsing.
(inferior-python-mode): Update docstring.
(python-shell-readline-completer): New option.
(python-shell-readline-completer-delims): New variable indicating the
word delimiters of readline completer.
(python-shell--readline-jedi-setup-code)
(python-shell--readline-ipython-setup-code): New internal variables.
(python-shell-completion-native-setup): Setup a suitable readline
completer and set the completer delimiter.
(python-shell-completion-native-get-completions): Convert output
string to completions properly.
(python-shell--get-multiline-input)
(python-shell--extra-completion-context)
(python-shell-completion-extra-context): New functions.
(python-shell-completion-at-point): Send text beginning from the line
start if the completion backend does not need word splitting.  Remove
the detection of import statement because it is not needed anymore.
Create proper completion table based on completions returned from
different backends.

* test/lisp/progmodes/python-tests.el (python-tests--completion-module):
(python-tests--completion-parameters):
(python-tests--completion-extra-context): New helper functions.
(python-shell-completion-at-point-jedi-completer):
(python-shell-completion-at-point-ipython): New tests.

* etc/NEWS: Announce the change.
---
 etc/NEWS                            |   6 +
 lisp/progmodes/python.el            | 309 +++++++++++++++++++++++-----
 test/lisp/progmodes/python-tests.el |  78 +++++++
 3 files changed, 341 insertions(+), 52 deletions(-)

diff --git a/etc/NEWS b/etc/NEWS
index a1874313502..9f05bdf4705 100644
--- a/etc/NEWS
+++ b/etc/NEWS
@@ -1047,6 +1047,12 @@ instead of:
 This allows the user to specify command line arguments to the non
 interactive Python interpreter specified by 'python-interpreter'.
 
+*** New user option 'python-shell-readline-completer'.
+This allows the user to specify the readline completer used for Python
+shell completion.  The default is 'auto', which means a suitable
+completer will be configured automatically according to the Python
+interpreter.
+
 ** use-package
 
 +++
diff --git a/lisp/progmodes/python.el b/lisp/progmodes/python.el
index ff799e1e662..8e7998ee8f6 100644
--- a/lisp/progmodes/python.el
+++ b/lisp/progmodes/python.el
@@ -128,9 +128,9 @@
 ;; receiving escape sequences (with some limitations, i.e. completion
 ;; in blocks does not work).  The code executed for the "fallback"
 ;; completion can be found in `python-shell-completion-setup-code' and
-;; `python-shell-completion-string-code' variables.  Their default
-;; values enable completion for both CPython and IPython, and probably
-;; any readline based shell (it's known to work with PyPy).  If your
+;; `python-shell-completion-get-completions'.  Their default values
+;; enable completion for both CPython and IPython, and probably any
+;; readline based shell (it's known to work with PyPy).  If your
 ;; Python installation lacks readline (like CPython for Windows),
 ;; installing pyreadline (URL `https://ipython.org/pyreadline.html')
 ;; should suffice.  To troubleshoot why you are not getting any
@@ -3601,7 +3601,6 @@ inferior-python-mode
 `python-shell-prompt-block-regexp',
 `python-shell-font-lock-enable',
 `python-shell-completion-setup-code',
-`python-shell-completion-string-code',
 `python-eldoc-setup-code',
 `python-ffap-setup-code' can
 customize this mode for different Python interpreters.
@@ -4319,8 +4318,9 @@ python-shell-completion-setup-code
     completions = []
     completer = None
 
+    import json
     try:
-        import readline
+        import readline, re
 
         try:
             import __builtin__
@@ -4331,16 +4331,29 @@ python-shell-completion-setup-code
 
         is_ipython = ('__IPYTHON__' in builtins or
                       '__IPYTHON__active' in builtins)
-        splits = text.split()
-        is_module = splits and splits[0] in ('from', 'import')
-
-        if is_ipython and is_module:
-            from IPython.core.completerlib import module_completion
-            completions = module_completion(text.strip())
-        elif is_ipython and '__IP' in builtins:
-            completions = __IP.complete(text)
-        elif is_ipython and 'get_ipython' in builtins:
-            completions = get_ipython().Completer.all_completions(text)
+
+        if is_ipython and 'get_ipython' in builtins:
+            def filter_c(prefix, c):
+                if re.match('_+(i?[0-9]+)?$', c):
+                    return False
+                elif c[0] == '%' and not re.match('[%a-zA-Z]+$', prefix):
+                    return False
+                return True
+
+            import IPython
+            try:
+                if IPython.version_info[0] >= 6:
+                    from IPython.core.completer import provisionalcompleter
+                    with provisionalcompleter():
+                        completions = [
+                            [c.text, c.start, c.end, c.type or '?', c.signature or '']
+                             for c in get_ipython().Completer.completions(text, len(text))
+                             if filter_c(text, c.text)]
+                else:
+                    part, matches = get_ipython().Completer.complete(line_buffer=text)
+                    completions = [text + m[len(part):] for m in matches if filter_c(text, m)]
+            except:
+                pass
         else:
             # Try to reuse current completer.
             completer = readline.get_completer()
@@ -4363,7 +4376,7 @@ python-shell-completion-setup-code
     finally:
         if getattr(completer, 'PYTHON_EL_WRAPPED', False):
             completer.print_mode = True
-    return completions"
+    return json.dumps(completions)"
   "Code used to setup completion in inferior Python processes."
   :type 'string)
 
@@ -4404,6 +4417,78 @@ python-shell-completion-native-try-output-timeout
   :version "25.1"
   :type 'float)
 
+(defcustom python-shell-readline-completer 'auto
+  "The readline completer used for Python shell completion.
+If the value is non-nil, Python shell will setup the readline
+completer unless it has been set elsewhere (e.g. in the
+PYTHONSTARTUP file).  Below are possible values:
+- `auto': the completer is determined according to the
+interpreter.  Specifically, the IPython completer, defined in
+`python-shell--readline-ipython-setup-code', is used when the
+interpreter is ipython, otherwise the Jedi completer is used.
+- a string: Python code to setup the readline.  See
+`python-shell--readline-jedi-setup-code' for reference.
+- `nil': Python shell will do nothing.
+
+In any case, if the completer is not set successfully in the end,
+fallback to the built-in rlcompleter."
+  :type '(choice (const  :tag "Automatic" auto)
+                 (const  :tag "No configuration" nil)
+                 (string :tag "Python setup code"))
+  :version "30.1")
+
+(defvar python-shell-readline-completer-delims nil
+  "Word delimiters used by the readline completer.
+It is automatically set by Python shell.")
+
+(defconst python-shell--readline-jedi-setup-code
+  "
+def __PYTHON_EL_setup_readline_completer():
+    from jedi.utils import setup_readline
+    setup_readline()"
+  "Code used to setup readline completer with Jedi.")
+
+(defconst python-shell--readline-ipython-setup-code
+  "
+def __PYTHON_EL_setup_readline_completer():
+    import readline, re, json, IPython
+
+    class __ipython_RL:
+        def __init__(self, v):
+            self.version = v
+
+        def filter(self, prefix, c):
+            if re.match('_+(i?[0-9]+)?$', c):
+                return False
+            elif c[0] == '%' and not re.match('[%a-zA-Z]+$', prefix):
+                return False
+            return True
+
+        def complete(self, text, state):
+            if state == 0:
+                try:
+                    if self.version >= 6:
+                        from IPython.core.completer import provisionalcompleter
+                        with provisionalcompleter():
+                            self.matches = [json.dumps([
+                                [c.text, c.start, c.end, c.type or '?', c.signature or '']
+                                for c in get_ipython().Completer.completions(text, len(text))
+                                if self.filter(text, c.text)])]
+                    else:
+                        part, matches = get_ipython().Completer.complete(line_buffer=text)
+                        self.matches = [text + m[len(part):] for m in matches
+                                        if self.filter(text, m)]
+                except Exception:
+                    pass
+            try:
+                return self.matches[state]
+            except IndexError:
+                return None
+
+    readline.set_completer(__ipython_RL(IPython.version_info[0]).complete)
+    readline.set_completer_delims('')"
+  "Code used to setup readline completer for IPython.")
+
 (defvar python-shell-completion-native-redirect-buffer
   " *Python completions redirect*"
   "Buffer to be used to redirect output of readline commands.")
@@ -4427,7 +4512,20 @@ python-shell-completion-native-try
 (defun python-shell-completion-native-setup ()
   "Try to setup native completion, return non-nil on success."
   (let* ((process (python-shell-get-process))
-         (output (python-shell-send-string-no-output "
+         (completer (pcase python-shell-readline-completer
+                      ('auto
+                       (if (string-match-p "ipython[23]?\\'" python-shell-interpreter)
+                           python-shell--readline-ipython-setup-code
+                         python-shell--readline-jedi-setup-code))
+                      ((pred stringp) python-shell-readline-completer)
+                      (_ "")))
+         (output (python-shell-send-string-no-output
+                  (concat "
+try:
+    del __PYTHON_EL_setup_readline_completer
+except:
+    pass
+" completer "
 def __PYTHON_EL_native_completion_setup():
     try:
         import readline
@@ -4500,11 +4598,23 @@ python-shell-completion-native-setup
                 else:
                     return completion
 
+        def is_rlcompleter(completer):
+            try:
+                if completer.__self__.__module__ == 'rlcompleter':
+                    return True
+                else:
+                    return False
+            except Exception:
+                return False
+
         completer = readline.get_completer()
 
-        if not completer:
-            # Used as last resort to avoid breaking customizations.
-            import rlcompleter
+        if not completer or is_rlcompleter(completer):
+            try:
+                __PYTHON_EL_setup_readline_completer()
+            except:
+                # Used as last resort to avoid breaking customizations.
+                import rlcompleter
             completer = readline.get_completer()
 
         if completer and not getattr(completer, 'PYTHON_EL_WRAPPED', False):
@@ -4539,9 +4649,13 @@ python-shell-completion-native-setup
         print ('python.el: native completion setup failed, %s: %s'
                % sys.exc_info()[:2])
 
-__PYTHON_EL_native_completion_setup()" process)))
+__PYTHON_EL_native_completion_setup()") process)))
     (when (string-match-p "python\\.el: native completion setup loaded"
                           output)
+      (setq-local python-shell-readline-completer-delims
+                  (string-trim-right
+                   (python-shell-send-string-no-output
+                    "import readline; print(readline.get_completer_delims())")))
       (python-shell-completion-native-try))))
 
 (defun python-shell-completion-native-turn-off (&optional msg)
@@ -4609,6 +4723,8 @@ python-shell-completion-native-get-completions
     (let* ((original-filter-fn (process-filter process))
            (redirect-buffer (get-buffer-create
                              python-shell-completion-native-redirect-buffer))
+           (sep (if (string= python-shell-readline-completer-delims "")
+                    "[\n\r]+" "[ \f\t\n\r\v()]+"))
            (trigger "\t")
            (new-input (concat input trigger))
            (input-length
@@ -4651,28 +4767,72 @@ python-shell-completion-native-get-completions
                      process python-shell-completion-native-output-timeout
                      comint-redirect-finished-regexp)
                 (re-search-backward "0__dummy_completion__" nil t)
-                (cl-remove-duplicates
-                 (split-string
-                  (buffer-substring-no-properties
-                   (line-beginning-position) (point-min))
-                  "[ \f\t\n\r\v()]+" t)
-                 :test #'string=))))
+                (let ((str (buffer-substring-no-properties
+                            (line-beginning-position) (point-min))))
+                  (if (string= "[" (substring str 0 1))
+                      (condition-case nil
+                          (python--parse-json-array str)
+                        (t (cl-remove-duplicates (split-string str sep t)
+                                                 :test #'string=)))
+                    (cl-remove-duplicates (split-string str sep t)
+                                          :test #'string=))))))
         (set-process-filter process original-filter-fn)))))
 
 (defun python-shell-completion-get-completions (process input)
   "Get completions of INPUT using PROCESS."
   (with-current-buffer (process-buffer process)
-    (let ((completions
-           (python-util-strip-string
-            (python-shell-send-string-no-output
-             (format
-              "%s\nprint(';'.join(__PYTHON_EL_get_completions(%s)))"
+    (python--parse-json-array
+     (python-shell-send-string-no-output
+      (format "%s\nprint(__PYTHON_EL_get_completions(%s))"
               python-shell-completion-setup-code
               (python-shell--encode-string input))
-             process))))
-      (when (> (length completions) 2)
-        (split-string completions
-                      "^'\\|^\"\\|;\\|'$\\|\"$" t)))))
+      process))))
+
+(defun python-shell--get-multiline-input ()
+  "Return lines at a multi-line input in Python shell."
+  (save-excursion
+    (let ((p (point)) lines)
+      (when (progn
+              (beginning-of-line)
+              (looking-back python-shell-prompt-block-regexp (pos-bol)))
+        (push (buffer-substring-no-properties (point) p) lines)
+        (while (progn (comint-previous-prompt 1)
+                      (looking-back python-shell-prompt-block-regexp (pos-bol)))
+          (push (buffer-substring-no-properties (point) (pos-eol)) lines))
+        (push (buffer-substring-no-properties (point) (pos-eol)) lines))
+      lines)))
+
+(defun python-shell--extra-completion-context ()
+  "Get extra completion context of current input in Python shell."
+  (let ((lines (python-shell--get-multiline-input))
+        (python-indent-guess-indent-offset nil))
+    (when (not (zerop (length lines)))
+      (with-temp-buffer
+        (delay-mode-hooks
+          (insert (string-join lines "\n"))
+          (python-mode)
+          (python-shell-completion-extra-context))))))
+
+(defun python-shell-completion-extra-context (&optional pos)
+  "Get extra completion context at position POS in Python buffer.
+If optional argument POS is nil, use current position.
+
+Readline completers could use current line as the completion
+context, which may be insufficient.  In this function, extra
+context (e.g. multi-line function call) is found and reformatted
+as one line, which is required by native completion."
+  (let (bound p)
+    (save-excursion
+      (and pos (goto-char pos))
+      (setq bound (pos-bol))
+      (python-nav-up-list -1)
+      (when (and (< (point) bound)
+                 (looking-back
+                  (python-rx (+ (or "." symbol-name))) (pos-bol) t))
+        (setq p (match-beginning 0))))
+    (when p
+      (replace-regexp-in-string
+       "\n[ \t]*" "" (buffer-substring-no-properties p (1- bound))))))
 
 (defvar-local python-shell--capf-cache nil
   "Variable to store cached completions and invalidation keys.")
@@ -4687,21 +4847,26 @@ python-shell-completion-at-point
                          ;; Working on a shell buffer: use prompt end.
                          (cdr (python-util-comint-last-prompt))
                        (line-beginning-position)))
-         (import-statement
-          (when (string-match-p
-                 (rx (* space) word-start (or "from" "import") word-end space)
-                 (buffer-substring-no-properties line-start (point)))
-            (buffer-substring-no-properties line-start (point))))
+         (no-delims
+          (and (not (if is-shell-buffer
+                        (eq 'font-lock-comment-face
+                            (get-text-property (1- (point)) 'face))
+                      (python-syntax-context 'comment)))
+               (with-current-buffer (process-buffer process)
+                 (if python-shell-completion-native-enable
+                     (string= python-shell-readline-completer-delims "")
+                   (string-match-p "ipython[23]?\\'" python-shell-interpreter)))))
          (start
           (if (< (point) line-start)
               (point)
             (save-excursion
-              (if (not (re-search-backward
-                        (python-rx
-                         (or whitespace open-paren close-paren
-                             string-delimiter simple-operator))
-                        line-start
-                        t 1))
+              (if (or no-delims
+                      (not (re-search-backward
+                            (python-rx
+                             (or whitespace open-paren close-paren
+                                 string-delimiter simple-operator))
+                            line-start
+                            t 1)))
                   line-start
                 (forward-char (length (match-string-no-properties 0)))
                 (point)))))
@@ -4741,18 +4906,58 @@ python-shell-completion-at-point
                   (t #'python-shell-completion-native-get-completions))))
          (prev-prompt (car python-shell--capf-cache))
          (re (or (cadr python-shell--capf-cache) regexp-unmatchable))
-         (prefix (buffer-substring-no-properties start end)))
+         (prefix (buffer-substring-no-properties start end))
+         (prefix-offset 0)
+         (extra-context (when no-delims
+                          (if is-shell-buffer
+                              (python-shell--extra-completion-context)
+                            (python-shell-completion-extra-context))))
+         (extra-offset (length extra-context)))
+    (unless (zerop extra-offset)
+      (setq prefix (concat extra-context prefix)))
     ;; To invalidate the cache, we check if the prompt position or the
     ;; completion prefix changed.
     (unless (and (equal prev-prompt (car prompt-boundaries))
-                 (string-match re prefix))
+                 (string-match re prefix)
+                 (setq prefix-offset (- (length prefix) (match-end 1))))
       (setq python-shell--capf-cache
             `(,(car prompt-boundaries)
               ,(if (string-empty-p prefix)
                    regexp-unmatchable
-                 (concat "\\`" (regexp-quote prefix) "\\(?:\\sw\\|\\s_\\)*\\'"))
-              ,@(funcall completion-fn process (or import-statement prefix)))))
-    (list start end (cddr python-shell--capf-cache))))
+                 (concat "\\`\\(" (regexp-quote prefix) "\\)\\(?:\\sw\\|\\s_\\)*\\'"))
+              ,@(funcall completion-fn process prefix))))
+    (let ((cands (cddr python-shell--capf-cache)))
+      (cond
+       ((stringp (car cands))
+        (if no-delims
+            ;; Reduce completion candidates due to long prefix.
+            (if-let ((Lp (length prefix))
+                     ((string-match "\\(\\sw\\|\\s_\\)+\\'" prefix))
+                     (L (match-beginning 0)))
+                ;; If extra-offset is not zero:
+                ;;                  start              end
+                ;; o------------------o---------o-------o
+                ;; |<- extra-offset ->|
+                ;; |<----------- L ------------>|
+                ;;                          new-start
+                (list (+ start L (- extra-offset)) end
+                      (mapcar (lambda (s) (substring s L)) cands))
+              (list end end (mapcar (lambda (s) (substring s Lp)) cands)))
+          (list start end cands)))
+       ;; python-shell-completion(-native)-get-completions may produce
+       ;; a list of (text start end type signature) for completion.
+       ;; See `python-shell--readline-ipython-setup-code' and
+       ;; `python-shell-completion-setup-code'.
+       ((consp (car cands))
+        (list (+ start (nth 1 (car cands)) (- extra-offset))
+              ;; Candidates may be cached, so the end position should
+              ;; be adjusted according to current completion prefix.
+              (+ start (nth 2 (car cands)) (- extra-offset) prefix-offset)
+              cands
+              :annotation-function
+              (lambda (c) (concat " " (nth 3 (assoc c cands))))
+              :company-docsig
+              (lambda (c) (nth 4 (assoc c cands)))))))))
 
 (define-obsolete-function-alias
   'python-shell-completion-complete-at-point
diff --git a/test/lisp/progmodes/python-tests.el b/test/lisp/progmodes/python-tests.el
index 97ffd5fe20f..ae23ff8ebe2 100644
--- a/test/lisp/progmodes/python-tests.el
+++ b/test/lisp/progmodes/python-tests.el
@@ -4787,6 +4787,84 @@ python-shell-completion-at-point-native-1
      (end-of-line 0)
      (should-not (nth 2 (python-shell-completion-at-point))))))
 
+(defun python-tests--completion-module ()
+  "Check if modules can be completed in Python shell."
+  (insert "import datet")
+  (completion-at-point)
+  (beginning-of-line)
+  (should (looking-at-p "import datetime"))
+  (kill-line)
+  (insert "from datet")
+  (completion-at-point)
+  (beginning-of-line)
+  (should (looking-at-p "from datetime"))
+  (end-of-line)
+  (insert " import timed")
+  (completion-at-point)
+  (beginning-of-line)
+  (should (looking-at-p "from datetime import timedelta"))
+  (kill-line))
+
+(defun python-tests--completion-parameters ()
+  "Check if parameters can be completed in Python shell."
+  (insert "import re")
+  (comint-send-input)
+  (python-tests-shell-wait-for-prompt)
+  (insert "re.split('b', 'abc', maxs")
+  (completion-at-point)
+  (should (string= "re.split('b', 'abc', maxsplit="
+                   (buffer-substring (line-beginning-position) (point))))
+  (insert "0, ")
+  (should (python-shell-completion-at-point))
+  ;; Test if cache is used.
+  (cl-letf (((symbol-function 'python-shell-completion-get-completions)
+             'ignore)
+            ((symbol-function 'python-shell-completion-native-get-completions)
+             'ignore))
+    (insert "fla")
+    (completion-at-point)
+    (should (string= "re.split('b', 'abc', maxsplit=0, flags="
+                     (buffer-substring (line-beginning-position) (point)))))
+  (beginning-of-line)
+  (kill-line))
+
+(defun python-tests--completion-extra-context ()
+  "Check if extra context is used for completion."
+  (insert "re.split('b', 'abc',")
+  (comint-send-input)
+  (python-tests-shell-wait-for-prompt)
+  (insert "maxs")
+  (completion-at-point)
+  (should (string= "maxsplit="
+                   (buffer-substring (line-beginning-position) (point)))))
+
+(ert-deftest python-shell-completion-at-point-jedi-completer ()
+  "Check if Python shell completion works with Jedi."
+  (skip-unless (executable-find python-tests-shell-interpreter))
+  (python-tests-with-temp-buffer-with-shell
+   ""
+   (python-shell-with-shell-buffer
+     (python-shell-completion-native-turn-on)
+     (skip-unless (string= python-shell-readline-completer-delims ""))
+     (python-tests--completion-module)
+     (python-tests--completion-parameters)
+     (python-tests--completion-extra-context))))
+
+(ert-deftest python-shell-completion-at-point-ipython ()
+  "Check if Python shell completion works for IPython."
+  (let ((python-shell-interpreter "ipython")
+        (python-shell-interpreter-args "-i --simple-prompt"))
+    (skip-unless (executable-find python-shell-interpreter))
+    (python-tests-with-temp-buffer-with-shell
+     ""
+     (python-shell-with-shell-buffer
+       (python-shell-completion-native-turn-off)
+       (python-tests--completion-module)
+       (python-tests--completion-parameters)
+       (python-shell-completion-native-turn-on)
+       (python-tests--completion-module)
+       (python-tests--completion-parameters)
+       (python-tests--completion-extra-context)))))
 
 
 ;;; PDB Track integration
-- 
2.25.1