Forward port r68941 adding itertools.compress().

1 files changed, 157 insertions, 0 deletions

diff --git a/Modules/itertoolsmodule.c b/Modules/itertoolsmodule.c
index dcf6aba..bee08de 100644
--- a/Modules/itertoolsmodule.c
+++ b/Modules/itertoolsmodule.c
@@ -2331,6 +2331,162 @@ static PyTypeObject permutations_type = {
 };
 
 
+/* 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 @@ PyInit_itertools(void)
 		&islice_type,
 		&starmap_type,
 		&chain_type,
+		&compress_type,
 		&filterfalse_type,
 		&count_type,
 		&ziplongest_type,