Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Commit 6b3b0fc

Browse files
rhettingerrhettinger
committed
Forward port r68941 adding itertools.compress().
1 parent ace6733 commit 6b3b0fc

5 files changed

Lines changed: 206 additions & 13 deletions

File tree

‎Doc/library/collections.rst‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -286,7 +286,7 @@ counts less than one::
286286
Section 4.6.3, Exercise 19*\.
287287

288288
* To enumerate all distinct multisets of a given size over a given set of
289-
elements, see the :func:`combinations_with_replacement` function in the
289+
elements, see :func:`combinations_with_replacement` in the
290290
:ref:`itertools-recipes` for itertools::
291291

292292
map(Counter, combinations_with_replacement('ABC', 2)) --> AA AB AC BB BC CC

‎Doc/library/itertools.rst‎

Lines changed: 14 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -133,6 +133,20 @@ loops that truncate the stream.
133133
The number of items returned is ``n! / r! / (n-r)!`` when ``0 <= r <= n``
134134
or zero when ``r > n``.
135135

136+
.. function:: compress(data, selectors)
137+
138+
Make an iterator that filters elements from *data* returning only those that
139+
have a corresponding element in *selectors* that evaluates to ``True``.
140+
Stops when either the *data* or *selectors* iterables have been exhausted.
141+
Equivalent to::
142+
143+
def compress(data, selectors):
144+
# compress('ABCDEF', [1,0,1,0,1,1]) --> A C E F
145+
return (d for d, s in zip(data, selectors) if s)
146+
147+
.. versionadded:: 2.7
148+
149+
136150
.. function:: count([n])
137151

138152
Make an iterator that returns consecutive integers starting with *n*. If not
@@ -594,10 +608,6 @@ which incur interpreter overhead.
594608
s = list(iterable)
595609
return chain.from_iterable(combinations(s, r) for r in range(len(s)+1))
596610

597-
def compress(data, selectors):
598-
"compress('ABCDEF', [1,0,1,0,1,1]) --> A C E F"
599-
return (d for d, s in zip(data, selectors) if s)
600-
601611
def combinations_with_replacement(iterable, r):
602612
"combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC"
603613
# number items returned: (n+r-1)! / r! / (n-1)!

‎Lib/test/test_itertools.py‎

Lines changed: 32 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -195,6 +195,21 @@ def permutations2(iterable, r=None):
195195
self.assertEqual(len(set(map(id, permutations('abcde', 3)))), 1)
196196
self.assertNotEqual(len(set(map(id, list(permutations('abcde', 3))))), 1)
197197

198+
def test_compress(self):
199+
self.assertEqual(list(compress('ABCDEF', [1,0,1,0,1,1])), list('ACEF'))
200+
self.assertEqual(list(compress('ABCDEF', [0,0,0,0,0,0])), list(''))
201+
self.assertEqual(list(compress('ABCDEF', [1,1,1,1,1,1])), list('ABCDEF'))
202+
self.assertEqual(list(compress('ABCDEF', [1,0,1])), list('AC'))
203+
self.assertEqual(list(compress('ABC', [0,1,1,1,1,1])), list('BC'))
204+
n = 10000
205+
data = chain.from_iterable(repeat(range(6), n))
206+
selectors = chain.from_iterable(repeat((0, 1)))
207+
self.assertEqual(list(compress(data, selectors)), [1,3,5] * n)
208+
self.assertRaises(TypeError, compress, None, range(6)) # 1st arg not iterable
209+
self.assertRaises(TypeError, compress, range(6), None) # 2nd arg not iterable
210+
self.assertRaises(TypeError, compress, range(6)) # too few args
211+
self.assertRaises(TypeError, compress, range(6), None) # too many args
212+
198213
def test_count(self):
199214
self.assertEqual(lzip('abc',count()), [('a', 0), ('b', 1), ('c', 2)])
200215
self.assertEqual(lzip('abc',count(3)), [('a', 3), ('b', 4), ('c', 5)])
@@ -715,6 +730,9 @@ def test_combinations(self):
715730
self.assertEqual(list(combinations(range(4), 3)),
716731
[(0,1,2), (0,1,3), (0,2,3), (1,2,3)])
717732

733+
def test_compress(self):
734+
self.assertEqual(list(compress('ABCDEF', [1,0,1,0,1,1])), list('ACEF'))
735+
718736
def test_count(self):
719737
self.assertEqual(list(islice(count(10), 5)), [10, 11, 12, 13, 14])
720738

@@ -795,6 +813,10 @@ def test_combinations(self):
795813
a = []
796814
self.makecycle(combinations([1,2,a,3], 3), a)
797815

