[Python-checkins] r68961 - in python/branches/py3k: Doc/library/collections.rst Doc/library/itertools.rst Lib/test/test_itertools.py Misc/NEWS Modules/itertoolsmodule.c

Mon Jan 26 03:56:58 CET 2009

Author: raymond.hettinger
Date: Mon Jan 26 03:56:58 2009
New Revision: 68961
Log:
Forward port r68941 adding itertools.compress().
Modified:
 python/branches/py3k/Doc/library/collections.rst
 python/branches/py3k/Doc/library/itertools.rst
 python/branches/py3k/Lib/test/test_itertools.py
 python/branches/py3k/Misc/NEWS
 python/branches/py3k/Modules/itertoolsmodule.c
Modified: python/branches/py3k/Doc/library/collections.rst
==============================================================================

--- python/branches/py3k/Doc/library/collections.rst	(original)
+++ python/branches/py3k/Doc/library/collections.rst	Mon Jan 26 03:56:58 2009
@@ -286,7 +286,7 @@
 Section 4.6.3, Exercise 19*\.
 
 * To enumerate all distinct multisets of a given size over a given set of
- elements, see the :func:`combinations_with_replacement` function in the
+ elements, see :func:`combinations_with_replacement` in the
 :ref:`itertools-recipes` for itertools::
 
 map(Counter, combinations_with_replacement('ABC', 2)) --> AA AB AC BB BC CC
Modified: python/branches/py3k/Doc/library/itertools.rst
==============================================================================
--- python/branches/py3k/Doc/library/itertools.rst	(original)
+++ python/branches/py3k/Doc/library/itertools.rst	Mon Jan 26 03:56:58 2009
@@ -133,6 +133,20 @@
 The number of items returned is ``n! / r! / (n-r)!`` when ``0 <= r <= n``
 or zero when ``r > n``.
 
+.. function:: compress(data, selectors)
+
+ Make an iterator that filters elements from *data* returning only those that
+ have a corresponding element in *selectors* that evaluates to ``True``.
+ Stops when either the *data* or *selectors* iterables have been exhausted.
+ Equivalent to::
+
+ def compress(data, selectors):
+ # compress('ABCDEF', [1,0,1,0,1,1]) --> A C E F
+ return (d for d, s in zip(data, selectors) if s)
+
+ .. versionadded:: 2.7
+
+
 .. function:: count([n])
 
 Make an iterator that returns consecutive integers starting with *n*. If not
@@ -594,10 +608,6 @@
 s = list(iterable)
 return chain.from_iterable(combinations(s, r) for r in range(len(s)+1))
 
- def compress(data, selectors):
- "compress('ABCDEF', [1,0,1,0,1,1]) --> A C E F"
- return (d for d, s in zip(data, selectors) if s)
-
 def combinations_with_replacement(iterable, r):
 "combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC"
 # number items returned: (n+r-1)! / r! / (n-1)!
Modified: python/branches/py3k/Lib/test/test_itertools.py
==============================================================================
--- python/branches/py3k/Lib/test/test_itertools.py	(original)
+++ python/branches/py3k/Lib/test/test_itertools.py	Mon Jan 26 03:56:58 2009
@@ -195,6 +195,21 @@
 self.assertEqual(len(set(map(id, permutations('abcde', 3)))), 1)
 self.assertNotEqual(len(set(map(id, list(permutations('abcde', 3))))), 1)
 
+ def test_compress(self):
+ self.assertEqual(list(compress('ABCDEF', [1,0,1,0,1,1])), list('ACEF'))
+ self.assertEqual(list(compress('ABCDEF', [0,0,0,0,0,0])), list(''))
+ self.assertEqual(list(compress('ABCDEF', [1,1,1,1,1,1])), list('ABCDEF'))
+ self.assertEqual(list(compress('ABCDEF', [1,0,1])), list('AC'))
+ self.assertEqual(list(compress('ABC', [0,1,1,1,1,1])), list('BC'))
+ n = 10000
+ data = chain.from_iterable(repeat(range(6), n))
+ selectors = chain.from_iterable(repeat((0, 1)))
+ self.assertEqual(list(compress(data, selectors)), [1,3,5] * n)
+ self.assertRaises(TypeError, compress, None, range(6)) # 1st arg not iterable
+ self.assertRaises(TypeError, compress, range(6), None) # 2nd arg not iterable
+ self.assertRaises(TypeError, compress, range(6)) # too few args
+ self.assertRaises(TypeError, compress, range(6), None) # too many args
+
 def test_count(self):
 self.assertEqual(lzip('abc',count()), [('a', 0), ('b', 1), ('c', 2)])
 self.assertEqual(lzip('abc',count(3)), [('a', 3), ('b', 4), ('c', 5)])
