Issue #4705: Fix the -u ("unbuffered binary stdout and stderr") command-line

antoine.pitrou · antoine.pitrou · commit 4e1117377e0c · 2009-01-09T18:53:14.000Z
flag to work properly. Furthermore, when specifying -u, the text stdout and stderr streams have line-by-line buffering enabled (the default being to buffer arbitrary chunks of data). Patch by Victor Stinner, test by me. git-svn-id: http://svn.python.org/projects/python/branches/py3k@68451 6015fed2-1504-0410-9fe1-9d1591cc4771
diff --git a/Include/pydebug.h b/Include/pydebug.h
@@ -18,6 +18,7 @@ PyAPI_DATA(int) Py_IgnoreEnvironmentFlag;
 PyAPI_DATA(int) Py_DivisionWarningFlag;
 PyAPI_DATA(int) Py_DontWriteBytecodeFlag;
 PyAPI_DATA(int) Py_NoUserSiteDirectory;
+PyAPI_DATA(int) Py_UnbufferedStdioFlag;
 
 /* this is a wrapper around getenv() that pays attention to
    Py_IgnoreEnvironmentFlag.  It should be used for getting variables like
diff --git a/Lib/test/test_cmd_line.py b/Lib/test/test_cmd_line.py
@@ -142,6 +142,23 @@ def test_run_code(self):
                 self.exit_code('-c', command),
                 0)
 
+    def test_unbuffered_output(self):
+        # Test expected operation of the '-u' switch
+        for stream in ('stdout', 'stderr'):
+            # Binary is unbuffered
+            code = ("import os, sys; sys.%s.buffer.write(b'x'); os._exit(0)"
+                % stream)
+            data, rc = self.start_python_and_exit_code('-u', '-c', code)
+            self.assertEqual(rc, 0)
+            self.assertEqual(data, b'x', "binary %s not unbuffered" % stream)
+            # Text is line-buffered
+            code = ("import os, sys; sys.%s.write('x\\n'); os._exit(0)"
+                % stream)
+            data, rc = self.start_python_and_exit_code('-u', '-c', code)
+            self.assertEqual(rc, 0)
+            self.assertEqual(data.strip(), b'x',
+                "text %s not line-buffered" % stream)
+
 
 def test_main():
     test.support.run_unittest(CmdLineTest)
diff --git a/Misc/NEWS b/Misc/NEWS
@@ -12,6 +12,11 @@ What's New in Python 3.1 alpha 0
 Core and Builtins
 -----------------
 
+- Issue #4705: Fix the -u ("unbuffered binary stdout and stderr") command-line
+  flag to work properly. Furthermore, when specifying -u, the text stdout
+  and stderr streams have line-by-line buffering enabled (the default being
+  to buffer arbitrary chunks of data).
+
 - The internal table, _PyLong_DigitValue, is now an array of unsigned chars
   instead of ints (reducing its size from 4 to 8 times thereby reducing
   Python's overall memory).
diff --git a/Modules/main.c b/Modules/main.c
@@ -292,7 +292,6 @@ Py_Main(int argc, wchar_t **argv)
 	wchar_t *module = NULL;
 	FILE *fp = stdin;
 	char *p;
-	int unbuffered = 0;
 	int skipfirstline = 0;
 	int stdin_is_interactive = 0;
 	int help = 0;
@@ -374,7 +373,7 @@ Py_Main(int argc, wchar_t **argv)
 			break;
 
 		case 'u':
-			unbuffered++;
+			Py_UnbufferedStdioFlag = 1;
 			saw_unbuffered_flag = 1;
 			break;
 
@@ -423,7 +422,7 @@ Py_Main(int argc, wchar_t **argv)
 		Py_InspectFlag = 1;
 	if (!saw_unbuffered_flag &&
 	    (p = Py_GETENV("PYTHONUNBUFFERED")) && *p != '\0')
-		unbuffered = 1;
+		Py_UnbufferedStdioFlag = 1;
 
 	if (!Py_NoUserSiteDirectory &&
 	    (p = Py_GETENV("PYTHONNOUSERSITE")) && *p != '\0')
@@ -444,7 +443,7 @@ Py_Main(int argc, wchar_t **argv)
 
 	stdin_is_interactive = Py_FdIsInteractive(stdin, (char *)0);
 
-	if (unbuffered) {
+	if (Py_UnbufferedStdioFlag) {
 #if defined(MS_WINDOWS) || defined(__CYGWIN__)
 		_setmode(fileno(stdin), O_BINARY);
 		_setmode(fileno(stdout), O_BINARY);
diff --git a/Python/pythonrun.c b/Python/pythonrun.c
@@ -88,6 +88,7 @@ int Py_UseClassExceptionsFlag = 1; /* Needed by bltinmodule.c: deprecated */
 int Py_FrozenFlag; /* Needed by getpath.c */
 int Py_IgnoreEnvironmentFlag; /* e.g. PYTHONPATH, PYTHONHOME */
 int Py_NoUserSiteDirectory = 0; /* for -s and site.py */
