Add alternate constructor for itertools.chain().

3 files changed, 34 insertions, 1 deletions

diff --git a/Lib/test/test_itertools.py b/Lib/test/test_itertools.py
index 41e9362..500afef 100644
--- a/Lib/test/test_itertools.py
+++ b/Lib/test/test_itertools.py
@@ -52,6 +52,13 @@ class TestBasicOps(unittest.TestCase):
         self.assertEqual(take(4, chain('abc', 'def')), list('abcd'))
         self.assertRaises(TypeError, list,chain(2, 3))
 
+    def test_chain_from_iterable(self):
+        self.assertEqual(list(chain.from_iterable(['abc', 'def'])), list('abcdef'))
+        self.assertEqual(list(chain.from_iterable(['abc'])), list('abc'))
+        self.assertEqual(list(chain.from_iterable([''])), [])
+        self.assertEqual(take(4, chain.from_iterable(['abc', 'def'])), list('abcd'))
+        self.assertRaises(TypeError, list, chain.from_iterable([2, 3]))
+
     def test_combinations(self):
         self.assertRaises(TypeError, combinations, 'abc')   # missing r argument
         self.assertRaises(TypeError, combinations, 'abc', 2, 1) # too many arguments
diff --git a/Misc/NEWS b/Misc/NEWS
index c28fd85..6df4ebc 100644
--- a/Misc/NEWS
+++ b/Misc/NEWS
@@ -1247,6 +1247,8 @@ Extension Modules
 - itertools.starmap() now accepts any iterable input. Previously, it required
   the function inputs to be tuples.
 
+- itertools.chain() now has an alterate constructor, chain.from_iterable().
+
 - Issue #1646: Make socket support TIPC. The socket module now has support
   for TIPC under Linux, see http://tipc.sf.net/ for more information.
 
diff --git a/Modules/itertoolsmodule.c b/Modules/itertoolsmodule.c
index 3b8339c..f29077a 100644
--- a/Modules/itertoolsmodule.c
+++ b/Modules/itertoolsmodule.c
@@ -1638,6 +1638,18 @@ chain_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
 	return chain_new_internal(type, source);
 }
 
+static PyObject *
+chain_new_from_iterable(PyTypeObject *type, PyObject *arg)
+{
+	PyObject *source;
+	
+	source = PyObject_GetIter(arg);
+	if (source == NULL)
+		return NULL;
+
+	return chain_new_internal(type, source);
+}
+
 static void
 chain_dealloc(chainobject *lz)
 {
@@ -1696,6 +1708,18 @@ Return a chain object whose .next() method returns elements from the\n\
 first iterable until it is exhausted, then elements from the next\n\
 iterable, until all of the iterables are exhausted.");
 
+PyDoc_STRVAR(chain_from_iterable_doc,
+"chain.from_iterable(iterable) --> chain object\n\
+\n\
+Alternate chain() contructor taking a single iterable argument\n\
+that evaluates lazily.");
+
+static PyMethodDef chain_methods[] = {
+	{"from_iterable", (PyCFunction) chain_new_from_iterable,	METH_O | METH_CLASS,
+		chain_from_iterable_doc},
+	{NULL,		NULL}	/* sentinel */
+};
+
 static PyTypeObject chain_type = {
 	PyVarObject_HEAD_INIT(NULL, 0)
 	"itertools.chain",		/* tp_name */
@@ -1726,7 +1750,7 @@ static PyTypeObject chain_type = {
 	0,				/* tp_weaklistoffset */
 	PyObject_SelfIter,		/* tp_iter */
 	(iternextfunc)chain_next,	/* tp_iternext */
-	0,				/* tp_methods */
+	chain_methods,			/* tp_methods */
 	0,				/* tp_members */
 	0,				/* tp_getset */
 	0,				/* tp_base */