@@ -715,6 +730,9 @@
 self.assertEqual(list(combinations(range(4), 3)),
 [(0,1,2), (0,1,3), (0,2,3), (1,2,3)])
 
+ def test_compress(self):
+ self.assertEqual(list(compress('ABCDEF', [1,0,1,0,1,1])), list('ACEF'))
+
 def test_count(self):
 self.assertEqual(list(islice(count(10), 5)), [10, 11, 12, 13, 14])
 
@@ -795,6 +813,10 @@
 a = []
 self.makecycle(combinations([1,2,a,3], 3), a)
 
+ def test_compress(self):
+ a = []
+ self.makecycle(compress('ABCDEF', [1,0,1,0,1,0]), a)
+
 def test_cycle(self):
 a = []
 self.makecycle(cycle([a]*2), a)
@@ -948,6 +970,15 @@
 self.assertRaises(TypeError, list, chain(N(s)))
 self.assertRaises(ZeroDivisionError, list, chain(E(s)))
 
+ def test_compress(self):
+ for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
+ n = len(s)
+ for g in (G, I, Ig, S, L, R):
+ self.assertEqual(list(compress(g(s), repeat(1))), list(g(s)))
+ self.assertRaises(TypeError, compress, X(s), repeat(1))
+ self.assertRaises(TypeError, compress, N(s), repeat(1))
+ self.assertRaises(ZeroDivisionError, list, compress(E(s), repeat(1)))
+
 def test_product(self):
 for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
 self.assertRaises(TypeError, product, X(s))
@@ -1144,7 +1175,7 @@
 def test_keywords_in_subclass(self):
 # count is not subclassable...
 for cls in (repeat, zip, filter, filterfalse, chain, map,
- starmap, islice, takewhile, dropwhile, cycle):
+ starmap, islice, takewhile, dropwhile, cycle, compress):
 class Subclass(cls):
 def __init__(self, newarg=None, *args):
 cls.__init__(self, *args)
@@ -1281,10 +1312,6 @@
 ... s = list(iterable)
 ... return chain.from_iterable(combinations(s, r) for r in range(len(s)+1))
 
->>> def compress(data, selectors):
-... "compress('ABCDEF', [1,0,1,0,1,1]) --> A C E F"
-... return (d for d, s in zip(data, selectors) if s)
-
 >>> def combinations_with_replacement(iterable, r):
 ... "combinations_with_replacement('ABC', 3) --> AA AB AC BB BC CC"
 ... pool = tuple(iterable)
@@ -1380,9 +1407,6 @@
 >>> list(powerset([1,2,3]))
 [(), (1,), (2,), (3,), (1, 2), (1, 3), (2, 3), (1, 2, 3)]
 
->>> list(compress('abcdef', [1,0,1,0,1,1]))
-['a', 'c', 'e', 'f']
-
 >>> list(combinations_with_replacement('abc', 2))
 [('a', 'a'), ('a', 'b'), ('a', 'c'), ('b', 'b'), ('b', 'c'), ('c', 'c')]
 
Modified: python/branches/py3k/Misc/NEWS
==============================================================================
--- python/branches/py3k/Misc/NEWS	(original)
+++ python/branches/py3k/Misc/NEWS	Mon Jan 26 03:56:58 2009
@@ -150,6 +150,8 @@
 
 - Issue #4863: distutils.mwerkscompiler has been removed.
 
+- Added a new function: itertools.compress().
+
 - Fix and properly document the multiprocessing module's logging
 support, expose the internal levels and provide proper usage
 examples.
Modified: python/branches/py3k/Modules/itertoolsmodule.c
==============================================================================
--- python/branches/py3k/Modules/itertoolsmodule.c	(original)
+++ python/branches/py3k/Modules/itertoolsmodule.c	Mon Jan 26 03:56:58 2009
@@ -2331,6 +2331,162 @@
 };
 
 