816+
def test_compress(self):
817+
a = []
818+
self.makecycle(compress('ABCDEF', [1,0,1,0,1,0]), a)
819+
798820
def test_cycle(self):
799821
a = []
800822
self.makecycle(cycle([a]*2), a)
@@ -948,6 +970,15 @@ def test_chain(self):
948970
self.assertRaises(TypeError, list, chain(N(s)))
949971
self.assertRaises(ZeroDivisionError, list, chain(E(s)))
950972

973+
def test_compress(self):
974+
for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
975+
n = len(s)
976+
for g in (G, I, Ig, S, L, R):
977+
self.assertEqual(list(compress(g(s), repeat(1))), list(g(s)))
978+
self.assertRaises(TypeError, compress, X(s), repeat(1))
979+
self.assertRaises(TypeError, compress, N(s), repeat(1))
980+
self.assertRaises(ZeroDivisionError, list, compress(E(s), repeat(1)))
981+
951982
def test_product(self):
952983
for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
953984
self.assertRaises(TypeError, product, X(s))
@@ -1144,7 +1175,7 @@ class SubclassWithKwargsTest(unittest.TestCase):
11441175
def test_keywords_in_subclass(self):
11451176
# count is not subclassable...
11461177
for cls in (repeat, zip, filter, filterfalse, chain, map,
1147-
starmap, islice, takewhile, dropwhile, cycle):
1178+
starmap, islice, takewhile, dropwhile, cycle, compress):
11481179
class Subclass(cls):
11491180
def __init__(self, newarg=None, *args):
11501181
cls.__init__(self, *args)
@@ -1281,10 +1312,6 @@ def __init__(self, newarg=None, *args):
12811312
... s = list(iterable)
12821313
... return chain.from_iterable(combinations(s, r) for r in range(len(s)+1))
12831314
1284-
>>> def compress(data, selectors):
1285-
... "compress('ABCDEF', [1,0,1,0,1,1]) --> A C E F"
1286-
... return (d for d, s in zip(data, selectors) if s)
1287-
12881315
>>> def combinations_with_replacement(iterable, r):
12891316
... "combinations_with_replacement('ABC', 3) --> AA AB AC BB BC CC"
12901317
... pool = tuple(iterable)
@@ -1380,9 +1407,6 @@ def __init__(self, newarg=None, *args):
13801407
>>> list(powerset([1,2,3]))
13811408
[(), (1,), (2,), (3,), (1, 2), (1, 3), (2, 3), (1, 2, 3)]
13821409
1383-
>>> list(compress('abcdef', [1,0,1,0,1,1]))
1384-
['a', 'c', 'e', 'f']
1385-
13861410
>>> list(combinations_with_replacement('abc', 2))
13871411
[('a', 'a'), ('a', 'b'), ('a', 'c'), ('b', 'b'), ('b', 'c'), ('c', 'c')]
13881412

‎Misc/NEWS‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -150,6 +150,8 @@ Library
150150

151151
- Issue #4863: distutils.mwerkscompiler has been removed.
152152

153+
- Added a new function: itertools.compress().
154+
153155
- Fix and properly document the multiprocessing module's logging
154156
support, expose the internal levels and provide proper usage
155157
examples.

‎Modules/itertoolsmodule.c‎

Lines changed: 157 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2331,6 +2331,162 @@ static PyTypeObject permutations_type = {
23312331
};
23322332

23332333

