itertoolz.pyx 50 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003100410051006100710081009101010111012101310141015101610171018101910201021102210231024102510261027102810291030103110321033103410351036103710381039104010411042104310441045104610471048104910501051105210531054105510561057105810591060106110621063106410651066106710681069107010711072107310741075107610771078107910801081108210831084108510861087108810891090109110921093109410951096109710981099110011011102110311041105110611071108110911101111111211131114111511161117111811191120112111221123112411251126112711281129113011311132113311341135113611371138113911401141114211431144114511461147114811491150115111521153115411551156115711581159116011611162116311641165116611671168116911701171117211731174117511761177117811791180118111821183118411851186118711881189119011911192119311941195119611971198119912001201120212031204120512061207120812091210121112121213121412151216121712181219122012211222122312241225122612271228122912301231123212331234123512361237123812391240124112421243124412451246124712481249125012511252125312541255125612571258125912601261126212631264126512661267126812691270127112721273127412751276127712781279128012811282128312841285128612871288128912901291129212931294129512961297129812991300130113021303130413051306130713081309131013111312131313141315131613171318131913201321132213231324132513261327132813291330133113321333133413351336133713381339134013411342134313441345134613471348134913501351135213531354135513561357135813591360136113621363136413651366136713681369137013711372137313741375137613771378137913801381138213831384138513861387138813891390139113921393139413951396139713981399140014011402140314041405140614071408140914101411141214131414141514161417141814191420142114221423142414251426142714281429143014311432143314341435143614371438143914401441144214431444144514461447144814491450145114521453145414551456145714581459146014611462146314641465146614671468146914701471147214731474147514761477147814791480148114821483148414851486148714881489149014911492149314941495149614971498149915001501150215031504150515061507150815091510151115121513151415151516151715181519152015211522152315241525152615271528152915301531153215331534153515361537153815391540154115421543154415451546154715481549155015511552155315541555155615571558155915601561156215631564156515661567156815691570157115721573157415751576157715781579158015811582158315841585158615871588158915901591159215931594159515961597159815991600160116021603160416051606160716081609161016111612161316141615161616171618161916201621162216231624162516261627162816291630163116321633163416351636163716381639164016411642164316441645164616471648164916501651165216531654165516561657165816591660166116621663166416651666166716681669167016711672167316741675167616771678167916801681168216831684168516861687168816891690169116921693169416951696169716981699170017011702170317041705170617071708170917101711171217131714171517161717171817191720172117221723172417251726172717281729173017311732173317341735173617371738173917401741174217431744174517461747174817491750175117521753175417551756175717581759176017611762176317641765176617671768176917701771177217731774177517761777177817791780178117821783178417851786178717881789179017911792179317941795179617971798179918001801180218031804180518061807180818091810
  1. from cpython.dict cimport PyDict_GetItem, PyDict_SetItem
  2. from cpython.exc cimport PyErr_Clear, PyErr_GivenExceptionMatches, PyErr_Occurred
  3. from cpython.list cimport PyList_Append, PyList_GET_ITEM, PyList_GET_SIZE
  4. from cpython.object cimport PyObject_RichCompareBool, Py_NE
  5. from cpython.ref cimport PyObject, Py_INCREF, Py_XDECREF
  6. from cpython.sequence cimport PySequence_Check
  7. from cpython.set cimport PySet_Add, PySet_Contains
  8. from cpython.tuple cimport PyTuple_GET_ITEM, PyTuple_GetSlice, PyTuple_New, PyTuple_SET_ITEM
  9. # Locally defined bindings that differ from `cython.cpython` bindings
  10. from cytoolz.cpython cimport PtrIter_Next, PtrObject_GetItem
  11. from collections import deque
  12. from heapq import heapify, heappop, heapreplace
  13. from itertools import chain, islice, zip_longest
  14. from operator import itemgetter
  15. from cytoolz.utils import no_default
  16. __all__ = ['remove', 'accumulate', 'groupby', 'merge_sorted', 'interleave',
  17. 'unique', 'isiterable', 'isdistinct', 'take', 'drop', 'take_nth',
  18. 'first', 'second', 'nth', 'last', 'get', 'concat', 'concatv',
  19. 'mapcat', 'cons', 'interpose', 'frequencies', 'reduceby', 'iterate',
  20. 'sliding_window', 'partition', 'partition_all', 'count', 'pluck',
  21. 'join', 'tail', 'diff', 'topk', 'peek', 'peekn', 'random_sample']
  22. cpdef object identity(object x):
  23. return x
  24. cdef class remove:
  25. """ remove(predicate, seq)
  26. Return those items of sequence for which predicate(item) is False
  27. >>> def iseven(x):
  28. ... return x % 2 == 0
  29. >>> list(remove(iseven, [1, 2, 3, 4]))
  30. [1, 3]
  31. """
  32. def __cinit__(self, object predicate, object seq):
  33. self.predicate = predicate
  34. self.iter_seq = iter(seq)
  35. def __iter__(self):
  36. return self
  37. def __next__(self):
  38. cdef object val
  39. val = next(self.iter_seq)
  40. while self.predicate(val):
  41. val = next(self.iter_seq)
  42. return val
  43. cdef class accumulate:
  44. """ accumulate(binop, seq, initial='__no__default__')
  45. Repeatedly apply binary function to a sequence, accumulating results
  46. >>> from operator import add, mul
  47. >>> list(accumulate(add, [1, 2, 3, 4, 5]))
  48. [1, 3, 6, 10, 15]
  49. >>> list(accumulate(mul, [1, 2, 3, 4, 5]))
  50. [1, 2, 6, 24, 120]
  51. Accumulate is similar to ``reduce`` and is good for making functions like
  52. cumulative sum:
  53. >>> from functools import partial, reduce
  54. >>> sum = partial(reduce, add)
  55. >>> cumsum = partial(accumulate, add)
  56. Accumulate also takes an optional argument that will be used as the first
  57. value. This is similar to reduce.
  58. >>> list(accumulate(add, [1, 2, 3], -1))
  59. [-1, 0, 2, 5]
  60. >>> list(accumulate(add, [], 1))
  61. [1]
  62. See Also:
  63. itertools.accumulate : In standard itertools for Python 3.2+
  64. """
  65. def __cinit__(self, object binop, object seq, object initial='__no__default__'):
  66. self.binop = binop
  67. self.iter_seq = iter(seq)
  68. self.result = self # sentinel
  69. self.initial = initial
  70. def __iter__(self):
  71. return self
  72. def __next__(self):
  73. if self.result is self:
  74. if self.initial != no_default:
  75. self.result = self.initial
  76. else:
  77. self.result = next(self.iter_seq)
  78. else:
  79. self.result = self.binop(self.result, next(self.iter_seq))
  80. return self.result
  81. cdef inline object _groupby_core(dict d, object key, object item):
  82. cdef PyObject *obj = PyDict_GetItem(d, key)
  83. if obj is NULL:
  84. val = []
  85. PyList_Append(val, item)
  86. PyDict_SetItem(d, key, val)
  87. else:
  88. PyList_Append(<object>obj, item)
  89. cpdef dict groupby(object key, object seq):
  90. """
  91. Group a collection by a key function
  92. >>> names = ['Alice', 'Bob', 'Charlie', 'Dan', 'Edith', 'Frank']
  93. >>> groupby(len, names) # doctest: +SKIP
  94. {3: ['Bob', 'Dan'], 5: ['Alice', 'Edith', 'Frank'], 7: ['Charlie']}
  95. >>> iseven = lambda x: x % 2 == 0
  96. >>> groupby(iseven, [1, 2, 3, 4, 5, 6, 7, 8]) # doctest: +SKIP
  97. {False: [1, 3, 5, 7], True: [2, 4, 6, 8]}
  98. Non-callable keys imply grouping on a member.
  99. >>> groupby('gender', [{'name': 'Alice', 'gender': 'F'},
  100. ... {'name': 'Bob', 'gender': 'M'},
  101. ... {'name': 'Charlie', 'gender': 'M'}]) # doctest:+SKIP
  102. {'F': [{'gender': 'F', 'name': 'Alice'}],
  103. 'M': [{'gender': 'M', 'name': 'Bob'},
  104. {'gender': 'M', 'name': 'Charlie'}]}
  105. Not to be confused with ``itertools.groupby``
  106. See Also:
  107. countby
  108. """
  109. cdef dict d = {}
  110. cdef object item, keyval
  111. cdef Py_ssize_t i, N
  112. if callable(key):
  113. for item in seq:
  114. keyval = key(item)
  115. _groupby_core(d, keyval, item)
  116. elif isinstance(key, list):
  117. N = PyList_GET_SIZE(key)
  118. for item in seq:
  119. keyval = PyTuple_New(N)
  120. for i in range(N):
  121. val = <object>PyList_GET_ITEM(key, i)
  122. val = item[val]
  123. Py_INCREF(val)
  124. PyTuple_SET_ITEM(keyval, i, val)
  125. _groupby_core(d, keyval, item)
  126. else:
  127. for item in seq:
  128. keyval = item[key]
  129. _groupby_core(d, keyval, item)
  130. return d
  131. cdef object _merge_sorted_binary(object seqs):
  132. mid = len(seqs) // 2
  133. L1 = seqs[:mid]
  134. if len(L1) == 1:
  135. seq1 = iter(L1[0])
  136. else:
  137. seq1 = _merge_sorted_binary(L1)
  138. L2 = seqs[mid:]
  139. if len(L2) == 1:
  140. seq2 = iter(L2[0])
  141. else:
  142. seq2 = _merge_sorted_binary(L2)
  143. try:
  144. val2 = next(seq2)
  145. except StopIteration:
  146. return seq1
  147. return _merge_sorted(seq1, seq2, val2)
  148. cdef class _merge_sorted:
  149. def __cinit__(self, seq1, seq2, val2):
  150. self.seq1 = seq1
  151. self.seq2 = seq2
  152. self.val1 = None
  153. self.val2 = val2
  154. self.loop = 0
  155. def __iter__(self):
  156. return self
  157. def __next__(self):
  158. if self.loop == 0:
  159. try:
  160. self.val1 = next(self.seq1)
  161. except StopIteration:
  162. self.loop = 2
  163. return self.val2
  164. if self.val2 < self.val1:
  165. self.loop = 1
  166. return self.val2
  167. return self.val1
  168. elif self.loop == 1:
  169. try:
  170. self.val2 = next(self.seq2)
  171. except StopIteration:
  172. self.loop = 3
  173. return self.val1
  174. if self.val2 < self.val1:
  175. return self.val2
  176. self.loop = 0
  177. return self.val1
  178. elif self.loop == 2:
  179. return next(self.seq2)
  180. return next(self.seq1)
  181. cdef object _merge_sorted_binary_key(object seqs, object key):
  182. mid = len(seqs) // 2
  183. L1 = seqs[:mid]
  184. if len(L1) == 1:
  185. seq1 = iter(L1[0])
  186. else:
  187. seq1 = _merge_sorted_binary_key(L1, key)
  188. L2 = seqs[mid:]
  189. if len(L2) == 1:
  190. seq2 = iter(L2[0])
  191. else:
  192. seq2 = _merge_sorted_binary_key(L2, key)
  193. try:
  194. val2 = next(seq2)
  195. except StopIteration:
  196. return seq1
  197. return _merge_sorted_key(seq1, seq2, val2, key)
  198. cdef class _merge_sorted_key:
  199. def __cinit__(self, seq1, seq2, val2, key):
  200. self.seq1 = seq1
  201. self.seq2 = seq2
  202. self.key = key
  203. self.val1 = None
  204. self.key1 = None
  205. self.val2 = val2
  206. self.key2 = key(val2)
  207. self.loop = 0
  208. def __iter__(self):
  209. return self
  210. def __next__(self):
  211. if self.loop == 0:
  212. try:
  213. self.val1 = next(self.seq1)
  214. except StopIteration:
  215. self.loop = 2
  216. return self.val2
  217. self.key1 = self.key(self.val1)
  218. if self.key2 < self.key1:
  219. self.loop = 1
  220. return self.val2
  221. return self.val1
  222. elif self.loop == 1:
  223. try:
  224. self.val2 = next(self.seq2)
  225. except StopIteration:
  226. self.loop = 3
  227. return self.val1
  228. self.key2 = self.key(self.val2)
  229. if self.key2 < self.key1:
  230. return self.val2
  231. self.loop = 0
  232. return self.val1
  233. elif self.loop == 2:
  234. return next(self.seq2)
  235. return next(self.seq1)
  236. cdef object c_merge_sorted(object seqs, object key=None):
  237. if len(seqs) == 0:
  238. return iter([])
  239. elif len(seqs) == 1:
  240. return iter(seqs[0])
  241. elif key is None:
  242. return _merge_sorted_binary(seqs)
  243. return _merge_sorted_binary_key(seqs, key)
  244. def merge_sorted(*seqs, **kwargs):
  245. """
  246. Merge and sort a collection of sorted collections
  247. This works lazily and only keeps one value from each iterable in memory.
  248. >>> list(merge_sorted([1, 3, 5], [2, 4, 6]))
  249. [1, 2, 3, 4, 5, 6]
  250. >>> ''.join(merge_sorted('abc', 'abc', 'abc'))
  251. 'aaabbbccc'
  252. The "key" function used to sort the input may be passed as a keyword.
  253. >>> list(merge_sorted([2, 3], [1, 3], key=lambda x: x // 3))
  254. [2, 1, 3, 3]
  255. """
  256. if 'key' in kwargs:
  257. return c_merge_sorted(seqs, kwargs['key'])
  258. return c_merge_sorted(seqs)
  259. cdef class interleave:
  260. """ interleave(seqs)
  261. Interleave a sequence of sequences
  262. >>> list(interleave([[1, 2], [3, 4]]))
  263. [1, 3, 2, 4]
  264. >>> ''.join(interleave(('ABC', 'XY')))
  265. 'AXBYC'
  266. Both the individual sequences and the sequence of sequences may be infinite
  267. Returns a lazy iterator
  268. """
  269. def __cinit__(self, seqs):
  270. self.iters = [iter(seq) for seq in seqs]
  271. self.newiters = []
  272. self.i = 0
  273. self.n = PyList_GET_SIZE(self.iters)
  274. def __iter__(self):
  275. return self
  276. def __next__(self):
  277. # This implementation is similar to what is done in `toolz` in that we
  278. # construct a new list of iterators, `self.newiters`, when a value is
  279. # successfully retrieved from an iterator from `self.iters`.
  280. cdef PyObject *obj
  281. cdef object val
  282. if self.i == self.n:
  283. self.n = PyList_GET_SIZE(self.newiters)
  284. self.i = 0
  285. if self.n == 0:
  286. raise StopIteration
  287. self.iters = self.newiters
  288. self.newiters = []
  289. val = <object>PyList_GET_ITEM(self.iters, self.i)
  290. self.i += 1
  291. obj = PtrIter_Next(val)
  292. # TODO: optimization opportunity. Previously, it was possible to
  293. # continue on given exceptions, `self.pass_exceptions`, which is
  294. # why this code is structured this way. Time to clean up?
  295. while obj is NULL:
  296. obj = PyErr_Occurred()
  297. if obj is not NULL:
  298. val = <object>obj
  299. PyErr_Clear()
  300. raise val
  301. if self.i == self.n:
  302. self.n = PyList_GET_SIZE(self.newiters)
  303. self.i = 0
  304. if self.n == 0:
  305. raise StopIteration
  306. self.iters = self.newiters
  307. self.newiters = []
  308. val = <object>PyList_GET_ITEM(self.iters, self.i)
  309. self.i += 1
  310. obj = PtrIter_Next(val)
  311. PyList_Append(self.newiters, val)
  312. val = <object>obj
  313. Py_XDECREF(obj)
  314. return val
  315. cdef class _unique_key:
  316. def __cinit__(self, object seq, object key):
  317. self.iter_seq = iter(seq)
  318. self.key = key
  319. self.seen = set()
  320. def __iter__(self):
  321. return self
  322. def __next__(self):
  323. cdef object item, tag
  324. item = next(self.iter_seq)
  325. tag = self.key(item)
  326. while PySet_Contains(self.seen, tag):
  327. item = next(self.iter_seq)
  328. tag = self.key(item)
  329. PySet_Add(self.seen, tag)
  330. return item
  331. cdef class _unique_identity:
  332. def __cinit__(self, object seq):
  333. self.iter_seq = iter(seq)
  334. self.seen = set()
  335. def __iter__(self):
  336. return self
  337. def __next__(self):
  338. cdef object item
  339. item = next(self.iter_seq)
  340. while PySet_Contains(self.seen, item):
  341. item = next(self.iter_seq)
  342. PySet_Add(self.seen, item)
  343. return item
  344. cpdef object unique(object seq, object key=None):
  345. """
  346. Return only unique elements of a sequence
  347. >>> tuple(unique((1, 2, 3)))
  348. (1, 2, 3)
  349. >>> tuple(unique((1, 2, 1, 3)))
  350. (1, 2, 3)
  351. Uniqueness can be defined by key keyword
  352. >>> tuple(unique(['cat', 'mouse', 'dog', 'hen'], key=len))
  353. ('cat', 'mouse')
  354. """
  355. if key is None:
  356. return _unique_identity(seq)
  357. else:
  358. return _unique_key(seq, key)
  359. cpdef object isiterable(object x):
  360. """
  361. Is x iterable?
  362. >>> isiterable([1, 2, 3])
  363. True
  364. >>> isiterable('abc')
  365. True
  366. >>> isiterable(5)
  367. False
  368. """
  369. try:
  370. iter(x)
  371. return True
  372. except TypeError:
  373. pass
  374. return False
  375. cpdef object isdistinct(object seq):
  376. """
  377. All values in sequence are distinct
  378. >>> isdistinct([1, 2, 3])
  379. True
  380. >>> isdistinct([1, 2, 1])
  381. False
  382. >>> isdistinct("Hello")
  383. False
  384. >>> isdistinct("World")
  385. True
  386. """
  387. if iter(seq) is seq:
  388. seen = set()
  389. for item in seq:
  390. if PySet_Contains(seen, item):
  391. return False
  392. seen.add(item)
  393. return True
  394. else:
  395. return len(seq) == len(set(seq))
  396. cpdef object take(Py_ssize_t n, object seq):
  397. """
  398. The first n elements of a sequence
  399. >>> list(take(2, [10, 20, 30, 40, 50]))
  400. [10, 20]
  401. See Also:
  402. drop
  403. tail
  404. """
  405. return islice(seq, n)
  406. cpdef object tail(Py_ssize_t n, object seq):
  407. """
  408. The last n elements of a sequence
  409. >>> tail(2, [10, 20, 30, 40, 50])
  410. [40, 50]
  411. See Also:
  412. drop
  413. take
  414. """
  415. if PySequence_Check(seq):
  416. return seq[-n:]
  417. return tuple(deque(seq, n))
  418. cpdef object drop(Py_ssize_t n, object seq):
  419. """
  420. The sequence following the first n elements
  421. >>> list(drop(2, [10, 20, 30, 40, 50]))
  422. [30, 40, 50]
  423. See Also:
  424. take
  425. tail
  426. """
  427. if n < 0:
  428. raise ValueError('n argument for drop() must be non-negative')
  429. cdef Py_ssize_t i
  430. cdef object iter_seq
  431. iter_seq = iter(seq)
  432. try:
  433. for i in range(n):
  434. next(iter_seq)
  435. except StopIteration:
  436. pass
  437. return iter_seq
  438. cpdef object take_nth(Py_ssize_t n, object seq):
  439. """
  440. Every nth item in seq
  441. >>> list(take_nth(2, [10, 20, 30, 40, 50]))
  442. [10, 30, 50]
  443. """
  444. return islice(seq, 0, None, n)
  445. cpdef object first(object seq):
  446. """
  447. The first element in a sequence
  448. >>> first('ABC')
  449. 'A'
  450. """
  451. return next(iter(seq))
  452. cpdef object second(object seq):
  453. """
  454. The second element in a sequence
  455. >>> second('ABC')
  456. 'B'
  457. """
  458. seq = iter(seq)
  459. next(seq)
  460. return next(seq)
  461. cpdef object nth(Py_ssize_t n, object seq):
  462. """
  463. The nth element in a sequence
  464. >>> nth(1, 'ABC')
  465. 'B'
  466. """
  467. if PySequence_Check(seq):
  468. return seq[n]
  469. if n < 0:
  470. raise ValueError('"n" must be positive when indexing an iterator')
  471. seq = iter(seq)
  472. while n > 0:
  473. n -= 1
  474. next(seq)
  475. return next(seq)
  476. cpdef object last(object seq):
  477. """
  478. The last element in a sequence
  479. >>> last('ABC')
  480. 'C'
  481. """
  482. cdef object val
  483. if PySequence_Check(seq):
  484. return seq[-1]
  485. val = no_default
  486. for val in seq:
  487. pass
  488. if val == no_default:
  489. raise IndexError
  490. return val
  491. cpdef object rest(object seq):
  492. seq = iter(seq)
  493. next(seq)
  494. return seq
  495. cdef tuple _get_exceptions = (IndexError, KeyError, TypeError)
  496. cdef tuple _get_list_exc = (IndexError, KeyError)
  497. cpdef object get(object ind, object seq, object default='__no__default__'):
  498. """
  499. Get element in a sequence or dict
  500. Provides standard indexing
  501. >>> get(1, 'ABC') # Same as 'ABC'[1]
  502. 'B'
  503. Pass a list to get multiple values
  504. >>> get([1, 2], 'ABC') # ('ABC'[1], 'ABC'[2])
  505. ('B', 'C')
  506. Works on any value that supports indexing/getitem
  507. For example here we see that it works with dictionaries
  508. >>> phonebook = {'Alice': '555-1234',
  509. ... 'Bob': '555-5678',
  510. ... 'Charlie':'555-9999'}
  511. >>> get('Alice', phonebook)
  512. '555-1234'
  513. >>> get(['Alice', 'Bob'], phonebook)
  514. ('555-1234', '555-5678')
  515. Provide a default for missing values
  516. >>> get(['Alice', 'Dennis'], phonebook, None)
  517. ('555-1234', None)
  518. See Also:
  519. pluck
  520. """
  521. cdef Py_ssize_t i
  522. cdef object val
  523. cdef tuple result
  524. cdef PyObject *obj
  525. if isinstance(ind, list):
  526. i = PyList_GET_SIZE(ind)
  527. result = PyTuple_New(i)
  528. # List of indices, no default
  529. if default == no_default:
  530. for i, val in enumerate(ind):
  531. val = seq[val]
  532. Py_INCREF(val)
  533. PyTuple_SET_ITEM(result, i, val)
  534. return result
  535. # List of indices with default
  536. for i, val in enumerate(ind):
  537. obj = PtrObject_GetItem(seq, val)
  538. if obj is NULL:
  539. val = <object>PyErr_Occurred()
  540. PyErr_Clear()
  541. if not PyErr_GivenExceptionMatches(val, _get_list_exc):
  542. raise val
  543. Py_INCREF(default)
  544. PyTuple_SET_ITEM(result, i, default)
  545. else:
  546. val = <object>obj
  547. PyTuple_SET_ITEM(result, i, val)
  548. return result
  549. obj = PtrObject_GetItem(seq, ind)
  550. if obj is NULL:
  551. val = <object>PyErr_Occurred()
  552. PyErr_Clear()
  553. if default == no_default:
  554. raise val
  555. if PyErr_GivenExceptionMatches(val, _get_exceptions):
  556. return default
  557. raise val
  558. val = <object>obj
  559. Py_XDECREF(obj)
  560. return val
  561. cpdef object concat(object seqs):
  562. """
  563. Concatenate zero or more iterables, any of which may be infinite.
  564. An infinite sequence will prevent the rest of the arguments from
  565. being included.
  566. We use chain.from_iterable rather than ``chain(*seqs)`` so that seqs
  567. can be a generator.
  568. >>> list(concat([[], [1], [2, 3]]))
  569. [1, 2, 3]
  570. See also:
  571. itertools.chain.from_iterable equivalent
  572. """
  573. return chain.from_iterable(seqs)
  574. def concatv(*seqs):
  575. """
  576. Variadic version of concat
  577. >>> list(concatv([], ["a"], ["b", "c"]))
  578. ['a', 'b', 'c']
  579. See also:
  580. itertools.chain
  581. """
  582. return chain.from_iterable(seqs)
  583. cpdef object mapcat(object func, object seqs):
  584. """
  585. Apply func to each sequence in seqs, concatenating results.
  586. >>> list(mapcat(lambda s: [c.upper() for c in s],
  587. ... [["a", "b"], ["c", "d", "e"]]))
  588. ['A', 'B', 'C', 'D', 'E']
  589. """
  590. return concat(map(func, seqs))
  591. cpdef object cons(object el, object seq):
  592. """
  593. Add el to beginning of (possibly infinite) sequence seq.
  594. >>> list(cons(1, [2, 3]))
  595. [1, 2, 3]
  596. """
  597. return chain((el,), seq)
  598. cdef class interpose:
  599. """ interpose(el, seq)
  600. Introduce element between each pair of elements in seq
  601. >>> list(interpose("a", [1, 2, 3]))
  602. [1, 'a', 2, 'a', 3]
  603. """
  604. def __cinit__(self, object el, object seq):
  605. self.el = el
  606. self.iter_seq = iter(seq)
  607. self.do_el = False
  608. try:
  609. self.val = next(self.iter_seq)
  610. except StopIteration:
  611. self.do_el = True
  612. def __iter__(self):
  613. return self
  614. def __next__(self):
  615. if self.do_el:
  616. self.val = next(self.iter_seq)
  617. self.do_el = False
  618. return self.el
  619. else:
  620. self.do_el = True
  621. return self.val
  622. cpdef dict frequencies(object seq):
  623. """
  624. Find number of occurrences of each value in seq
  625. >>> frequencies(['cat', 'cat', 'ox', 'pig', 'pig', 'cat']) #doctest: +SKIP
  626. {'cat': 3, 'ox': 1, 'pig': 2}
  627. See Also:
  628. countby
  629. groupby
  630. """
  631. cdef dict d = {}
  632. cdef PyObject *obj
  633. cdef Py_ssize_t val
  634. for item in seq:
  635. obj = PyDict_GetItem(d, item)
  636. if obj is NULL:
  637. d[item] = 1
  638. else:
  639. val = <object>obj
  640. d[item] = val + 1
  641. return d
  642. cdef inline object _reduceby_core(dict d, object key, object item, object binop,
  643. object init, bint skip_init, bint call_init):
  644. cdef PyObject *obj = PyDict_GetItem(d, key)
  645. if obj is not NULL:
  646. PyDict_SetItem(d, key, binop(<object>obj, item))
  647. elif skip_init:
  648. PyDict_SetItem(d, key, item)
  649. elif call_init:
  650. PyDict_SetItem(d, key, binop(init(), item))
  651. else:
  652. PyDict_SetItem(d, key, binop(init, item))
  653. cpdef dict reduceby(object key, object binop, object seq, object init='__no__default__'):
  654. """
  655. Perform a simultaneous groupby and reduction
  656. The computation:
  657. >>> result = reduceby(key, binop, seq, init) # doctest: +SKIP
  658. is equivalent to the following:
  659. >>> def reduction(group): # doctest: +SKIP
  660. ... return reduce(binop, group, init) # doctest: +SKIP
  661. >>> groups = groupby(key, seq) # doctest: +SKIP
  662. >>> result = valmap(reduction, groups) # doctest: +SKIP
  663. But the former does not build the intermediate groups, allowing it to
  664. operate in much less space. This makes it suitable for larger datasets
  665. that do not fit comfortably in memory
  666. The ``init`` keyword argument is the default initialization of the
  667. reduction. This can be either a constant value like ``0`` or a callable
  668. like ``lambda : 0`` as might be used in ``defaultdict``.
  669. Simple Examples
  670. ---------------
  671. >>> from operator import add, mul
  672. >>> iseven = lambda x: x % 2 == 0
  673. >>> data = [1, 2, 3, 4, 5]
  674. >>> reduceby(iseven, add, data) # doctest: +SKIP
  675. {False: 9, True: 6}
  676. >>> reduceby(iseven, mul, data) # doctest: +SKIP
  677. {False: 15, True: 8}
  678. Complex Example
  679. ---------------
  680. >>> projects = [{'name': 'build roads', 'state': 'CA', 'cost': 1000000},
  681. ... {'name': 'fight crime', 'state': 'IL', 'cost': 100000},
  682. ... {'name': 'help farmers', 'state': 'IL', 'cost': 2000000},
  683. ... {'name': 'help farmers', 'state': 'CA', 'cost': 200000}]
  684. >>> reduceby('state', # doctest: +SKIP
  685. ... lambda acc, x: acc + x['cost'],
  686. ... projects, 0)
  687. {'CA': 1200000, 'IL': 2100000}
  688. Example Using ``init``
  689. ----------------------
  690. >>> def set_add(s, i):
  691. ... s.add(i)
  692. ... return s
  693. >>> reduceby(iseven, set_add, [1, 2, 3, 4, 1, 2, 3], set) # doctest: +SKIP
  694. {True: set([2, 4]),
  695. False: set([1, 3])}
  696. """
  697. cdef dict d = {}
  698. cdef object item, keyval
  699. cdef Py_ssize_t i, N
  700. cdef bint skip_init = init == no_default
  701. cdef bint call_init = callable(init)
  702. if callable(key):
  703. for item in seq:
  704. keyval = key(item)
  705. _reduceby_core(d, keyval, item, binop, init, skip_init, call_init)
  706. elif isinstance(key, list):
  707. N = PyList_GET_SIZE(key)
  708. for item in seq:
  709. keyval = PyTuple_New(N)
  710. for i in range(N):
  711. val = <object>PyList_GET_ITEM(key, i)
  712. val = item[val]
  713. Py_INCREF(val)
  714. PyTuple_SET_ITEM(keyval, i, val)
  715. _reduceby_core(d, keyval, item, binop, init, skip_init, call_init)
  716. else:
  717. for item in seq:
  718. keyval = item[key]
  719. _reduceby_core(d, keyval, item, binop, init, skip_init, call_init)
  720. return d
  721. cdef class iterate:
  722. """ iterate(func, x)
  723. Repeatedly apply a function func onto an original input
  724. Yields x, then func(x), then func(func(x)), then func(func(func(x))), etc..
  725. >>> def inc(x): return x + 1
  726. >>> counter = iterate(inc, 0)
  727. >>> next(counter)
  728. 0
  729. >>> next(counter)
  730. 1
  731. >>> next(counter)
  732. 2
  733. >>> double = lambda x: x * 2
  734. >>> powers_of_two = iterate(double, 1)
  735. >>> next(powers_of_two)
  736. 1
  737. >>> next(powers_of_two)
  738. 2
  739. >>> next(powers_of_two)
  740. 4
  741. >>> next(powers_of_two)
  742. 8
  743. """
  744. def __cinit__(self, object func, object x):
  745. self.func = func
  746. self.x = x
  747. self.val = self # sentinel
  748. def __iter__(self):
  749. return self
  750. def __next__(self):
  751. if self.val is self:
  752. self.val = self.x
  753. else:
  754. self.x = self.func(self.x)
  755. return self.x
  756. cdef class sliding_window:
  757. """ sliding_window(n, seq)
  758. A sequence of overlapping subsequences
  759. >>> list(sliding_window(2, [1, 2, 3, 4]))
  760. [(1, 2), (2, 3), (3, 4)]
  761. This function creates a sliding window suitable for transformations like
  762. sliding means / smoothing
  763. >>> mean = lambda seq: float(sum(seq)) / len(seq)
  764. >>> list(map(mean, sliding_window(2, [1, 2, 3, 4])))
  765. [1.5, 2.5, 3.5]
  766. """
  767. def __cinit__(self, Py_ssize_t n, object seq):
  768. cdef Py_ssize_t i
  769. self.iterseq = iter(seq)
  770. self.prev = PyTuple_New(n)
  771. Py_INCREF(None)
  772. PyTuple_SET_ITEM(self.prev, 0, None)
  773. for i, seq in enumerate(islice(self.iterseq, n-1), 1):
  774. Py_INCREF(seq)
  775. PyTuple_SET_ITEM(self.prev, i, seq)
  776. self.n = n
  777. def __iter__(self):
  778. return self
  779. def __next__(self):
  780. cdef tuple current
  781. cdef object item
  782. cdef Py_ssize_t i
  783. item = next(self.iterseq)
  784. current = PyTuple_New(self.n)
  785. Py_INCREF(item)
  786. PyTuple_SET_ITEM(current, self.n-1, item)
  787. for i in range(1, self.n):
  788. item = self.prev[i]
  789. Py_INCREF(item)
  790. PyTuple_SET_ITEM(current, i-1, item)
  791. self.prev = current
  792. return current
  793. no_pad = '__no__pad__'
  794. cpdef object partition(Py_ssize_t n, object seq, object pad='__no__pad__'):
  795. """
  796. Partition sequence into tuples of length n
  797. >>> list(partition(2, [1, 2, 3, 4]))
  798. [(1, 2), (3, 4)]
  799. If the length of ``seq`` is not evenly divisible by ``n``, the final tuple
  800. is dropped if ``pad`` is not specified, or filled to length ``n`` by pad:
  801. >>> list(partition(2, [1, 2, 3, 4, 5]))
  802. [(1, 2), (3, 4)]
  803. >>> list(partition(2, [1, 2, 3, 4, 5], pad=None))
  804. [(1, 2), (3, 4), (5, None)]
  805. See Also:
  806. partition_all
  807. """
  808. args = [iter(seq)] * n
  809. if pad == '__no__pad__':
  810. return zip(*args)
  811. else:
  812. return zip_longest(*args, fillvalue=pad)
  813. cdef class partition_all:
  814. """ partition_all(n, seq)
  815. Partition all elements of sequence into tuples of length at most n
  816. The final tuple may be shorter to accommodate extra elements.
  817. >>> list(partition_all(2, [1, 2, 3, 4]))
  818. [(1, 2), (3, 4)]
  819. >>> list(partition_all(2, [1, 2, 3, 4, 5]))
  820. [(1, 2), (3, 4), (5,)]
  821. See Also:
  822. partition
  823. """
  824. def __cinit__(self, Py_ssize_t n, object seq):
  825. self.n = n
  826. self.iterseq = iter(seq)
  827. def __iter__(self):
  828. return self
  829. def __next__(self):
  830. cdef tuple result
  831. cdef object item
  832. cdef Py_ssize_t i = 0
  833. result = PyTuple_New(self.n)
  834. for item in self.iterseq:
  835. Py_INCREF(item)
  836. PyTuple_SET_ITEM(result, i, item)
  837. i += 1
  838. if i == self.n:
  839. return result
  840. # iterable exhausted before filling the tuple
  841. if i == 0:
  842. raise StopIteration
  843. for j in range(i, self.n):
  844. Py_INCREF(None)
  845. PyTuple_SET_ITEM(result, j, None)
  846. return PyTuple_GetSlice(result, 0, i)
  847. cpdef object count(object seq):
  848. """
  849. Count the number of items in seq
  850. Like the builtin ``len`` but works on lazy sequences.
  851. Not to be confused with ``itertools.count``
  852. See also:
  853. len
  854. """
  855. if iter(seq) is not seq and hasattr(seq, '__len__'):
  856. return len(seq)
  857. cdef Py_ssize_t i = 0
  858. for _ in seq:
  859. i += 1
  860. return i
  861. cdef class _pluck_index:
  862. def __cinit__(self, object ind, object seqs):
  863. self.ind = ind
  864. self.iterseqs = iter(seqs)
  865. def __iter__(self):
  866. return self
  867. def __next__(self):
  868. val = next(self.iterseqs)
  869. return val[self.ind]
  870. cdef class _pluck_index_default:
  871. def __cinit__(self, object ind, object seqs, object default):
  872. self.ind = ind
  873. self.iterseqs = iter(seqs)
  874. self.default = default
  875. def __iter__(self):
  876. return self
  877. def __next__(self):
  878. cdef PyObject *obj
  879. cdef object val
  880. val = next(self.iterseqs)
  881. obj = PtrObject_GetItem(val, self.ind)
  882. if obj is NULL:
  883. val = <object>PyErr_Occurred()
  884. PyErr_Clear()
  885. if not PyErr_GivenExceptionMatches(val, _get_exceptions):
  886. raise val
  887. return self.default
  888. val = <object>obj
  889. Py_XDECREF(obj)
  890. return val
  891. cdef class _pluck_list:
  892. def __cinit__(self, list ind not None, object seqs):
  893. self.ind = ind
  894. self.iterseqs = iter(seqs)
  895. self.n = len(ind)
  896. def __iter__(self):
  897. return self
  898. def __next__(self):
  899. cdef Py_ssize_t i
  900. cdef tuple result
  901. cdef object val, seq
  902. seq = next(self.iterseqs)
  903. result = PyTuple_New(self.n)
  904. for i, val in enumerate(self.ind):
  905. val = seq[val]
  906. Py_INCREF(val)
  907. PyTuple_SET_ITEM(result, i, val)
  908. return result
  909. cdef class _pluck_list_default:
  910. def __cinit__(self, list ind not None, object seqs, object default):
  911. self.ind = ind
  912. self.iterseqs = iter(seqs)
  913. self.default = default
  914. self.n = len(ind)
  915. def __iter__(self):
  916. return self
  917. def __next__(self):
  918. cdef Py_ssize_t i
  919. cdef object val, seq
  920. cdef tuple result
  921. seq = next(self.iterseqs)
  922. result = PyTuple_New(self.n)
  923. for i, val in enumerate(self.ind):
  924. obj = PtrObject_GetItem(seq, val)
  925. if obj is NULL:
  926. val = <object>PyErr_Occurred()
  927. PyErr_Clear()
  928. if not PyErr_GivenExceptionMatches(val, _get_list_exc):
  929. raise val
  930. Py_INCREF(self.default)
  931. PyTuple_SET_ITEM(result, i, self.default)
  932. else:
  933. val = <object>obj
  934. PyTuple_SET_ITEM(result, i, val)
  935. return result
  936. cpdef object pluck(object ind, object seqs, object default='__no__default__'):
  937. """
  938. plucks an element or several elements from each item in a sequence.
  939. ``pluck`` maps ``itertoolz.get`` over a sequence and returns one or more
  940. elements of each item in the sequence.
  941. This is equivalent to running `map(curried.get(ind), seqs)`
  942. ``ind`` can be either a single string/index or a list of strings/indices.
  943. ``seqs`` should be sequence containing sequences or dicts.
  944. e.g.
  945. >>> data = [{'id': 1, 'name': 'Cheese'}, {'id': 2, 'name': 'Pies'}]
  946. >>> list(pluck('name', data))
  947. ['Cheese', 'Pies']
  948. >>> list(pluck([0, 1], [[1, 2, 3], [4, 5, 7]]))
  949. [(1, 2), (4, 5)]
  950. See Also:
  951. get
  952. map
  953. """
  954. if isinstance(ind, list):
  955. if default != no_default:
  956. return _pluck_list_default(ind, seqs, default)
  957. if PyList_GET_SIZE(ind) < 10:
  958. return _pluck_list(ind, seqs)
  959. return map(itemgetter(*ind), seqs)
  960. if default == no_default:
  961. return _pluck_index(ind, seqs)
  962. return _pluck_index_default(ind, seqs, default)
  963. cdef class _getter_index:
  964. def __cinit__(self, object ind):
  965. self.ind = ind
  966. def __call__(self, object seq):
  967. return seq[self.ind]
  968. cdef class _getter_list:
  969. def __cinit__(self, list ind not None):
  970. self.ind = ind
  971. self.n = len(ind)
  972. def __call__(self, object seq):
  973. cdef Py_ssize_t i
  974. cdef tuple result
  975. cdef object val
  976. result = PyTuple_New(self.n)
  977. for i, val in enumerate(self.ind):
  978. val = seq[val]
  979. Py_INCREF(val)
  980. PyTuple_SET_ITEM(result, i, val)
  981. return result
  982. cdef class _getter_null:
  983. def __call__(self, object seq):
  984. return ()
  985. # TODO: benchmark getters (and compare against itemgetter)
  986. cpdef object getter(object index):
  987. if isinstance(index, list):
  988. if PyList_GET_SIZE(index) == 0:
  989. return _getter_null()
  990. elif PyList_GET_SIZE(index) < 10:
  991. return _getter_list(index)
  992. return itemgetter(*index)
  993. return _getter_index(index)
  994. cpdef object join(object leftkey, object leftseq,
  995. object rightkey, object rightseq,
  996. object left_default='__no__default__',
  997. object right_default='__no__default__'):
  998. """
  999. Join two sequences on common attributes
  1000. This is a semi-streaming operation. The LEFT sequence is fully evaluated
  1001. and placed into memory. The RIGHT sequence is evaluated lazily and so can
  1002. be arbitrarily large.
  1003. (Note: If right_default is defined, then unique keys of rightseq
  1004. will also be stored in memory.)
  1005. >>> friends = [('Alice', 'Edith'),
  1006. ... ('Alice', 'Zhao'),
  1007. ... ('Edith', 'Alice'),
  1008. ... ('Zhao', 'Alice'),
  1009. ... ('Zhao', 'Edith')]
  1010. >>> cities = [('Alice', 'NYC'),
  1011. ... ('Alice', 'Chicago'),
  1012. ... ('Dan', 'Sydney'),
  1013. ... ('Edith', 'Paris'),
  1014. ... ('Edith', 'Berlin'),
  1015. ... ('Zhao', 'Shanghai')]
  1016. >>> # Vacation opportunities
  1017. >>> # In what cities do people have friends?
  1018. >>> result = join(second, friends,
  1019. ... first, cities)
  1020. >>> for ((a, b), (c, d)) in sorted(unique(result)):
  1021. ... print((a, d))
  1022. ('Alice', 'Berlin')
  1023. ('Alice', 'Paris')
  1024. ('Alice', 'Shanghai')
  1025. ('Edith', 'Chicago')
  1026. ('Edith', 'NYC')
  1027. ('Zhao', 'Chicago')
  1028. ('Zhao', 'NYC')
  1029. ('Zhao', 'Berlin')
  1030. ('Zhao', 'Paris')
  1031. Specify outer joins with keyword arguments ``left_default`` and/or
  1032. ``right_default``. Here is a full outer join in which unmatched elements
  1033. are paired with None.
  1034. >>> identity = lambda x: x
  1035. >>> list(join(identity, [1, 2, 3],
  1036. ... identity, [2, 3, 4],
  1037. ... left_default=None, right_default=None))
  1038. [(2, 2), (3, 3), (None, 4), (1, None)]
  1039. Usually the key arguments are callables to be applied to the sequences. If
  1040. the keys are not obviously callable then it is assumed that indexing was
  1041. intended, e.g. the following is a legal change.
  1042. The join is implemented as a hash join and the keys of leftseq must be
  1043. hashable. Additionally, if right_default is defined, then keys of rightseq
  1044. must also be hashable.
  1045. >>> # result = join(second, friends, first, cities)
  1046. >>> result = join(1, friends, 0, cities) # doctest: +SKIP
  1047. """
  1048. if left_default == no_default and right_default == no_default:
  1049. if callable(rightkey):
  1050. return _inner_join_key(leftkey, leftseq, rightkey, rightseq,
  1051. left_default, right_default)
  1052. elif isinstance(rightkey, list):
  1053. return _inner_join_indices(leftkey, leftseq, rightkey, rightseq,
  1054. left_default, right_default)
  1055. else:
  1056. return _inner_join_index(leftkey, leftseq, rightkey, rightseq,
  1057. left_default, right_default)
  1058. elif left_default != no_default and right_default == no_default:
  1059. if callable(rightkey):
  1060. return _right_outer_join_key(leftkey, leftseq, rightkey, rightseq,
  1061. left_default, right_default)
  1062. elif isinstance(rightkey, list):
  1063. return _right_outer_join_indices(leftkey, leftseq, rightkey, rightseq,
  1064. left_default, right_default)
  1065. else:
  1066. return _right_outer_join_index(leftkey, leftseq, rightkey, rightseq,
  1067. left_default, right_default)
  1068. elif left_default == no_default and right_default != no_default:
  1069. if callable(rightkey):
  1070. return _left_outer_join_key(leftkey, leftseq, rightkey, rightseq,
  1071. left_default, right_default)
  1072. elif isinstance(rightkey, list):
  1073. return _left_outer_join_indices(leftkey, leftseq, rightkey, rightseq,
  1074. left_default, right_default)
  1075. else:
  1076. return _left_outer_join_index(leftkey, leftseq, rightkey, rightseq,
  1077. left_default, right_default)
  1078. else:
  1079. if callable(rightkey):
  1080. return _outer_join_key(leftkey, leftseq, rightkey, rightseq,
  1081. left_default, right_default)
  1082. elif isinstance(rightkey, list):
  1083. return _outer_join_indices(leftkey, leftseq, rightkey, rightseq,
  1084. left_default, right_default)
  1085. else:
  1086. return _outer_join_index(leftkey, leftseq, rightkey, rightseq,
  1087. left_default, right_default)
  1088. cdef class _join:
  1089. def __cinit__(self,
  1090. object leftkey, object leftseq,
  1091. object rightkey, object rightseq,
  1092. object left_default=no_default,
  1093. object right_default=no_default):
  1094. self.left_default = left_default
  1095. self.right_default = right_default
  1096. self._rightkey = rightkey
  1097. self.rightseq = iter(rightseq)
  1098. if isinstance(rightkey, list):
  1099. self.N = len(rightkey)
  1100. self.d = groupby(leftkey, leftseq)
  1101. self.seen_keys = set()
  1102. self.matches = []
  1103. self.right = None
  1104. self.is_rightseq_exhausted = False
  1105. def __iter__(self):
  1106. return self
  1107. cdef object rightkey(self):
  1108. pass
  1109. cdef class _right_outer_join(_join):
  1110. def __next__(self):
  1111. cdef PyObject *obj
  1112. if self.i == PyList_GET_SIZE(self.matches):
  1113. self.right = next(self.rightseq)
  1114. key = self.rightkey()
  1115. obj = PyDict_GetItem(self.d, key)
  1116. if obj is NULL:
  1117. return (self.left_default, self.right)
  1118. self.matches = <object>obj
  1119. self.i = 0
  1120. match = <object>PyList_GET_ITEM(self.matches, self.i) # skip error checking
  1121. self.i += 1
  1122. return (match, self.right)
  1123. cdef class _right_outer_join_key(_right_outer_join):
  1124. cdef object rightkey(self):
  1125. return self._rightkey(self.right)
  1126. cdef class _right_outer_join_index(_right_outer_join):
  1127. cdef object rightkey(self):
  1128. return self.right[self._rightkey]
  1129. cdef class _right_outer_join_indices(_right_outer_join):
  1130. cdef object rightkey(self):
  1131. keyval = PyTuple_New(self.N)
  1132. for i in range(self.N):
  1133. val = <object>PyList_GET_ITEM(self._rightkey, i)
  1134. val = self.right[val]
  1135. Py_INCREF(val)
  1136. PyTuple_SET_ITEM(keyval, i, val)
  1137. return keyval
  1138. cdef class _outer_join(_join):
  1139. def __next__(self):
  1140. cdef PyObject *obj
  1141. if not self.is_rightseq_exhausted:
  1142. if self.i == PyList_GET_SIZE(self.matches):
  1143. try:
  1144. self.right = next(self.rightseq)
  1145. except StopIteration:
  1146. self.is_rightseq_exhausted = True
  1147. self.keys = iter(self.d)
  1148. return next(self)
  1149. key = self.rightkey()
  1150. PySet_Add(self.seen_keys, key)
  1151. obj = PyDict_GetItem(self.d, key)
  1152. if obj is NULL:
  1153. return (self.left_default, self.right)
  1154. self.matches = <object>obj
  1155. self.i = 0
  1156. match = <object>PyList_GET_ITEM(self.matches, self.i) # skip error checking
  1157. self.i += 1
  1158. return (match, self.right)
  1159. else:
  1160. if self.i == PyList_GET_SIZE(self.matches):
  1161. key = next(self.keys)
  1162. while key in self.seen_keys:
  1163. key = next(self.keys)
  1164. obj = PyDict_GetItem(self.d, key)
  1165. self.matches = <object>obj
  1166. self.i = 0
  1167. match = <object>PyList_GET_ITEM(self.matches, self.i) # skip error checking
  1168. self.i += 1
  1169. return (match, self.right_default)
  1170. cdef class _outer_join_key(_outer_join):
  1171. cdef object rightkey(self):
  1172. return self._rightkey(self.right)
  1173. cdef class _outer_join_index(_outer_join):
  1174. cdef object rightkey(self):
  1175. return self.right[self._rightkey]
  1176. cdef class _outer_join_indices(_outer_join):
  1177. cdef object rightkey(self):
  1178. keyval = PyTuple_New(self.N)
  1179. for i in range(self.N):
  1180. val = <object>PyList_GET_ITEM(self._rightkey, i)
  1181. val = self.right[val]
  1182. Py_INCREF(val)
  1183. PyTuple_SET_ITEM(keyval, i, val)
  1184. return keyval
  1185. cdef class _left_outer_join(_join):
  1186. def __next__(self):
  1187. cdef PyObject *obj
  1188. if not self.is_rightseq_exhausted:
  1189. if self.i == PyList_GET_SIZE(self.matches):
  1190. obj = NULL
  1191. while obj is NULL:
  1192. try:
  1193. self.right = next(self.rightseq)
  1194. except StopIteration:
  1195. self.is_rightseq_exhausted = True
  1196. self.keys = iter(self.d)
  1197. return next(self)
  1198. key = self.rightkey()
  1199. PySet_Add(self.seen_keys, key)
  1200. obj = PyDict_GetItem(self.d, key)
  1201. self.matches = <object>obj
  1202. self.i = 0
  1203. match = <object>PyList_GET_ITEM(self.matches, self.i) # skip error checking
  1204. self.i += 1
  1205. return (match, self.right)
  1206. else:
  1207. if self.i == PyList_GET_SIZE(self.matches):
  1208. key = next(self.keys)
  1209. while key in self.seen_keys:
  1210. key = next(self.keys)
  1211. obj = PyDict_GetItem(self.d, key)
  1212. self.matches = <object>obj
  1213. self.i = 0
  1214. match = <object>PyList_GET_ITEM(self.matches, self.i) # skip error checking
  1215. self.i += 1
  1216. return (match, self.right_default)
  1217. cdef class _left_outer_join_key(_left_outer_join):
  1218. cdef object rightkey(self):
  1219. return self._rightkey(self.right)
  1220. cdef class _left_outer_join_index(_left_outer_join):
  1221. cdef object rightkey(self):
  1222. return self.right[self._rightkey]
  1223. cdef class _left_outer_join_indices(_left_outer_join):
  1224. cdef object rightkey(self):
  1225. keyval = PyTuple_New(self.N)
  1226. for i in range(self.N):
  1227. val = <object>PyList_GET_ITEM(self._rightkey, i)
  1228. val = self.right[val]
  1229. Py_INCREF(val)
  1230. PyTuple_SET_ITEM(keyval, i, val)
  1231. return keyval
  1232. cdef class _inner_join(_join):
  1233. def __next__(self):
  1234. cdef PyObject *obj = NULL
  1235. if self.i == PyList_GET_SIZE(self.matches):
  1236. while obj is NULL:
  1237. self.right = next(self.rightseq)
  1238. key = self.rightkey()
  1239. obj = PyDict_GetItem(self.d, key)
  1240. self.matches = <object>obj
  1241. self.i = 0
  1242. match = <object>PyList_GET_ITEM(self.matches, self.i) # skip error checking
  1243. self.i += 1
  1244. return (match, self.right)
  1245. cdef class _inner_join_key(_inner_join):
  1246. cdef object rightkey(self):
  1247. return self._rightkey(self.right)
  1248. cdef class _inner_join_index(_inner_join):
  1249. cdef object rightkey(self):
  1250. return self.right[self._rightkey]
  1251. cdef class _inner_join_indices(_inner_join):
  1252. cdef object rightkey(self):
  1253. keyval = PyTuple_New(self.N)
  1254. for i in range(self.N):
  1255. val = <object>PyList_GET_ITEM(self._rightkey, i)
  1256. val = self.right[val]
  1257. Py_INCREF(val)
  1258. PyTuple_SET_ITEM(keyval, i, val)
  1259. return keyval
  1260. cdef class _diff_key:
  1261. def __cinit__(self, object seqs, object key, object default=no_default):
  1262. self.N = len(seqs)
  1263. if self.N < 2:
  1264. raise TypeError('Too few sequences given (min 2 required)')
  1265. if default == no_default:
  1266. self.iters = zip(*seqs)
  1267. else:
  1268. self.iters = zip_longest(*seqs, fillvalue=default)
  1269. self.key = key
  1270. def __iter__(self):
  1271. return self
  1272. def __next__(self):
  1273. cdef object val, val2, items
  1274. cdef Py_ssize_t i
  1275. while True:
  1276. items = next(self.iters)
  1277. val = self.key(<object>PyTuple_GET_ITEM(items, 0))
  1278. for i in range(1, self.N):
  1279. val2 = self.key(<object>PyTuple_GET_ITEM(items, i))
  1280. if PyObject_RichCompareBool(val, val2, Py_NE):
  1281. return items
  1282. cdef class _diff_identity:
  1283. def __cinit__(self, object seqs, object default=no_default):
  1284. self.N = len(seqs)
  1285. if self.N < 2:
  1286. raise TypeError('Too few sequences given (min 2 required)')
  1287. if default == no_default:
  1288. self.iters = zip(*seqs)
  1289. else:
  1290. self.iters = zip_longest(*seqs, fillvalue=default)
  1291. def __iter__(self):
  1292. return self
  1293. def __next__(self):
  1294. cdef object val, val2, items
  1295. cdef Py_ssize_t i
  1296. while True:
  1297. items = next(self.iters)
  1298. val = <object>PyTuple_GET_ITEM(items, 0)
  1299. for i in range(1, self.N):
  1300. val2 = <object>PyTuple_GET_ITEM(items, i)
  1301. if PyObject_RichCompareBool(val, val2, Py_NE):
  1302. return items
  1303. cdef object c_diff(object seqs, object default=no_default, object key=None):
  1304. if key is None:
  1305. return _diff_identity(seqs, default=default)
  1306. else:
  1307. return _diff_key(seqs, key, default=default)
  1308. def diff(*seqs, **kwargs):
  1309. """
  1310. Return those items that differ between sequences
  1311. >>> list(diff([1, 2, 3], [1, 2, 10, 100]))
  1312. [(3, 10)]
  1313. Shorter sequences may be padded with a ``default`` value:
  1314. >>> list(diff([1, 2, 3], [1, 2, 10, 100], default=None))
  1315. [(3, 10), (None, 100)]
  1316. A ``key`` function may also be applied to each item to use during
  1317. comparisons:
  1318. >>> list(diff(['apples', 'bananas'], ['Apples', 'Oranges'], key=str.lower))
  1319. [('bananas', 'Oranges')]
  1320. """
  1321. N = len(seqs)
  1322. if N == 1 and isinstance(seqs[0], list):
  1323. seqs = seqs[0]
  1324. default = kwargs.get('default', no_default)
  1325. key = kwargs.get('key')
  1326. return c_diff(seqs, default=default, key=key)
  1327. cpdef object topk(Py_ssize_t k, object seq, object key=None):
  1328. """
  1329. Find the k largest elements of a sequence
  1330. Operates lazily in ``n*log(k)`` time
  1331. >>> topk(2, [1, 100, 10, 1000])
  1332. (1000, 100)
  1333. Use a key function to change sorted order
  1334. >>> topk(2, ['Alice', 'Bob', 'Charlie', 'Dan'], key=len)
  1335. ('Charlie', 'Alice')
  1336. See also:
  1337. heapq.nlargest
  1338. """
  1339. cdef object item, val, top
  1340. cdef object it = iter(seq)
  1341. cdef object _heapreplace = heapreplace
  1342. cdef Py_ssize_t i = k
  1343. cdef list pq = []
  1344. if key is not None and not callable(key):
  1345. key = getter(key)
  1346. if k < 2:
  1347. if k < 1:
  1348. return ()
  1349. top = list(take(1, it))
  1350. if len(top) == 0:
  1351. return ()
  1352. it = concatv(top, it)
  1353. if key is None:
  1354. return (max(it),)
  1355. else:
  1356. return (max(it, key=key),)
  1357. for item in it:
  1358. if key is None:
  1359. PyList_Append(pq, (item, i))
  1360. else:
  1361. PyList_Append(pq, (key(item), i, item))
  1362. i -= 1
  1363. if i == 0:
  1364. break
  1365. if i != 0:
  1366. pq.sort(reverse=True)
  1367. k = 0 if key is None else 2
  1368. return tuple([item[k] for item in pq])
  1369. heapify(pq)
  1370. top = pq[0][0]
  1371. if key is None:
  1372. for item in it:
  1373. if top < item:
  1374. _heapreplace(pq, (item, i))
  1375. top = pq[0][0]
  1376. i -= 1
  1377. else:
  1378. for item in it:
  1379. val = key(item)
  1380. if top < val:
  1381. _heapreplace(pq, (val, i, item))
  1382. top = pq[0][0]
  1383. i -= 1
  1384. pq.sort(reverse=True)
  1385. k = 0 if key is None else 2
  1386. return tuple([item[k] for item in pq])
  1387. cpdef object peek(object seq):
  1388. """
  1389. Retrieve the next element of a sequence
  1390. Returns the first element and an iterable equivalent to the original
  1391. sequence, still having the element retrieved.
  1392. >>> seq = [0, 1, 2, 3, 4]
  1393. >>> first, seq = peek(seq)
  1394. >>> first
  1395. 0
  1396. >>> list(seq)
  1397. [0, 1, 2, 3, 4]
  1398. """
  1399. iterator = iter(seq)
  1400. item = next(iterator)
  1401. return item, chain((item,), iterator)
  1402. cpdef object peekn(Py_ssize_t n, object seq):
  1403. """
  1404. Retrieve the next n elements of a sequence
  1405. Returns a tuple of the first n elements and an iterable equivalent
  1406. to the original, still having the elements retrieved.
  1407. >>> seq = [0, 1, 2, 3, 4]
  1408. >>> first_two, seq = peekn(2, seq)
  1409. >>> first_two
  1410. (0, 1)
  1411. >>> list(seq)
  1412. [0, 1, 2, 3, 4]
  1413. """
  1414. iterator = iter(seq)
  1415. peeked = tuple(take(n, iterator))
  1416. return peeked, chain(iter(peeked), iterator)
  1417. cdef class random_sample:
  1418. """ random_sample(prob, seq, random_state=None)
  1419. Return elements from a sequence with probability of prob
  1420. Returns a lazy iterator of random items from seq.
  1421. ``random_sample`` considers each item independently and without
  1422. replacement. See below how the first time it returned 13 items and the
  1423. next time it returned 6 items.
  1424. >>> seq = list(range(100))
  1425. >>> list(random_sample(0.1, seq)) # doctest: +SKIP
  1426. [6, 9, 19, 35, 45, 50, 58, 62, 68, 72, 78, 86, 95]
  1427. >>> list(random_sample(0.1, seq)) # doctest: +SKIP
  1428. [6, 44, 54, 61, 69, 94]
  1429. Providing an integer seed for ``random_state`` will result in
  1430. deterministic sampling. Given the same seed it will return the same sample
  1431. every time.
  1432. >>> list(random_sample(0.1, seq, random_state=2016))
  1433. [7, 9, 19, 25, 30, 32, 34, 48, 59, 60, 81, 98]
  1434. >>> list(random_sample(0.1, seq, random_state=2016))
  1435. [7, 9, 19, 25, 30, 32, 34, 48, 59, 60, 81, 98]
  1436. ``random_state`` can also be any object with a method ``random`` that
  1437. returns floats between 0.0 and 1.0 (exclusive).
  1438. >>> from random import Random
  1439. >>> randobj = Random(2016)
  1440. >>> list(random_sample(0.1, seq, random_state=randobj))
  1441. [7, 9, 19, 25, 30, 32, 34, 48, 59, 60, 81, 98]
  1442. """
  1443. def __cinit__(self, object prob, object seq, random_state=None):
  1444. float(prob)
  1445. self.prob = prob
  1446. self.iter_seq = iter(seq)
  1447. if not hasattr(random_state, 'random'):
  1448. from random import Random
  1449. random_state = Random(random_state)
  1450. self.random_func = random_state.random
  1451. def __iter__(self):
  1452. return self
  1453. def __next__(self):
  1454. while True:
  1455. if self.random_func() < self.prob:
  1456. return next(self.iter_seq)
  1457. next(self.iter_seq)