Skip to content

Commit 34f716f

Browse files
committed
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, python/cpython#12697
1 parent 1e99d20 commit 34f716f

File tree

2 files changed

+79
-14
lines changed

2 files changed

+79
-14
lines changed

py/_io/terminalwriter.py

Lines changed: 49 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -26,14 +26,56 @@
2626

2727
def _getdimensions():
2828
if py33:
29-
import shutil
30-
size = shutil.get_terminal_size()
31-
return size.lines, size.columns
29+
# Improved version of shutil.get_terminal_size that looks at stdin,
30+
# stderr, stdout. Ref: https://bugs.python.org/issue14841.
31+
fallback = (80, 24)
32+
# columns, lines are the working values
33+
try:
34+
columns = int(os.environ['COLUMNS'])
35+
except (KeyError, ValueError):
36+
columns = 0
37+
38+
try:
39+
lines = int(os.environ['LINES'])
40+
except (KeyError, ValueError):
41+
lines = 0
42+
43+
# only query if necessary
44+
if columns <= 0 or lines <= 0:
45+
try:
46+
os_get_terminal_size = os.get_terminal_size
47+
except AttributeError:
48+
size = os.terminal_size(fallback)
49+
else:
50+
for check in [sys.__stdin__, sys.__stderr__, sys.__stdout__]:
51+
try:
52+
size = os_get_terminal_size(check.fileno())
53+
except (AttributeError, ValueError, OSError):
54+
# fd is None, closed, detached, or not a terminal.
55+
continue
56+
else:
57+
break
58+
else:
59+
size = os.terminal_size(fallback)
60+
if columns <= 0:
61+
columns = size.columns
62+
if lines <= 0:
63+
lines = size.lines
64+
65+
return lines, columns
3266
else:
33-
import termios, fcntl, struct
34-
call = fcntl.ioctl(1, termios.TIOCGWINSZ, "\000" * 8)
35-
height, width = struct.unpack("hhhh", call)[:2]
36-
return height, width
67+
import termios
68+
import fcntl
69+
import struct
70+
for fd in (0, 2, 1):
71+
try:
72+
call = fcntl.ioctl(fd, termios.TIOCGWINSZ, "\000" * 8)
73+
except OSError:
74+
continue
75+
height, width = struct.unpack("hhhh", call)[:2]
76+
return height, width
77+
78+
return 24, 80
3779

3880

3981
def get_terminal_width():

testing/io_/test_terminalwriter.py

Lines changed: 30 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -10,23 +10,46 @@ def test_get_terminal_width():
1010
x = py.io.get_terminal_width
1111
assert x == terminalwriter.get_terminal_width
1212

13-
def test_getdimensions(monkeypatch):
13+
14+
@pytest.mark.parametrize("via_fd", (0, 1, 2))
15+
def test_getdimensions(via_fd, monkeypatch):
16+
mock_calls = []
17+
1418
if sys.version_info >= (3, 3):
15-
import shutil
1619
Size = namedtuple('Size', 'lines columns')
17-
monkeypatch.setattr(shutil, 'get_terminal_size', lambda: Size(60, 100))
20+
21+
def os_get_terminal_size(*args):
22+
mock_calls.append(args)
23+
fd = args[0]
24+
if fd != via_fd:
25+
raise ValueError
26+
return Size(60, 100)
27+
monkeypatch.setattr(os, 'get_terminal_size', os_get_terminal_size)
1828
assert terminalwriter._getdimensions() == (60, 100)
29+
1930
else:
2031
fcntl = py.test.importorskip("fcntl")
2132
import struct
22-
l = []
23-
monkeypatch.setattr(fcntl, 'ioctl', lambda *args: l.append(args))
33+
34+
def mock_ioctl(*args):
35+
mock_calls.append(args)
36+
fd = args[0]
37+
if fd != via_fd:
38+
raise OSError
39+
40+
monkeypatch.setattr(fcntl, 'ioctl', mock_ioctl)
2441
try:
2542
terminalwriter._getdimensions()
2643
except (TypeError, struct.error):
2744
pass
28-
assert len(l) == 1
29-
assert l[0][0] == 1
45+
46+
if via_fd == 0:
47+
assert len(mock_calls) == 1
48+
elif via_fd == 2:
49+
assert len(mock_calls) == 2
50+
elif via_fd == 1:
51+
assert len(mock_calls) == 3
52+
assert mock_calls[-1][0] == via_fd
3053

3154
def test_terminal_width_COLUMNS(monkeypatch):
3255
""" Dummy test for get_terminal_width

0 commit comments

Comments
 (0)