Fix PEX repl prompt for Linux PBS libedit. (#2503)

Previously the ansi terminal escape sequences were being displayed raw
instead of being interpreted as color codes by the terminal. This was
noticed when using a PEX scie since the Linux PBS CPython builds use a
an un-patched version of libedit that does not handle ansi terminal 
escape sequences unless finessed like we do now.

CHANGES.md

@@ -1,5 +1,14 @@
 # Release Notes
 
+## 2.16.1
+
+This release fixes the PEX repl for [Python Standalone Builds][PBS]
+Linux CPython PEX scies. These PEXes ship using a version of libedit
+for readline support that does not support naive use of ansi terminal
+escape sequences for prompt colorization.
+
+* Fix PEX repl prompt for Linux PBS libedit. (#2503)
+
 ## 2.16.0
 
 This release adds support for `--venv-system-site-packages` when

pex/repl/custom.py

@@ -22,6 +22,10 @@ def _try_enable_readline(
     history=False,  # type: bool
     history_file=None,  # type: Optional[str]
 ):
+    # type: (...) -> bool
+
+    libedit = False
+
     try:
         import readline
     except ImportError:
@@ -41,6 +45,7 @@ def _try_enable_readline(
         if "libedit" in readline.__doc__:
             # Mac can use libedit, and libedit has different config syntax.
             readline.parse_and_bind("bind ^I rl_complete")
+            libedit = True
         else:
             readline.parse_and_bind("tab: complete")
 
@@ -70,6 +75,8 @@ def _try_enable_readline(
 
             atexit.register(readline.write_history_file, histfile)
 
+    return libedit
+
 
 def repl_loop(
     banner=None,  # type: Optional[str]
@@ -78,10 +85,11 @@ def repl_loop(
     custom_commands=None,  # type: Optional[Mapping[str, Tuple[Callable, str]]]
     history=False,  # type: bool
     history_file=None,  # type: Optional[str]
+    use_libedit_color_prompt_workaround=False,  # type: bool
 ):
     # type: (...) -> Callable[[], Dict[str, Any]]
 
-    _try_enable_readline(history=history, history_file=history_file)
+    libedit = _try_enable_readline(history=history, history_file=history_file)
 
     _custom_commands = custom_commands or {}
 
@@ -108,14 +116,38 @@ def raw_input(self, prompt=""):
     repl = CustomREPL(locals=local)
     extra_args = {"exitmsg": ""} if sys.version_info[:2] >= (3, 6) else {}
 
-    repl_banner = _ANSI_RE.sub("", banner) if banner and not sys.stdout.isatty() else banner
+    def fixup_ansi(
+        text,  # type: str
+        prompt=False,  # type: bool
+    ):
+        # type: (...) -> str
+
+        if not sys.stdout.isatty():
+            text = _ANSI_RE.sub("", text)
+        elif prompt and libedit and use_libedit_color_prompt_workaround:
+            # Most versions of libedit do not support ansi terminal escape sequences, but they do
+            # support a readline back door where you can bracket characters that should be ignored
+            # for prompt width calculation but otherwise passed through unaltered to the terminal.
+            # This back door is defined in readline.h, which libedit implements, with the
+            # RL_PROMPT_START_IGNORE and RL_PROMPT_END_IGNORE character constants which we use
+            # below to bracket ansi terminal escape sequences.
+            #
+            # See:
+            # + https://tiswww.cwru.edu/php/chet/readline/readline.html#index-rl_005fexpand_005fprompt
+            # + https://git.savannah.gnu.org/cgit/readline.git/tree/readline.h?id=037d85f199a8c6e5b16689a46c8bc31b586a0c94#n884
+            text = _ANSI_RE.sub(
+                lambda match: "\001{ansi_sequence}\002".format(ansi_sequence=match.group(0)), text
+            )
+        return text + " " if prompt else text
+
+    repl_banner = fixup_ansi(banner) if banner else banner
 
     def loop():
         # type: () -> Dict[str, Any]
         if ps1:
-            sys.ps1 = ps1
+            sys.ps1 = fixup_ansi(ps1, prompt=True)
         if ps2:
-            sys.ps2 = ps2
+            sys.ps2 = fixup_ansi(ps2, prompt=True)
         repl.interact(banner=repl_banner, **extra_args)
         return local
 

pex/repl/pex_repl.py

@@ -32,6 +32,8 @@
     from pex.third_party import attr
 
 
+_ARGV0 = os.environ.get("SCIE_ARGV0", sys.argv[0])
+
 _PEX_CLI_RUN_ENV_VAR_NAME = "_PEX_CLI_RUN"
 _PEX_CLI_RUN_NO_ARGS_ENV_VAR_NAME = "_PEX_CLI_RUN_NO_ARGS"
 
@@ -41,9 +43,9 @@ def export_pex_cli_run(env=None):
     """Records the fact that the Pex CLI executable is being run with no arguments."""
 
     _env = cast("Dict[str, str]", env or os.environ)
-    _env[_PEX_CLI_RUN_ENV_VAR_NAME] = sys.argv[0]
+    _env[_PEX_CLI_RUN_ENV_VAR_NAME] = _ARGV0
     if len(sys.argv) == 1:
-        _env[_PEX_CLI_RUN_NO_ARGS_ENV_VAR_NAME] = sys.argv[0]
+        _env[_PEX_CLI_RUN_NO_ARGS_ENV_VAR_NAME] = _ARGV0
     return _env
 
 
@@ -128,6 +130,11 @@ def pex(json=False):
         },
         history=history,
         history_file=history_file,
+        # The PBS Linux CPythons used by PEX scies use a version of libedit that needs workarounds
+        # to support color prompts.
+        use_libedit_color_prompt_workaround=(
+            os.environ.get("SCIE") is not None and sys.platform.lower().startswith("linux")
+        ),
     )
 
 
@@ -143,7 +150,7 @@ def _create_repl_data(
     pex_info,  # type: PexInfo
     requirements,  # type: Sequence[str]
     activated_dists,  # type: Sequence[Distribution]
-    pex=sys.argv[0],  # type: str
+    pex=_ARGV0,  # type: str
     env=ENV,  # type: Variables
     venv=False,  # type: bool
 ):
@@ -239,16 +246,16 @@ def _create_repl_data(
     return _REPLData(
         banner=banner,
         pex_info_summary=os.linesep.join(pex_info_summary),
-        ps1=colors.yellow(">>> "),
-        ps2=colors.yellow("... "),
+        ps1=colors.yellow(">>>"),
+        ps2=colors.yellow("..."),
     )
 
 
 def create_pex_repl_exe(
     shebang,  # type: str
     pex_info,  # type: PexInfo
     activated_dists,  # type: Sequence[Distribution]
-    pex=sys.argv[0],  # type: str
+    pex=_ARGV0,  # type: str
     env=ENV,  # type: Variables
     venv=False,  # type: bool
 ):
@@ -306,7 +313,7 @@ def create_pex_repl(
     pex_info,  # type: PexInfo
     requirements,  # type: Sequence[str]
     activated_dists,  # type: Sequence[Distribution]
-    pex=sys.argv[0],  # type: str
+    pex=_ARGV0,  # type: str
     env=ENV,  # type: Variables
 ):
     # type: (...) -> Callable[[], Dict[str, Any]]

pex/version.py

@@ -1,4 +1,4 @@
 # Copyright 2015 Pex project contributors.
 # Licensed under the Apache License, Version 2.0 (see LICENSE).
 
-__version__ = "2.16.0"
+__version__ = "2.16.1"