Generalize file.writelines() to allow iterable objects.

5 files changed, 91 insertions, 35 deletions

diff --git a/Doc/lib/libstdtypes.tex b/Doc/lib/libstdtypes.tex
index a79e142..916a1cf 100644
--- a/Doc/lib/libstdtypes.tex
+++ b/Doc/lib/libstdtypes.tex
@@ -1312,8 +1312,10 @@ Files have the following methods:
   the \method{flush()} or \method{close()} method is called.
 \end{methoddesc}
 
-\begin{methoddesc}[file]{writelines}{list}
-  Write a list of strings to the file.  There is no return value.
+\begin{methoddesc}[file]{writelines}{sequence}
+  Write a sequence of strings to the file.  The sequence can be any
+  iterable object producing strings, typically a list of strings.
+  There is no return value.
   (The name is intended to match \method{readlines()};
   \method{writelines()} does not add line separators.)
 \end{methoddesc}
diff --git a/Lib/test/test_iter.py b/Lib/test/test_iter.py
index f6084cb..257b61d 100644
--- a/Lib/test/test_iter.py
+++ b/Lib/test/test_iter.py
@@ -641,6 +641,59 @@ class TestCase(unittest.TestCase):
             self.assertEqual(indexOf(iclass, i), i)
         self.assertRaises(ValueError, indexOf, iclass, -1)
 
+    # Test iterators with file.writelines().
+    def test_writelines(self):
+        f = file(TESTFN, "w")
+
+        try:
+            self.assertRaises(TypeError, f.writelines, None)
+            self.assertRaises(TypeError, f.writelines, 42)
+    
+            f.writelines(["1\n", "2\n"])
+            f.writelines(("3\n", "4\n"))
+            f.writelines({'5\n': None})
+            f.writelines({})
+
+            # Try a big chunk too.
+            class Iterator:
+                def __init__(self, start, finish):
+                    self.start = start
+                    self.finish = finish
+                    self.i = self.start
+
+                def next(self):
+                    if self.i >= self.finish:
+                        raise StopIteration
+                    result = str(self.i) + '\n'
+                    self.i += 1
+                    return result
+
+                def __iter__(self):
+                    return self
+
+            class Whatever:
+                def __init__(self, start, finish):
+                    self.start = start
+                    self.finish = finish
+
+                def __iter__(self):
+                    return Iterator(self.start, self.finish)
+                    
+            f.writelines(Whatever(6, 6+2000))            
+            f.close()
+
+            f = file(TESTFN)
+            expected = [str(i) + "\n" for i in range(1, 2006)]
+            self.assertEqual(list(f), expected)
+            
+        finally:
+            f.close()
+            try:
+                unlink(TESTFN)
+            except OSError:
+                pass
+
+
     # Test iterators on RHS of unpacking assignments.
     def test_unpack_iter(self):
         a, b = 1, 2
diff --git a/Misc/NEWS b/Misc/NEWS
index cf8e3fc..9a58c38 100644
--- a/Misc/NEWS
+++ b/Misc/NEWS
@@ -3,6 +3,8 @@ What's New in Python 2.2a4?
 
 Core
 
+- file.writelines() now accepts any iterable object producing strings.
+
 - PyUnicode_FromEncodedObject() now works very much like
   PyObject_Str(obj) in that it tries to use __str__/tp_str
   on the object if the object is not a string or buffer. This
diff --git a/Objects/fileobject.c b/Objects/fileobject.c
index b6c039c..d330924 100644
--- a/Objects/fileobject.c
+++ b/Objects/fileobject.c
@@ -1164,55 +1164,54 @@ file_write(PyFileObject *f, PyObject *args)
 }
 
 static PyObject *
