Make unicode.join() work nice with iterators. This also required a change

to string.join(), so that when the latter figures out in midstream that
it really needs unicode.join() instead, unicode.join() can actually get
all the sequence elements (i.e., there's no guarantee that the sequence
passed to string.join() can be iterated over *again* by unicode.join(),
so string.join() must not pass on the original sequence object anymore).

2 files changed, 23 insertions, 12 deletions

diff --git a/Objects/stringobject.c b/Objects/stringobject.c
index b905679..87d7c195 100644
--- a/Objects/stringobject.c
+++ b/Objects/stringobject.c
@@ -861,8 +861,15 @@ string_join(PyStringObject *self, PyObject *args)
 		item = PySequence_Fast_GET_ITEM(seq, i);
 		if (!PyString_Check(item)){
 			if (PyUnicode_Check(item)) {
+				/* Defer to Unicode join.
+				 * CAUTION:  There's no gurantee that the
+				 * original sequence can be iterated over
+				 * again, so we must pass seq here.
+				 */
+				PyObject *result;
+				result = PyUnicode_Join((PyObject *)self, seq);
 				Py_DECREF(seq);
-				return PyUnicode_Join((PyObject *)self, orig);
+				return result;
 			}
 			PyErr_Format(PyExc_TypeError,
 				     "sequence item %i: expected string,"
diff --git a/Objects/unicodeobject.c b/Objects/unicodeobject.c
index e52d628..5da4d2f 100644
--- a/Objects/unicodeobject.c
+++ b/Objects/unicodeobject.c
@@ -2724,10 +2724,11 @@ PyObject *PyUnicode_Join(PyObject *separator,
     int seqlen = 0;
     int sz = 100;
     int i;
+    PyObject *it;
 
-    seqlen = PySequence_Size(seq);
-    if (seqlen < 0 && PyErr_Occurred())
-	return NULL;
+    it = PyObject_GetIter(seq);
+    if (it == NULL)
+        return NULL;
 
     if (separator == NULL) {
 	Py_UNICODE blank = ' ';
@@ -2737,7 +2738,7 @@ PyObject *PyUnicode_Join(PyObject *separator,
     else {
 	separator = PyUnicode_FromObject(separator);
 	if (separator == NULL)
-	    return NULL;
+	    goto onError;
 	sep = PyUnicode_AS_UNICODE(separator);
 	seplen = PyUnicode_GET_SIZE(separator);
     }
@@ -2748,13 +2749,14 @@ PyObject *PyUnicode_Join(PyObject *separator,
     p = PyUnicode_AS_UNICODE(res);
     reslen = 0;
 
-    for (i = 0; i < seqlen; i++) {
+    for (i = 0; ; ++i) {
 	int itemlen;
-	PyObject *item;
-
-	item = PySequence_GetItem(seq, i);
-	if (item == NULL)
-	    goto onError;
+	PyObject *item = PyIter_Next(it);
+	if (item == NULL) {
+	    if (PyErr_Occurred())
+		goto onError;
+	    break;
+	}
 	if (!PyUnicode_Check(item)) {
 	    PyObject *v;
 	    v = PyUnicode_FromObject(item);
@@ -2784,11 +2786,13 @@ PyObject *PyUnicode_Join(PyObject *separator,
 	goto onError;
 
     Py_XDECREF(separator);
+    Py_DECREF(it);
     return (PyObject *)res;
 
  onError:
     Py_XDECREF(separator);
-    Py_DECREF(res);
+    Py_XDECREF(res);
+    Py_DECREF(it);
     return NULL;
 }