gh-128636: Fix crash in PyREPL when os.environ is overwritten with an invalid value for mac

Lib/_colorize.py

@@ -286,7 +286,10 @@ def decolor(text: str) -> str:
     return text
 
 
-def can_colorize(*, file: IO[str] | IO[bytes] | None = None) -> bool:
+def can_colorize(*, file: IO[str] | IO[bytes] | None = None, already_colorize: bool = False) -> bool:
+    if already_colorize:
+        return True
+
     if file is None:
         file = sys.stdout
 
@@ -331,6 +334,7 @@ def get_theme(
     tty_file: IO[str] | IO[bytes] | None = None,
     force_color: bool = False,
     force_no_color: bool = False,
+    already_colorize: bool = False,
 ) -> Theme:
     """Returns the currently set theme, potentially in a zero-color variant.
 
@@ -344,7 +348,8 @@ def get_theme(
     environment (including environment variable state and console configuration
     on Windows) can also change in the course of the application life cycle.
     """
-    if force_color or (not force_no_color and can_colorize(file=tty_file)):
+    if force_color or (not force_no_color and
+                       can_colorize(file=tty_file, already_colorize=already_colorize)):
         return _theme
     return theme_no_color
 

Lib/_pyrepl/reader.py

@@ -343,7 +343,7 @@ def calc_screen(self) -> list[str]:
                 screeninfo.append((0, []))
             pos -= line_len + 1
             prompt, prompt_len = self.process_prompt(prompt)
-            chars, char_widths = disp_str(line, colors, offset)
+            chars, char_widths = disp_str(line, colors, offset, already_colorize=self.can_colorize)
             wrapcount = (sum(char_widths) + prompt_len) // self.console.width
             if wrapcount == 0 or not char_widths:
                 offset += line_len + 1  # Takes all of the line plus the newline
@@ -491,7 +491,7 @@ def get_prompt(self, lineno: int, cursor_on_line: bool) -> str:
             prompt = self.ps1
 
         if self.can_colorize:
-            t = THEME()
+            t = THEME(already_colorize=self.can_colorize)
             prompt = f"{t.prompt}{prompt}{t.reset}"
         return prompt
 

Lib/_pyrepl/unix_console.py

@@ -159,6 +159,10 @@ def __init__(
         self.pollob.register(self.input_fd, select.POLLIN)
         self.terminfo = terminfo.TermInfo(term or None)
         self.term = term
+        self.is_mac = (
+            platform.system() == "Darwin"
+            and os.getenv("TERM_PROGRAM") == "Apple_Terminal"
+        )
 
         @overload
         def _my_getstr(cap: str, optional: Literal[False] = False) -> bytes: ...
@@ -339,7 +343,7 @@ def prepare(self):
         tcsetattr(self.input_fd, termios.TCSADRAIN, raw)
 
         # In macOS terminal we need to deactivate line wrap via ANSI escape code
-        if platform.system() == "Darwin" and os.getenv("TERM_PROGRAM") == "Apple_Terminal":
+        if self.is_mac:
             os.write(self.output_fd, b"\033[?7l")
 
         self.screen = []
@@ -370,7 +374,7 @@ def restore(self):
         self.flushoutput()
         tcsetattr(self.input_fd, termios.TCSADRAIN, self.__svtermstate)
 
-        if platform.system() == "Darwin" and os.getenv("TERM_PROGRAM") == "Apple_Terminal":
+        if self.is_mac:
             os.write(self.output_fd, b"\033[?7h")
 
         if hasattr(self, "old_sigwinch"):

Lib/_pyrepl/utils.py

@@ -266,6 +266,7 @@ def disp_str(
     colors: list[ColorSpan] | None = None,
     start_index: int = 0,
     force_color: bool = False,
+    already_colorize: bool = False,
 ) -> tuple[CharBuffer, CharWidths]:
     r"""Decompose the input buffer into a printable variant with applied colors.
 
@@ -306,7 +307,7 @@ def disp_str(
         # move past irrelevant spans
         colors.pop(0)
 
-    theme = THEME(force_color=force_color)
+    theme = THEME(force_color=force_color, already_colorize=already_colorize)
     pre_color = ""
     post_color = ""
     if colors and colors[0].span.start < start_index:

Lib/test/support/__init__.py

@@ -2898,7 +2898,7 @@ def force_color(color: bool):
     from .os_helper import EnvironmentVarGuard
 
     with (
-        swap_attr(_colorize, "can_colorize", lambda file=None: color),
+        swap_attr(_colorize, "can_colorize", lambda *, file=None, already_colorize=False: color),
         EnvironmentVarGuard() as env,
     ):
         env.unset("FORCE_COLOR", "NO_COLOR", "PYTHON_COLORS")

Lib/test/test_pyrepl/test_unix_console.py

@@ -303,3 +303,12 @@ def test_getheightwidth_with_invalid_environ(self, _os_write):
             self.assertIsInstance(console.getheightwidth(), tuple)
             os.environ = []
             self.assertIsInstance(console.getheightwidth(), tuple)
+
+    @unittest.skipUnless(sys.platform == "darwin", "requires OS X")
+    def test_restore_with_invalid_environ_on_macOS(self, _os_write):
+        # gh-128636 for macOS
+        console = UnixConsole(term="xterm")
+        with os_helper.EnvironmentVarGuard() as env:
+            os.environ = []
+            console.prepare()  # needed to call restore()
+            console.restore()  # this should succeed