-file_writelines(PyFileObject *f, PyObject *args)
+file_writelines(PyFileObject *f, PyObject *seq)
 {
 #define CHUNKSIZE 1000
 	PyObject *list, *line;
+	PyObject *it;	/* iter(seq) */
 	PyObject *result;
 	int i, j, index, len, nwritten, islist;
 
+	assert(seq != NULL);
 	if (f->f_fp == NULL)
 		return err_closed();
-	if (args == NULL || !PySequence_Check(args)) {
-		PyErr_SetString(PyExc_TypeError,
-			   "writelines() argument must be a sequence of strings");
-		return NULL;
-	}
-	islist = PyList_Check(args);
 
-	/* Strategy: slurp CHUNKSIZE lines into a private list,
-	   checking that they are all strings, then write that list
-	   without holding the interpreter lock, then come back for more. */
-	index = 0;
-	if (islist)
-		list = NULL;
+	result = NULL;
+	list = NULL;
+	islist = PyList_Check(seq);
+	if  (islist)
+		it = NULL;
 	else {
+		it = PyObject_GetIter(seq);
+		if (it == NULL) {
+			PyErr_SetString(PyExc_TypeError,
+				"writelines() requires an iterable argument");
+			return NULL;
+		}
+		/* From here on, fail by going to error, to reclaim "it". */
 		list = PyList_New(CHUNKSIZE);
 		if (list == NULL)
-			return NULL;
+			goto error;
 	}
-	result = NULL;
 
-	for (;;) {
+	/* Strategy: slurp CHUNKSIZE lines into a private list,
+	   checking that they are all strings, then write that list
+	   without holding the interpreter lock, then come back for more. */
+	for (index = 0; ; index += CHUNKSIZE) {
 		if (islist) {
 			Py_XDECREF(list);
-			list = PyList_GetSlice(args, index, index+CHUNKSIZE);
+			list = PyList_GetSlice(seq, index, index+CHUNKSIZE);
 			if (list == NULL)
-				return NULL;
+				goto error;
 			j = PyList_GET_SIZE(list);
 		}
 		else {
 			for (j = 0; j < CHUNKSIZE; j++) {
-				line = PySequence_GetItem(args, index+j);
+				line = PyIter_Next(it);
 				if (line == NULL) {
-					if (PyErr_ExceptionMatches(
-						PyExc_IndexError)) {
-						PyErr_Clear();
-						break;
-					}
-					/* Some other error occurred.
-					   XXX We may lose some output. */
-					goto error;
+					if (PyErr_Occurred())
+						goto error;
+					break;
 				}
 				PyList_SetItem(list, j, line);
 			}
@@ -1271,14 +1270,15 @@ file_writelines(PyFileObject *f, PyObject *args)
 
 		if (j < CHUNKSIZE)
 			break;
-		index += CHUNKSIZE;
 	}
 
 	Py_INCREF(Py_None);
 	result = Py_None;
   error:
 	Py_XDECREF(list);
+  	Py_XDECREF(it);
 	return result;
+#undef CHUNKSIZE
 }
 
 static char readline_doc[] =
@@ -1342,10 +1342,10 @@ static char xreadlines_doc[] =
 "often quicker, due to reading ahead internally.";
 
 static char writelines_doc[] =
-"writelines(list of strings) -> None.  Write the strings to the file.\n"
+"writelines(sequence_of_strings) -> None.  Write the strings to the file.\n"
 "\n"
-"Note that newlines are not added.  This is equivalent to calling write()\n"
-"for each string in the list.";
+"Note that newlines are not added.  The sequence can be any iterable object\n"
+"producing strings. This is equivalent to calling write() for each string.";
 
 static char flush_doc[] =
 "flush() -> None.  Flush the internal I/O buffer.";
diff --git a/Tools/scripts/ndiff.py b/Tools/scripts/ndiff.py
index 7ceccc5..6f0f9a9 100755
--- a/Tools/scripts/ndiff.py
+++ b/Tools/scripts/ndiff.py
@@ -118,8 +118,7 @@ def main(args):
 
 def restore(which):
     restored = difflib.restore(sys.stdin.readlines(), which)
-    for line in restored:
-        print line,
+    sys.stdout.writelines(restored)
 
 if __name__ == '__main__':
     args = sys.argv[1:]