2334+
/* compress object ************************************************************/
2335+
2336+
/* Equivalent to:
2337+
2338+
def compress(data, selectors):
2339+
"compress('ABCDEF', [1,0,1,0,1,1]) --> A C E F"
2340+
return (d for d, s in zip(data, selectors) if s)
2341+
*/
2342+
2343+
typedef struct {
2344+
PyObject_HEAD
2345+
PyObject *data;
2346+
PyObject *selectors;
2347+
} compressobject;
2348+
2349+
static PyTypeObject compress_type;
2350+
2351+
static PyObject *
2352+
compress_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
2353+
{
2354+
PyObject *seq1, *seq2;
2355+
PyObject *data=NULL, *selectors=NULL;
2356+
compressobject *lz;
2357+
2358+
if (type == &compress_type && !_PyArg_NoKeywords("compress()", kwds))
2359+
return NULL;
2360+
2361+
if (!PyArg_UnpackTuple(args, "compress", 2, 2, &seq1, &seq2))
2362+
return NULL;
2363+
2364+
data = PyObject_GetIter(seq1);
2365+
if (data == NULL)
2366+
goto fail;
2367+
selectors = PyObject_GetIter(seq2);
2368+
if (selectors == NULL)
2369+
goto fail;
2370+
2371+
/* create compressobject structure */
2372+
lz = (compressobject *)type->tp_alloc(type, 0);
2373+
if (lz == NULL)
2374+
goto fail;
2375+
lz->data = data;
2376+
lz->selectors = selectors;
2377+
return (PyObject *)lz;
2378+
2379+
fail:
2380+
Py_XDECREF(data);
2381+
Py_XDECREF(selectors);
2382+
return NULL;
2383+
}
2384+
2385+
static void
2386+
compress_dealloc(compressobject *lz)
2387+
{
2388+
PyObject_GC_UnTrack(lz);
2389+
Py_XDECREF(lz->data);
2390+
Py_XDECREF(lz->selectors);
2391+
Py_TYPE(lz)->tp_free(lz);
2392+
}
2393+
2394+
static int
2395+
compress_traverse(compressobject *lz, visitproc visit, void *arg)
2396+
{
2397+
Py_VISIT(lz->data);
2398+
Py_VISIT(lz->selectors);
2399+
return 0;
2400+
}
2401+
2402+
static PyObject *
2403+
compress_next(compressobject *lz)
2404+
{
2405+
PyObject *data = lz->data, *selectors = lz->selectors;
2406+
PyObject *datum, *selector;
2407+
PyObject *(*datanext)(PyObject *) = *Py_TYPE(data)->tp_iternext;
2408+
PyObject *(*selectornext)(PyObject *) = *Py_TYPE(selectors)->tp_iternext;
2409+
int ok;
2410+
2411+
while (1) {
2412+
/* Steps: get datum, get selector, evaluate selector.
2413+
Order is important (to match the pure python version
2414+
in terms of which input gets a chance to raise an
2415+
exception first).
2416+
*/
2417+
2418+
datum = datanext(data);
2419+
if (datum == NULL)
2420+
return NULL;
2421+
2422+
selector = selectornext(selectors);
2423+
if (selector == NULL) {
2424+
Py_DECREF(datum);
2425+
return NULL;
2426+
}
2427+
2428+
ok = PyObject_IsTrue(selector);
2429+
Py_DECREF(selector);
2430+
if (ok == 1)
2431+
return datum;
2432+
Py_DECREF(datum);
2433+
if (ok == -1)
2434+
return NULL;
2435+
}
2436+
}
2437+
2438+
PyDoc_STRVAR(compress_doc,
2439+
"compress(data sequence, selector sequence) --> iterator over selected data\n\
2440+
\n\
2441+
Return data elements corresponding to true selector elements.\n\
2442+
Forms a shorter iterator from selected data elements using the\n\
2443+
selectors to choose the data elements.");
2444+
2445+
static PyTypeObject compress_type = {
2446+
PyVarObject_HEAD_INIT(NULL, 0)
2447+
"itertools.compress", /* tp_name */
2448+
sizeof(compressobject), /* tp_basicsize */
2449+
0, /* tp_itemsize */
2450+
/* methods */
2451+
(destructor)compress_dealloc, /* tp_dealloc */
2452+
0, /* tp_print */
2453+
0, /* tp_getattr */
2454+
0, /* tp_setattr */
2455+
0, /* tp_compare */
2456+
0, /* tp_repr */
2457+
0, /* tp_as_number */
2458+
0, /* tp_as_sequence */
2459+
0, /* tp_as_mapping */
2460+
0, /* tp_hash */
2461+
0, /* tp_call */
2462+
0, /* tp_str */
2463+
PyObject_GenericGetAttr, /* tp_getattro */
2464+
0, /* tp_setattro */
2465+
0, /* tp_as_buffer */
2466+
Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
2467+
Py_TPFLAGS_BASETYPE, /* tp_flags */
2468+
compress_doc, /* tp_doc */
2469+
(traverseproc)compress_traverse, /* tp_traverse */
2470+
0, /* tp_clear */
2471+
0, /* tp_richcompare */
2472+
0, /* tp_weaklistoffset */
2473+
PyObject_SelfIter, /* tp_iter */
2474+
(iternextfunc)compress_next, /* tp_iternext */
2475+
0, /* tp_methods */
2476+
0, /* tp_members */
2477+
0, /* tp_getset */
2478+
0, /* tp_base */
2479+
0, /* tp_dict */
2480+
0, /* tp_descr_get */
2481+
0, /* tp_descr_set */
2482+
0, /* tp_dictoffset */
2483+
0, /* tp_init */
2484+
0, /* tp_alloc */
2485+
compress_new, /* tp_new */
2486+
PyObject_GC_Del, /* tp_free */
2487+
};
2488+
2489+
23342490
/* filterfalse object ************************************************************/
23352491

23362492
typedef struct {
@@ -3041,6 +3197,7 @@ PyInit_itertools(void)
30413197
&islice_type,
30423198
&starmap_type,
30433199
&chain_type,
3200+
&compress_type,
30443201
&filterfalse_type,
30453202
&count_type,
30463203
&ziplongest_type,

0 commit comments

Comments
 (0)