+/* compress object ************************************************************/
+
+/* Equivalent to:
+
+	def compress(data, selectors):
+		"compress('ABCDEF', [1,0,1,0,1,1]) --> A C E F"
+		return (d for d, s in zip(data, selectors) if s)
+*/
+
+typedef struct {
+	PyObject_HEAD
+	PyObject *data;
+	PyObject *selectors;
+} compressobject;
+
+static PyTypeObject compress_type;
+
+static PyObject *
+compress_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
+{
+	PyObject *seq1, *seq2;
+	PyObject *data=NULL, *selectors=NULL;
+	compressobject *lz;
+
+	if (type == &compress_type && !_PyArg_NoKeywords("compress()", kwds))
+		return NULL;
+
+	if (!PyArg_UnpackTuple(args, "compress", 2, 2, &seq1, &seq2))
+		return NULL;
+
+	data = PyObject_GetIter(seq1);
+	if (data == NULL)
+		goto fail;
+	selectors = PyObject_GetIter(seq2);
+	if (selectors == NULL)
+		goto fail;
+
+	/* create compressobject structure */
+	lz = (compressobject *)type->tp_alloc(type, 0);
+	if (lz == NULL)
+		goto fail;
+	lz->data = data;
+	lz->selectors = selectors;
+	return (PyObject *)lz;
+
+fail:
+	Py_XDECREF(data);
+	Py_XDECREF(selectors);
+	return NULL;
+}
+
+static void
+compress_dealloc(compressobject *lz)
+{
+	PyObject_GC_UnTrack(lz);
+	Py_XDECREF(lz->data);
+	Py_XDECREF(lz->selectors);
+	Py_TYPE(lz)->tp_free(lz);
+}
+
+static int
+compress_traverse(compressobject *lz, visitproc visit, void *arg)
+{
+	Py_VISIT(lz->data);
+	Py_VISIT(lz->selectors);
+	return 0;
+}
+
+static PyObject *
+compress_next(compressobject *lz)
+{
+	PyObject *data = lz->data, *selectors = lz->selectors;
+	PyObject *datum, *selector;
+	PyObject *(*datanext)(PyObject *) = *Py_TYPE(data)->tp_iternext;
+	PyObject *(*selectornext)(PyObject *) = *Py_TYPE(selectors)->tp_iternext;
+	int ok;
+
+	while (1) {
+		/* Steps: get datum, get selector, evaluate selector.
+		 Order is important (to match the pure python version
+		 in terms of which input gets a chance to raise an
+		 exception first).
+		*/
+
+		datum = datanext(data);
+		if (datum == NULL)
+			return NULL;
+
+		selector = selectornext(selectors);
+		if (selector == NULL) {
+			Py_DECREF(datum);
+			return NULL;
+		}
+
+		ok = PyObject_IsTrue(selector);
+		Py_DECREF(selector);
+		if (ok == 1)
+			return datum;
+		Py_DECREF(datum);
+		if (ok == -1)
+			return NULL;
+	}
+}
+
+PyDoc_STRVAR(compress_doc,
+"compress(data sequence, selector sequence) --> iterator over selected data\n\
+\n\
+Return data elements corresponding to true selector elements.\n\
+Forms a shorter iterator from selected data elements using the\n\
+selectors to choose the data elements.");
+
+static PyTypeObject compress_type = {
+	PyVarObject_HEAD_INIT(NULL, 0)
+	"itertools.compress",		/* tp_name */
+	sizeof(compressobject),		/* tp_basicsize */
+	0,							/* tp_itemsize */
+	/* methods */
+	(destructor)compress_dealloc,	/* tp_dealloc */
+	0,								/* tp_print */
+	0,								/* tp_getattr */
+	0,								/* tp_setattr */
+	0,								/* tp_compare */
+	0,								/* tp_repr */
+	0,								/* tp_as_number */
+	0,								/* tp_as_sequence */
+	0,								/* tp_as_mapping */
+	0,								/* tp_hash */
+	0,								/* tp_call */
+	0,								/* tp_str */
+	PyObject_GenericGetAttr,		/* tp_getattro */
+	0,								/* tp_setattro */
+	0,								/* tp_as_buffer */
+	Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
+		Py_TPFLAGS_BASETYPE,		/* tp_flags */
+	compress_doc,					/* tp_doc */
+	(traverseproc)compress_traverse,	/* tp_traverse */
+	0,								/* tp_clear */
+	0,								/* tp_richcompare */
+	0,								/* tp_weaklistoffset */
+	PyObject_SelfIter,				/* tp_iter */
+	(iternextfunc)compress_next,	/* tp_iternext */
+	0,								/* tp_methods */
+	0,								/* tp_members */
+	0,								/* tp_getset */
+	0,								/* tp_base */
+	0,								/* tp_dict */
+	0,								/* tp_descr_get */
+	0,								/* tp_descr_set */
+	0,								/* tp_dictoffset */
+	0,								/* tp_init */
+	0,								/* tp_alloc */
+	compress_new,					/* tp_new */
+	PyObject_GC_Del,				/* tp_free */
+};
+
+
 /* filterfalse object ************************************************************/
 
 typedef struct {
@@ -3041,6 +3197,7 @@
 		&islice_type,
 		&starmap_type,
 		&chain_type,
+		&compress_type,
 		&filterfalse_type,
 		&count_type,
 		&ziplongest_type,
