From 3554c6d80134b6048f09b9c175c893968b8ddbc2 Mon Sep 17 00:00:00 2001
From: kobarity <kobarity@gmail.com>
Date: Wed, 12 Jun 2024 01:09:21 +0900
Subject: [PATCH] Allow escape sequences in Python prompts

* lisp/progmodes/python.el (python-shell-prompt-detect): Use
Python's json package if available, and remove escape sequences
in prompts.
* test/lisp/progmodes/python-tests.el
(python-tests-interpreter-2-6-higher-p): New predicate
function.
(python-shell-prompt-detect-7): New test.  (Bug#71440)
---
 lisp/progmodes/python.el            |  9 +++++++--
 test/lisp/progmodes/python-tests.el | 28 ++++++++++++++++++++++++++++
 2 files changed, 35 insertions(+), 2 deletions(-)

diff --git a/lisp/progmodes/python.el b/lisp/progmodes/python.el
index 2664d71d358..ca5ecfab6ea 100644
--- a/lisp/progmodes/python.el
+++ b/lisp/progmodes/python.el
@@ -3104,8 +3104,13 @@ python-shell-prompt-detect
       (let* ((code (concat
                     "import sys\n"
                     "ps = [getattr(sys, 'ps%s' % i, '') for i in range(1,4)]\n"
+                    "try:\n"
+                    "    import json\n"
+                    "    ps_json = '\\n' + json.dumps(ps)\n"
+                    "except ImportError:\n"
                     ;; JSON is built manually for compatibility
-                    "ps_json = '\\n[\"%s\", \"%s\", \"%s\"]\\n' % tuple(ps)\n"
+                    "    ps_json = '\\n[\"%s\", \"%s\", \"%s\"]\\n' % tuple(ps)\n"
+                    "\n"
                     "print (ps_json)\n"
                     "sys.exit(0)\n"))
              (interpreter python-shell-interpreter)
@@ -3168,7 +3173,7 @@ python-shell-prompt-detect
             "Or alternatively in:\n"
             "  + `python-shell-prompt-input-regexps'\n"
             "  + `python-shell-prompt-output-regexps'")))
-        prompts))))
+        (mapcar #'ansi-color-filter-apply prompts)))))
 
 (defun python-shell-prompt-validate-regexps ()
   "Validate all user provided regexps for prompts.
diff --git a/test/lisp/progmodes/python-tests.el b/test/lisp/progmodes/python-tests.el
index ce103921454..04b9b448ec2 100644
--- a/test/lisp/progmodes/python-tests.el
+++ b/test/lisp/progmodes/python-tests.el
@@ -3820,6 +3820,17 @@ python-tests-interpreter-3-p
   (when (string= (car (split-string (cdr info) "\\.")) "3")
     (car info)))
 
+(defun python-tests-interpreter-2-6-higher-p (info)
+  "Check if the interpreter major version in INFO is 2.6 or higher.
+This function is intended to be used as the PRED argument of
+`python-tests-get-shell-interpreter'."
+  (let* ((version (split-string (cdr info) "\\."))
+         (major (string-to-number (car version)))
+         (minor (string-to-number (cadr version))))
+    (when (or (>= major 3)
+              (and (= major 2) (>= minor 6)))
+      (car info))))
+
 (ert-deftest python-shell-get-process-name-1 ()
   "Check process name calculation sans `buffer-file-name'."
   (python-tests-with-temp-buffer
@@ -4353,6 +4364,23 @@ python-shell-prompt-detect-6
            (should (not (get-buffer "*Warnings*"))))
        (ignore-errors (delete-file startup-file))))))
 
+(ert-deftest python-shell-prompt-detect-7 ()
+  "Check prompt autodetection with escape sequences.  Bug#71440."
+  (python-tests-with-shell-interpreter
+   #'python-tests-interpreter-2-6-higher-p
+   (let* ((process-environment process-environment)
+          (startup-code (concat "import sys\n"
+                                "sys.ps1 = '\033[32mpy> \033[0m'\n"
+                                "sys.ps2 = '\033[32m..> \033[0m'\n"
+                                "sys.ps3 = '\033[32mout \033[0m'\n"))
+          (startup-file (python-shell--save-temp-file startup-code)))
+     (unwind-protect
+         (progn
+           (setenv "PYTHONSTARTUP" startup-file)
+           (should python-shell-prompt-detect-enabled)
+           (should (equal (python-shell-prompt-detect) '("py> " "..> " "out "))))
+       (ignore-errors (delete-file startup-file))))))
+
 (ert-deftest python-shell-prompt-validate-regexps-1 ()
   "Check `python-shell-prompt-input-regexps' are validated."
   (let* ((python-shell-prompt-input-regexps '("\\("))
-- 
2.34.1

