Skip to content

Commit

Permalink
Fix _getdimensions for when stdout is redirected
Browse files Browse the repository at this point in the history
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,
   python/cpython#12697
  • Loading branch information
blueyed committed Aug 17, 2019
1 parent 1e99d20 commit 34f716f
Show file tree
Hide file tree
Showing 2 changed files with 79 additions and 14 deletions.
56 changes: 49 additions & 7 deletions py/_io/terminalwriter.py
Original file line number Diff line number Diff line change
Expand Up @@ -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():
Expand Down
37 changes: 30 additions & 7 deletions testing/io_/test_terminalwriter.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down

0 comments on commit 34f716f

Please sign in to comment.