From 34f716fe1fb4df9a5257b5f640ebb8a71e10aa88 Mon Sep 17 00:00:00 2001
From: Daniel Hahler <git@thequod.de>
Date: Fri, 5 Apr 2019 14:51:52 +0200
Subject: [PATCH] Fix _getdimensions for when stdout is redirected

This uses an improved version of `shutil.get_terminal_width` [1], and
also improves the code for before Python 3.3.

1: https://bugs.python.org/issue14841,
   https://github.com/python/cpython/pull/12697
---
 py/_io/terminalwriter.py           | 56 ++++++++++++++++++++++++++----
 testing/io_/test_terminalwriter.py | 37 ++++++++++++++++----
 2 files changed, 79 insertions(+), 14 deletions(-)

diff --git a/py/_io/terminalwriter.py b/py/_io/terminalwriter.py
index be559867..05118c1f 100644
--- a/py/_io/terminalwriter.py
+++ b/py/_io/terminalwriter.py
@@ -26,14 +26,56 @@
 
 def _getdimensions():
     if py33:
-        import shutil
-        size = shutil.get_terminal_size()
-        return size.lines, size.columns
+        # Improved version of shutil.get_terminal_size that looks at stdin,
+        # stderr, stdout.  Ref: https://bugs.python.org/issue14841.
+        fallback = (80, 24)
+        # columns, lines are the working values
+        try:
+            columns = int(os.environ['COLUMNS'])
+        except (KeyError, ValueError):
+            columns = 0
+
+        try:
+            lines = int(os.environ['LINES'])
+        except (KeyError, ValueError):
+            lines = 0
+
+        # only query if necessary
+        if columns <= 0 or lines <= 0:
+            try:
+                os_get_terminal_size = os.get_terminal_size
+            except AttributeError:
+                size = os.terminal_size(fallback)
+            else:
+                for check in [sys.__stdin__, sys.__stderr__, sys.__stdout__]:
+                    try:
+                        size = os_get_terminal_size(check.fileno())
+                    except (AttributeError, ValueError, OSError):
+                        # fd is None, closed, detached, or not a terminal.
+                        continue
+                    else:
+                        break
+                else:
+                    size = os.terminal_size(fallback)
+            if columns <= 0:
+                columns = size.columns
+            if lines <= 0:
+                lines = size.lines
+
+        return lines, columns
     else:
-        import termios, fcntl, struct
-        call = fcntl.ioctl(1, termios.TIOCGWINSZ, "\000" * 8)
-        height, width = struct.unpack("hhhh", call)[:2]
-        return height, width
+        import termios
+        import fcntl
+        import struct
+        for fd in (0, 2, 1):
+            try:
+                call = fcntl.ioctl(fd, termios.TIOCGWINSZ, "\000" * 8)
+            except OSError:
+                continue
+            height, width = struct.unpack("hhhh", call)[:2]
+            return height, width
+
+        return 24, 80
 
 
 def get_terminal_width():
diff --git a/testing/io_/test_terminalwriter.py b/testing/io_/test_terminalwriter.py
index 1eef7f7d..c138f9a1 100644
--- a/testing/io_/test_terminalwriter.py
+++ b/testing/io_/test_terminalwriter.py
@@ -10,23 +10,46 @@ def test_get_terminal_width():
     x = py.io.get_terminal_width
     assert x == terminalwriter.get_terminal_width
 
-def test_getdimensions(monkeypatch):
+
+@pytest.mark.parametrize("via_fd", (0, 1, 2))
+def test_getdimensions(via_fd, monkeypatch):
+    mock_calls = []
+
     if sys.version_info >= (3, 3):
-        import shutil
         Size = namedtuple('Size', 'lines columns')
-        monkeypatch.setattr(shutil, 'get_terminal_size', lambda: Size(60, 100))
+
+        def os_get_terminal_size(*args):
+            mock_calls.append(args)
+            fd = args[0]
+            if fd != via_fd:
+                raise ValueError
+            return Size(60, 100)
+        monkeypatch.setattr(os, 'get_terminal_size', os_get_terminal_size)
         assert terminalwriter._getdimensions() == (60, 100)
+
     else:
         fcntl = py.test.importorskip("fcntl")
         import struct
-        l = []
-        monkeypatch.setattr(fcntl, 'ioctl', lambda *args: l.append(args))
+
+        def mock_ioctl(*args):
+            mock_calls.append(args)
+            fd = args[0]
+            if fd != via_fd:
+                raise OSError
+
+        monkeypatch.setattr(fcntl, 'ioctl', mock_ioctl)
         try:
             terminalwriter._getdimensions()
         except (TypeError, struct.error):
             pass
-        assert len(l) == 1
-        assert l[0][0] == 1
+
+    if via_fd == 0:
+        assert len(mock_calls) == 1
+    elif via_fd == 2:
+        assert len(mock_calls) == 2
+    elif via_fd == 1:
+        assert len(mock_calls) == 3
+    assert mock_calls[-1][0] == via_fd
 
 def test_terminal_width_COLUMNS(monkeypatch):
     """ Dummy test for get_terminal_width