diff -r 8d6d068e9966 Lib/os.py --- a/Lib/os.py Thu Dec 15 19:24:49 2011 -0500 +++ b/Lib/os.py Sat Dec 17 02:32:15 2011 +0100 @@ -829,3 +829,38 @@ raise TypeError("invalid fd type (%s, expected integer)" % type(fd)) import io return io.open(fd, *args, **kwargs) + +def get_terminal_size(columns=80, rows=25): + """Get the size of the terminal window. + + First, the environment variables COLUMNS and ROWS are checked. If + they are defined and the values are valid positive integers, they + are used. + + When COLUMNS or ROWS are not defined, which is the common case, + the terminal connected to sys.stdout is queried. + + If the terminal size cannot be sucessfully queried, either because + the system doesn't support querying, or we are not connected to a + terminal, the default values given in the parameters are used. + + The value returned is a named tuple of type TerminalSize. + + Right now the values are not cached, but this might change. + """ + try: + columns = int(environ['COLUMNS']) + except (KeyError, ValueError): + columns = 0 + + try: + rows = int(environ['ROWS']) + except (KeyError, ValueError): + rows = 0 + + if columns <= 0 or rows <= 0: + size = query_terminal_size() + columns = columns if columns > 0 else size.columns + rows = rows if rows > 0 else size.rows + + return TerminalSize((columns, rows)) diff -r 8d6d068e9966 Lib/test/test_termsize.py --- /dev/null Thu Jan 01 00:00:00 1970 +0000 +++ b/Lib/test/test_termsize.py Sat Dec 17 02:32:15 2011 +0100 @@ -0,0 +1,37 @@ +import unittest +import os +import contextlib +from test import support + +@contextlib.contextmanager +def set_environ(**variables): + old = {var:os.environ.get(var, None) for var in variables} + os.environ.update(variables) + try: + yield + finally: + for var, value in old: + if value is None: + del os.environ[var] + else: + os.environ[var] = value + +class TermsizeTests(unittest.TestCase): + def test_does_not_crash(): + # There's no easy portable way to actually check the size of + # the terminal, so let's check if it returns something + # sensible instead. + size = os.get_terminal_size() + self.assertTrue(0 < size.columns) + self.assertTrue(0 < size.rows) + + def test_os_environ_first(): + # Check if environment variables have precedence + + with set_environ(COLUMNS=777): + size = os.get_terminal_size() + self.assertTrue(size.columns == 777) + + with set_environ(ROWS=888): + size = os.get_terminal_size() + self.assertTrue(size.rows == 888) diff -r 8d6d068e9966 Modules/posixmodule.c --- a/Modules/posixmodule.c Thu Dec 15 19:24:49 2011 -0500 +++ b/Modules/posixmodule.c Sat Dec 17 02:32:15 2011 +0100 @@ -125,6 +125,24 @@ #include #endif +#ifdef HAVE_SYS_IOCTL_H +#include +#define TERMSIZE_USE_IOCTL +#else + +#ifdef HAVE_CONIO_H +#include +#include +#define TERMSIZE_USE_CONIO + +#else + +#define TERMSIZE_USE_NOTIMPLEMENTED + +#endif /* HAVE_CONIO_H */ +#endif /* HAVE_SYS_IOCTL_H */ + + /* Various compilers have only certain posix functions */ /* XXX Gosh I wish these were all moved into pyconfig.h */ #if defined(PYCC_VACPP) && defined(PYOS_OS2) @@ -10495,6 +10513,109 @@ #endif /* USE_XATTRS */ + +/* Terminal size querying */ + +PyDoc_STRVAR(termsize__doc__, + "Return the size of the terminal window as (columns, rows).\n" \ + "\n" \ + "The optional argument fd (default 1, or standard output) specifies\n" \ + "which file descriptor should be queried.\n" \ + "\n" \ + "If the file descriptor is not connected to a terminal, an IOError\n" \ + "is thrown. If the function is not available for this system, a\n" \ + "NotImplementedError is thrown.\n" \ + "\n" \ + "get_terminal_size is the highlevel function which should normally\n" \ + "be used, query_terminal_size is the low-level implementation."); + +static PyTypeObject TerminalSizeType; + +PyDoc_STRVAR(TerminalSize_docstring, + "A tuple of (columns, rows) for holding terminal window size"); + +static PyStructSequence_Field TerminalSize_fields[] = { + {"columns", "width of the terminal window in characters"}, + {"rows", "height of the terminal window in characters"}, + {NULL, NULL} +}; + +static PyStructSequence_Desc TerminalSize_desc = { + "os.TerminalSize", + TerminalSize_docstring, + TerminalSize_fields, + 2, +}; + +static PyObject* +query_terminal_size(PyObject *self, PyObject *args) +{ + int columns, rows; + PyObject *termsize; + int fd = 1; /* stdout */ + + if (!PyArg_ParseTuple(args, "|i", &fd)) + return NULL; + +#ifdef TERMSIZE_USE_NOTIMPLEMENTED + return PyErr_Format(PyExc_NotImplementedError, + "not imlemented for this system, sorry"); +#endif /* TERMSIZE_USE_NOTIMPLEMENTED */ + +#ifdef TERMSIZE_USE_IOCTL + { + struct winsize w; + if (ioctl(fd, TIOCGWINSZ, &w)) { + switch (errno) { + case EBADF: + return PyErr_Format(PyExc_ValueError, "bad file descriptor"); + default: + return PyErr_SetFromErrno(PyExc_IOError); + } + } + columns = w.ws_col; + rows = w.ws_row; + } +#endif /* TERMSIZE_USE_IOCTL */ + +#ifdef TERMSIZE_USE_CONIO + { HANDLE handle; CONSOLE_SCREEN_BUFFER_INFO csbi; + DWORD nhandle; + switch (fd) { + case 0: nhandle = STD_INPUT_HANDLE; + break; + case 1: nhandle = STD_OUTPUT_HANDLE; + break; + case 2: nhandle = STD_ERROR_HANDLE; + break; + default: + return PyErr_Format(PyExc_ValueError, "bad file descriptor"); + } + handle = GetStdHandle(nhandle); + if (handle == INVALID_HANDLE_VALUE) + return PyErr_Format(PyExc_ValueError, "invalid handle value"); + + if (!GetConsoleScreenBufferInfo(handle, &csbi)) + return PyErr_Format(PyExc_IOError, "error %i", (int) GetLastError()); + + columns = csbi.srWindow.Right - csbi.srWindow.Left + 1; + rows = csbi.srWindow.Bottom - csbi.srWindow.Top + 1; + } +#endif /* TERMSIZE_USE_CONIO */ + + termsize = PyStructSequence_New(&TerminalSizeType); + if (termsize == NULL) + return NULL; + PyStructSequence_SET_ITEM(termsize, 0, PyLong_FromLong(columns)); + PyStructSequence_SET_ITEM(termsize, 1, PyLong_FromLong(rows)); + if (PyErr_Occurred()) { + Py_CLEAR(termsize); + return NULL; + } + return termsize; +} + + static PyMethodDef posix_methods[] = { {"access", posix_access, METH_VARARGS, posix_access__doc__}, #ifdef HAVE_TTYNAME @@ -10962,6 +11083,7 @@ {"llistxattr", posix_llistxattr, METH_VARARGS, posix_llistxattr__doc__}, {"flistxattr", posix_flistxattr, METH_VARARGS, posix_flistxattr__doc__}, #endif + {"query_terminal_size", query_terminal_size, METH_VARARGS, termsize__doc__}, {NULL, NULL} /* Sentinel */ }; @@ -11556,6 +11678,10 @@ PyStructSequence_InitType(&SchedParamType, &sched_param_desc); SchedParamType.tp_new = sched_param_new; #endif + + /* initialize TerminalSize_info */ + PyStructSequence_InitType(&TerminalSizeType, &TerminalSize_desc); + Py_INCREF(&TerminalSizeType); } #if defined(HAVE_WAITID) && !defined(__APPLE__) Py_INCREF((PyObject*) &WaitidResultType); @@ -11610,6 +11736,9 @@ #endif /* __APPLE__ */ + + PyModule_AddObject(m, "TerminalSize", (PyObject*) &TerminalSizeType); + return m; } diff -r 8d6d068e9966 configure --- a/configure Thu Dec 15 19:24:49 2011 -0500 +++ b/configure Sat Dec 17 02:32:15 2011 +0100 @@ -6147,7 +6147,7 @@ sys/audioio.h sys/xattr.h sys/bsdtty.h sys/event.h sys/file.h sys/loadavg.h \ sys/lock.h sys/mkdev.h sys/modem.h \ sys/param.h sys/select.h sys/sendfile.h sys/socket.h sys/statvfs.h \ -sys/stat.h sys/termio.h sys/time.h \ +sys/stat.h sys/termio.h sys/time.h sys/ioctl.h \ sys/times.h sys/types.h sys/uio.h sys/un.h sys/utsname.h sys/wait.h pty.h \ libutil.h sys/resource.h netpacket/packet.h sysexits.h bluetooth.h \ bluetooth/bluetooth.h linux/tipc.h spawn.h util.h diff -r 8d6d068e9966 configure.in --- a/configure.in Thu Dec 15 19:24:49 2011 -0500 +++ b/configure.in Sat Dec 17 02:32:15 2011 +0100 @@ -1337,7 +1337,7 @@ sys/audioio.h sys/xattr.h sys/bsdtty.h sys/event.h sys/file.h sys/loadavg.h \ sys/lock.h sys/mkdev.h sys/modem.h \ sys/param.h sys/select.h sys/sendfile.h sys/socket.h sys/statvfs.h \ -sys/stat.h sys/termio.h sys/time.h \ +sys/stat.h sys/termio.h sys/time.h sys/ioctl.h \ sys/times.h sys/types.h sys/uio.h sys/un.h sys/utsname.h sys/wait.h pty.h \ libutil.h sys/resource.h netpacket/packet.h sysexits.h bluetooth.h \ bluetooth/bluetooth.h linux/tipc.h spawn.h util.h) diff -r 8d6d068e9966 pyconfig.h.in --- a/pyconfig.h.in Thu Dec 15 19:24:49 2011 -0500 +++ b/pyconfig.h.in Sat Dec 17 02:32:15 2011 +0100 @@ -902,6 +902,9 @@ /* Define to 1 if you have the header file. */ #undef HAVE_SYS_FILE_H +/* Define to 1 if you have the header file. */ +#undef HAVE_SYS_IOCTL_H + /* Define to 1 if you have the header file. */ #undef HAVE_SYS_LOADAVG_H