+int Py_UnbufferedStdioFlag = 0; /* Unbuffered binary std{in,out,err} */
 
 /* PyModule_GetWarningsModule is no longer necessary as of 2.6
 since _warnings is builtin.  This API should not be used. */
@@ -728,6 +729,75 @@ initsite(void)
 	}
 }
 
+static PyObject*
+create_stdio(PyObject* io,
+	int fd, int write_mode, char* name,
+	char* encoding, char* errors)
+{
+	PyObject *buf = NULL, *stream = NULL, *text = NULL, *raw = NULL;
+	const char* mode;
+	const PyObject *line_buffering;
+	int buffering;
+
+	if (Py_UnbufferedStdioFlag)
+		buffering = 0;
+	else
+		buffering = -1;
+	if (write_mode)
+		mode = "wb";
+	else
+		mode = "rb";
+	buf = PyObject_CallMethod(io, "open", "isiOOOi",
+				  fd, mode, buffering,
+				  Py_None, Py_None, Py_None, 0);
+	if (buf == NULL)
+		goto error;
+
+	if (!Py_UnbufferedStdioFlag) {
+		raw = PyObject_GetAttrString(buf, "raw");
+		if (raw == NULL)
+			goto error;
+	}
+	else {
+		raw = buf;
+		Py_INCREF(raw);
+	}
+
+	text = PyUnicode_FromString(name);
+	if (text == NULL || PyObject_SetAttrString(raw, "_name", text) < 0)
+		goto error;
+	Py_CLEAR(raw);
+	Py_CLEAR(text);
+
+	if (Py_UnbufferedStdioFlag)
+		line_buffering = Py_True;
+	else
+		line_buffering = Py_False;
+	stream = PyObject_CallMethod(io, "TextIOWrapper", "OsssO",
+				     buf, encoding, errors,
+				     "\n", line_buffering);
+	Py_CLEAR(buf);
+	if (stream == NULL)
+		goto error;
+
+	if (write_mode)
+		mode = "w";
+	else
+		mode = "r";
+	text = PyUnicode_FromString(mode);
+	if (!text || PyObject_SetAttrString(stream, "mode", text) < 0)
+		goto error;
+	Py_CLEAR(text);
+	return stream;
+
+error:
+	Py_XDECREF(buf);
+	Py_XDECREF(stream);
+	Py_XDECREF(text);
+	Py_XDECREF(raw);
+	return NULL;
+}
+
 /* Initialize sys.stdin, stdout, stderr and builtins.open */
 static int
 initstdio(void)
@@ -794,10 +864,9 @@ initstdio(void)
 #endif
 	}
 	else {
-		if (!(std = PyFile_FromFd(fd, "<stdin>", "r", -1, encoding, 
-					  errors, "\n", 0))) {
+		std = create_stdio(iomod, fd, 0, "<stdin>", encoding, errors);
+		if (std == NULL)
 			goto error;
-		}
 	} /* if (fd < 0) */
 	PySys_SetObject("__stdin__", std);
 	PySys_SetObject("stdin", std);
@@ -814,10 +883,9 @@ initstdio(void)
 #endif
 	}
 	else {
-		if (!(std = PyFile_FromFd(fd, "<stdout>", "w", -1, encoding, 
-					  errors, "\n", 0))) {
+		std = create_stdio(iomod, fd, 1, "<stdout>", encoding, errors);
+		if (std == NULL)
 			goto error;
-		}
 	} /* if (fd < 0) */
 	PySys_SetObject("__stdout__", std);
 	PySys_SetObject("stdout", std);
@@ -835,10 +903,9 @@ initstdio(void)
 #endif
 	}
 	else {
-		if (!(std = PyFile_FromFd(fd, "<stderr>", "w", -1, encoding,
-					  "backslashreplace", "\n", 0))) {
+		std = create_stdio(iomod, fd, 1, "<stderr>", encoding, "backslashreplace");
+		if (std == NULL)
 			goto error;
-		}
 	} /* if (fd < 0) */
 
 	/* Same as hack above, pre-import stderr's codec to avoid recursion
