itertoolsmodule.c 132 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003100410051006100710081009101010111012101310141015101610171018101910201021102210231024102510261027102810291030103110321033103410351036103710381039104010411042104310441045104610471048104910501051105210531054105510561057105810591060106110621063106410651066106710681069107010711072107310741075107610771078107910801081108210831084108510861087108810891090109110921093109410951096109710981099110011011102110311041105110611071108110911101111111211131114111511161117111811191120112111221123112411251126112711281129113011311132113311341135113611371138113911401141114211431144114511461147114811491150115111521153115411551156115711581159116011611162116311641165116611671168116911701171117211731174117511761177117811791180118111821183118411851186118711881189119011911192119311941195119611971198119912001201120212031204120512061207120812091210121112121213121412151216121712181219122012211222122312241225122612271228122912301231123212331234123512361237123812391240124112421243124412451246124712481249125012511252125312541255125612571258125912601261126212631264126512661267126812691270127112721273127412751276127712781279128012811282128312841285128612871288128912901291129212931294129512961297129812991300130113021303130413051306130713081309131013111312131313141315131613171318131913201321132213231324132513261327132813291330133113321333133413351336133713381339134013411342134313441345134613471348134913501351135213531354135513561357135813591360136113621363136413651366136713681369137013711372137313741375137613771378137913801381138213831384138513861387138813891390139113921393139413951396139713981399140014011402140314041405140614071408140914101411141214131414141514161417141814191420142114221423142414251426142714281429143014311432143314341435143614371438143914401441144214431444144514461447144814491450145114521453145414551456145714581459146014611462146314641465146614671468146914701471147214731474147514761477147814791480148114821483148414851486148714881489149014911492149314941495149614971498149915001501150215031504150515061507150815091510151115121513151415151516151715181519152015211522152315241525152615271528152915301531153215331534153515361537153815391540154115421543154415451546154715481549155015511552155315541555155615571558155915601561156215631564156515661567156815691570157115721573157415751576157715781579158015811582158315841585158615871588158915901591159215931594159515961597159815991600160116021603160416051606160716081609161016111612161316141615161616171618161916201621162216231624162516261627162816291630163116321633163416351636163716381639164016411642164316441645164616471648164916501651165216531654165516561657165816591660166116621663166416651666166716681669167016711672167316741675167616771678167916801681168216831684168516861687168816891690169116921693169416951696169716981699170017011702170317041705170617071708170917101711171217131714171517161717171817191720172117221723172417251726172717281729173017311732173317341735173617371738173917401741174217431744174517461747174817491750175117521753175417551756175717581759176017611762176317641765176617671768176917701771177217731774177517761777177817791780178117821783178417851786178717881789179017911792179317941795179617971798179918001801180218031804180518061807180818091810181118121813181418151816181718181819182018211822182318241825182618271828182918301831183218331834183518361837183818391840184118421843184418451846184718481849185018511852185318541855185618571858185918601861186218631864186518661867186818691870187118721873187418751876187718781879188018811882188318841885188618871888188918901891189218931894189518961897189818991900190119021903190419051906190719081909191019111912191319141915191619171918191919201921192219231924192519261927192819291930193119321933193419351936193719381939194019411942194319441945194619471948194919501951195219531954195519561957195819591960196119621963196419651966196719681969197019711972197319741975197619771978197919801981198219831984198519861987198819891990199119921993199419951996199719981999200020012002200320042005200620072008200920102011201220132014201520162017201820192020202120222023202420252026202720282029203020312032203320342035203620372038203920402041204220432044204520462047204820492050205120522053205420552056205720582059206020612062206320642065206620672068206920702071207220732074207520762077207820792080208120822083208420852086208720882089209020912092209320942095209620972098209921002101210221032104210521062107210821092110211121122113211421152116211721182119212021212122212321242125212621272128212921302131213221332134213521362137213821392140214121422143214421452146214721482149215021512152215321542155215621572158215921602161216221632164216521662167216821692170217121722173217421752176217721782179218021812182218321842185218621872188218921902191219221932194219521962197219821992200220122022203220422052206220722082209221022112212221322142215221622172218221922202221222222232224222522262227222822292230223122322233223422352236223722382239224022412242224322442245224622472248224922502251225222532254225522562257225822592260226122622263226422652266226722682269227022712272227322742275227622772278227922802281228222832284228522862287228822892290229122922293229422952296229722982299230023012302230323042305230623072308230923102311231223132314231523162317231823192320232123222323232423252326232723282329233023312332233323342335233623372338233923402341234223432344234523462347234823492350235123522353235423552356235723582359236023612362236323642365236623672368236923702371237223732374237523762377237823792380238123822383238423852386238723882389239023912392239323942395239623972398239924002401240224032404240524062407240824092410241124122413241424152416241724182419242024212422242324242425242624272428242924302431243224332434243524362437243824392440244124422443244424452446244724482449245024512452245324542455245624572458245924602461246224632464246524662467246824692470247124722473247424752476247724782479248024812482248324842485248624872488248924902491249224932494249524962497249824992500250125022503250425052506250725082509251025112512251325142515251625172518251925202521252225232524252525262527252825292530253125322533253425352536253725382539254025412542254325442545254625472548254925502551255225532554255525562557255825592560256125622563256425652566256725682569257025712572257325742575257625772578257925802581258225832584258525862587258825892590259125922593259425952596259725982599260026012602260326042605260626072608260926102611261226132614261526162617261826192620262126222623262426252626262726282629263026312632263326342635263626372638263926402641264226432644264526462647264826492650265126522653265426552656265726582659266026612662266326642665266626672668266926702671267226732674267526762677267826792680268126822683268426852686268726882689269026912692269326942695269626972698269927002701270227032704270527062707270827092710271127122713271427152716271727182719272027212722272327242725272627272728272927302731273227332734273527362737273827392740274127422743274427452746274727482749275027512752275327542755275627572758275927602761276227632764276527662767276827692770277127722773277427752776277727782779278027812782278327842785278627872788278927902791279227932794279527962797279827992800280128022803280428052806280728082809281028112812281328142815281628172818281928202821282228232824282528262827282828292830283128322833283428352836283728382839284028412842284328442845284628472848284928502851285228532854285528562857285828592860286128622863286428652866286728682869287028712872287328742875287628772878287928802881288228832884288528862887288828892890289128922893289428952896289728982899290029012902290329042905290629072908290929102911291229132914291529162917291829192920292129222923292429252926292729282929293029312932293329342935293629372938293929402941294229432944294529462947294829492950295129522953295429552956295729582959296029612962296329642965296629672968296929702971297229732974297529762977297829792980298129822983298429852986298729882989299029912992299329942995299629972998299930003001300230033004300530063007300830093010301130123013301430153016301730183019302030213022302330243025302630273028302930303031303230333034303530363037303830393040304130423043304430453046304730483049305030513052305330543055305630573058305930603061306230633064306530663067306830693070307130723073307430753076307730783079308030813082308330843085308630873088308930903091309230933094309530963097309830993100310131023103310431053106310731083109311031113112311331143115311631173118311931203121312231233124312531263127312831293130313131323133313431353136313731383139314031413142314331443145314631473148314931503151315231533154315531563157315831593160316131623163316431653166316731683169317031713172317331743175317631773178317931803181318231833184318531863187318831893190319131923193319431953196319731983199320032013202320332043205320632073208320932103211321232133214321532163217321832193220322132223223322432253226322732283229323032313232323332343235323632373238323932403241324232433244324532463247324832493250325132523253325432553256325732583259326032613262326332643265326632673268326932703271327232733274327532763277327832793280328132823283328432853286328732883289329032913292329332943295329632973298329933003301330233033304330533063307330833093310331133123313331433153316331733183319332033213322332333243325332633273328332933303331333233333334333533363337333833393340334133423343334433453346334733483349335033513352335333543355335633573358335933603361336233633364336533663367336833693370337133723373337433753376337733783379338033813382338333843385338633873388338933903391339233933394339533963397339833993400340134023403340434053406340734083409341034113412341334143415341634173418341934203421342234233424342534263427342834293430343134323433343434353436343734383439344034413442344334443445344634473448344934503451345234533454345534563457345834593460346134623463346434653466346734683469347034713472347334743475347634773478347934803481348234833484348534863487348834893490349134923493349434953496349734983499350035013502350335043505350635073508350935103511351235133514351535163517351835193520352135223523352435253526352735283529353035313532353335343535353635373538353935403541354235433544354535463547354835493550355135523553355435553556355735583559356035613562356335643565356635673568356935703571357235733574357535763577357835793580358135823583358435853586358735883589359035913592359335943595359635973598359936003601360236033604360536063607360836093610361136123613361436153616361736183619362036213622362336243625362636273628362936303631363236333634363536363637363836393640364136423643364436453646364736483649365036513652365336543655365636573658365936603661366236633664366536663667366836693670367136723673367436753676367736783679368036813682368336843685368636873688368936903691369236933694369536963697369836993700370137023703370437053706370737083709371037113712371337143715371637173718371937203721372237233724372537263727372837293730373137323733373437353736373737383739374037413742374337443745374637473748374937503751375237533754375537563757375837593760376137623763376437653766376737683769377037713772377337743775377637773778377937803781378237833784378537863787378837893790379137923793379437953796379737983799380038013802380338043805380638073808380938103811381238133814381538163817381838193820382138223823382438253826382738283829383038313832383338343835383638373838383938403841384238433844384538463847384838493850385138523853385438553856385738583859386038613862386338643865386638673868386938703871387238733874387538763877387838793880388138823883388438853886388738883889389038913892389338943895389638973898389939003901390239033904390539063907390839093910391139123913391439153916391739183919392039213922392339243925392639273928392939303931393239333934393539363937393839393940394139423943394439453946394739483949395039513952395339543955395639573958395939603961396239633964396539663967396839693970397139723973397439753976397739783979398039813982398339843985398639873988398939903991399239933994399539963997399839994000400140024003400440054006400740084009401040114012401340144015401640174018401940204021402240234024402540264027402840294030403140324033403440354036403740384039404040414042404340444045404640474048404940504051405240534054405540564057405840594060406140624063406440654066406740684069407040714072407340744075407640774078407940804081408240834084408540864087408840894090409140924093409440954096409740984099410041014102410341044105410641074108410941104111411241134114411541164117411841194120412141224123412441254126412741284129413041314132413341344135413641374138413941404141414241434144414541464147414841494150415141524153415441554156415741584159416041614162416341644165416641674168416941704171417241734174417541764177417841794180418141824183418441854186418741884189419041914192419341944195419641974198419942004201420242034204420542064207420842094210421142124213421442154216421742184219422042214222422342244225422642274228422942304231423242334234423542364237423842394240424142424243424442454246424742484249425042514252425342544255425642574258425942604261426242634264426542664267426842694270427142724273427442754276427742784279428042814282428342844285428642874288428942904291429242934294429542964297429842994300430143024303430443054306430743084309431043114312431343144315431643174318431943204321432243234324432543264327432843294330433143324333433443354336433743384339434043414342434343444345434643474348434943504351435243534354435543564357435843594360436143624363436443654366436743684369437043714372437343744375437643774378437943804381438243834384438543864387438843894390439143924393439443954396439743984399440044014402440344044405440644074408440944104411441244134414441544164417441844194420442144224423442444254426442744284429443044314432443344344435443644374438443944404441444244434444444544464447444844494450445144524453445444554456445744584459446044614462446344644465446644674468446944704471447244734474447544764477447844794480448144824483448444854486448744884489449044914492449344944495449644974498449945004501450245034504450545064507450845094510451145124513451445154516451745184519452045214522452345244525452645274528452945304531453245334534453545364537453845394540454145424543454445454546454745484549455045514552455345544555455645574558455945604561456245634564456545664567456845694570457145724573457445754576457745784579458045814582458345844585458645874588458945904591459245934594459545964597459845994600460146024603460446054606460746084609461046114612461346144615461646174618461946204621462246234624462546264627462846294630463146324633463446354636463746384639464046414642464346444645464646474648464946504651465246534654465546564657465846594660466146624663466446654666466746684669467046714672467346744675467646774678467946804681468246834684468546864687468846894690469146924693469446954696469746984699470047014702470347044705470647074708470947104711471247134714471547164717471847194720472147224723472447254726472747284729473047314732473347344735473647374738473947404741474247434744474547464747474847494750475147524753475447554756475747584759476047614762476347644765476647674768476947704771
  1. #define PY_SSIZE_T_CLEAN
  2. #include "Python.h"
  3. #include "pycore_call.h" // _PyObject_CallNoArgs()
  4. #include "pycore_long.h" // _PyLong_GetZero()
  5. #include "pycore_moduleobject.h" // _PyModule_GetState()
  6. #include "pycore_typeobject.h" // _PyType_GetModuleState()
  7. #include "pycore_object.h" // _PyObject_GC_TRACK()
  8. #include "pycore_tuple.h" // _PyTuple_ITEMS()
  9. #include "structmember.h" // PyMemberDef
  10. #include <stddef.h> // offsetof()
  11. /* Itertools module written and maintained
  12. by Raymond D. Hettinger <python@rcn.com>
  13. */
  14. typedef struct {
  15. PyTypeObject *accumulate_type;
  16. PyTypeObject *batched_type;
  17. PyTypeObject *chain_type;
  18. PyTypeObject *combinations_type;
  19. PyTypeObject *compress_type;
  20. PyTypeObject *count_type;
  21. PyTypeObject *cwr_type;
  22. PyTypeObject *cycle_type;
  23. PyTypeObject *dropwhile_type;
  24. PyTypeObject *filterfalse_type;
  25. PyTypeObject *groupby_type;
  26. PyTypeObject *_grouper_type;
  27. PyTypeObject *islice_type;
  28. PyTypeObject *pairwise_type;
  29. PyTypeObject *permutations_type;
  30. PyTypeObject *product_type;
  31. PyTypeObject *repeat_type;
  32. PyTypeObject *starmap_type;
  33. PyTypeObject *takewhile_type;
  34. PyTypeObject *tee_type;
  35. PyTypeObject *teedataobject_type;
  36. PyTypeObject *ziplongest_type;
  37. } itertools_state;
  38. static inline itertools_state *
  39. get_module_state(PyObject *mod)
  40. {
  41. void *state = _PyModule_GetState(mod);
  42. assert(state != NULL);
  43. return (itertools_state *)state;
  44. }
  45. static inline itertools_state *
  46. get_module_state_by_cls(PyTypeObject *cls)
  47. {
  48. void *state = _PyType_GetModuleState(cls);
  49. assert(state != NULL);
  50. return (itertools_state *)state;
  51. }
  52. static struct PyModuleDef itertoolsmodule;
  53. static inline itertools_state *
  54. find_state_by_type(PyTypeObject *tp)
  55. {
  56. PyObject *mod = PyType_GetModuleByDef(tp, &itertoolsmodule);
  57. assert(mod != NULL);
  58. return get_module_state(mod);
  59. }
  60. /*[clinic input]
  61. module itertools
  62. class itertools.groupby "groupbyobject *" "clinic_state()->groupby_type"
  63. class itertools._grouper "_grouperobject *" "clinic_state()->_grouper_type"
  64. class itertools.teedataobject "teedataobject *" "clinic_state()->teedataobject_type"
  65. class itertools._tee "teeobject *" "clinic_state()->tee_type"
  66. class itertools.batched "batchedobject *" "clinic_state()->batched_type"
  67. class itertools.cycle "cycleobject *" "clinic_state()->cycle_type"
  68. class itertools.dropwhile "dropwhileobject *" "clinic_state()->dropwhile_type"
  69. class itertools.takewhile "takewhileobject *" "clinic_state()->takewhile_type"
  70. class itertools.starmap "starmapobject *" "clinic_state()->starmap_type"
  71. class itertools.chain "chainobject *" "clinic_state()->chain_type"
  72. class itertools.combinations "combinationsobject *" "clinic_state()->combinations_type"
  73. class itertools.combinations_with_replacement "cwr_object *" "clinic_state()->cwr_type"
  74. class itertools.permutations "permutationsobject *" "clinic_state()->permutations_type"
  75. class itertools.accumulate "accumulateobject *" "clinic_state()->accumulate_type"
  76. class itertools.compress "compressobject *" "clinic_state()->compress_type"
  77. class itertools.filterfalse "filterfalseobject *" "clinic_state()->filterfalse_type"
  78. class itertools.count "countobject *" "clinic_state()->count_type"
  79. class itertools.pairwise "pairwiseobject *" "clinic_state()->pairwise_type"
  80. [clinic start generated code]*/
  81. /*[clinic end generated code: output=da39a3ee5e6b4b0d input=aa48fe4de9d4080f]*/
  82. #define clinic_state() (find_state_by_type(type))
  83. #define clinic_state_by_cls() (get_module_state_by_cls(base_tp))
  84. #include "clinic/itertoolsmodule.c.h"
  85. #undef clinic_state_by_cls
  86. #undef clinic_state
  87. /* Deprecation of pickle support: GH-101588 *********************************/
  88. #define ITERTOOL_PICKLE_DEPRECATION \
  89. if (PyErr_WarnEx( \
  90. PyExc_DeprecationWarning, \
  91. "Pickle, copy, and deepcopy support will be " \
  92. "removed from itertools in Python 3.14.", 1) < 0) { \
  93. return NULL; \
  94. }
  95. /* batched object ************************************************************/
  96. /* Note: The built-in zip() function includes a "strict" argument
  97. that was needed because that function would silently truncate data,
  98. and there was no easy way for a user to detect the data loss.
  99. The same reasoning does not apply to batched() which never drops data.
  100. Instead, batched() produces a shorter tuple which can be handled
  101. as the user sees fit. If requested, it would be reasonable to add
  102. "fillvalue" support which had demonstrated value in zip_longest().
  103. For now, the API is kept simple and clean.
  104. */
  105. typedef struct {
  106. PyObject_HEAD
  107. PyObject *it;
  108. Py_ssize_t batch_size;
  109. } batchedobject;
  110. /*[clinic input]
  111. @classmethod
  112. itertools.batched.__new__ as batched_new
  113. iterable: object
  114. n: Py_ssize_t
  115. Batch data into tuples of length n. The last batch may be shorter than n.
  116. Loops over the input iterable and accumulates data into tuples
  117. up to size n. The input is consumed lazily, just enough to
  118. fill a batch. The result is yielded as soon as a batch is full
  119. or when the input iterable is exhausted.
  120. >>> for batch in batched('ABCDEFG', 3):
  121. ... print(batch)
  122. ...
  123. ('A', 'B', 'C')
  124. ('D', 'E', 'F')
  125. ('G',)
  126. [clinic start generated code]*/
  127. static PyObject *
  128. batched_new_impl(PyTypeObject *type, PyObject *iterable, Py_ssize_t n)
  129. /*[clinic end generated code: output=7ebc954d655371b6 input=ffd70726927c5129]*/
  130. {
  131. PyObject *it;
  132. batchedobject *bo;
  133. if (n < 1) {
  134. /* We could define the n==0 case to return an empty iterator
  135. but that is at odds with the idea that batching should
  136. never throw-away input data.
  137. */
  138. PyErr_SetString(PyExc_ValueError, "n must be at least one");
  139. return NULL;
  140. }
  141. it = PyObject_GetIter(iterable);
  142. if (it == NULL) {
  143. return NULL;
  144. }
  145. /* create batchedobject structure */
  146. bo = (batchedobject *)type->tp_alloc(type, 0);
  147. if (bo == NULL) {
  148. Py_DECREF(it);
  149. return NULL;
  150. }
  151. bo->batch_size = n;
  152. bo->it = it;
  153. return (PyObject *)bo;
  154. }
  155. static void
  156. batched_dealloc(batchedobject *bo)
  157. {
  158. PyTypeObject *tp = Py_TYPE(bo);
  159. PyObject_GC_UnTrack(bo);
  160. Py_XDECREF(bo->it);
  161. tp->tp_free(bo);
  162. Py_DECREF(tp);
  163. }
  164. static int
  165. batched_traverse(batchedobject *bo, visitproc visit, void *arg)
  166. {
  167. Py_VISIT(Py_TYPE(bo));
  168. Py_VISIT(bo->it);
  169. return 0;
  170. }
  171. static PyObject *
  172. batched_next(batchedobject *bo)
  173. {
  174. Py_ssize_t i;
  175. Py_ssize_t n = bo->batch_size;
  176. PyObject *it = bo->it;
  177. PyObject *item;
  178. PyObject *result;
  179. if (it == NULL) {
  180. return NULL;
  181. }
  182. result = PyTuple_New(n);
  183. if (result == NULL) {
  184. return NULL;
  185. }
  186. iternextfunc iternext = *Py_TYPE(it)->tp_iternext;
  187. PyObject **items = _PyTuple_ITEMS(result);
  188. for (i=0 ; i < n ; i++) {
  189. item = iternext(it);
  190. if (item == NULL) {
  191. goto null_item;
  192. }
  193. items[i] = item;
  194. }
  195. return result;
  196. null_item:
  197. if (PyErr_Occurred()) {
  198. if (!PyErr_ExceptionMatches(PyExc_StopIteration)) {
  199. /* Input raised an exception other than StopIteration */
  200. Py_CLEAR(bo->it);
  201. Py_DECREF(result);
  202. return NULL;
  203. }
  204. PyErr_Clear();
  205. }
  206. if (i == 0) {
  207. Py_CLEAR(bo->it);
  208. Py_DECREF(result);
  209. return NULL;
  210. }
  211. _PyTuple_Resize(&result, i);
  212. return result;
  213. }
  214. static PyType_Slot batched_slots[] = {
  215. {Py_tp_dealloc, batched_dealloc},
  216. {Py_tp_getattro, PyObject_GenericGetAttr},
  217. {Py_tp_doc, (void *)batched_new__doc__},
  218. {Py_tp_traverse, batched_traverse},
  219. {Py_tp_iter, PyObject_SelfIter},
  220. {Py_tp_iternext, batched_next},
  221. {Py_tp_alloc, PyType_GenericAlloc},
  222. {Py_tp_new, batched_new},
  223. {Py_tp_free, PyObject_GC_Del},
  224. {0, NULL},
  225. };
  226. static PyType_Spec batched_spec = {
  227. .name = "itertools.batched",
  228. .basicsize = sizeof(batchedobject),
  229. .flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_BASETYPE |
  230. Py_TPFLAGS_IMMUTABLETYPE),
  231. .slots = batched_slots,
  232. };
  233. /* pairwise object ***********************************************************/
  234. typedef struct {
  235. PyObject_HEAD
  236. PyObject *it;
  237. PyObject *old;
  238. } pairwiseobject;
  239. /*[clinic input]
  240. @classmethod
  241. itertools.pairwise.__new__ as pairwise_new
  242. iterable: object
  243. /
  244. Return an iterator of overlapping pairs taken from the input iterator.
  245. s -> (s0,s1), (s1,s2), (s2, s3), ...
  246. [clinic start generated code]*/
  247. static PyObject *
  248. pairwise_new_impl(PyTypeObject *type, PyObject *iterable)
  249. /*[clinic end generated code: output=9f0267062d384456 input=6e7c3cddb431a8d6]*/
  250. {
  251. PyObject *it;
  252. pairwiseobject *po;
  253. it = PyObject_GetIter(iterable);
  254. if (it == NULL) {
  255. return NULL;
  256. }
  257. po = (pairwiseobject *)type->tp_alloc(type, 0);
  258. if (po == NULL) {
  259. Py_DECREF(it);
  260. return NULL;
  261. }
  262. po->it = it;
  263. po->old = NULL;
  264. return (PyObject *)po;
  265. }
  266. static void
  267. pairwise_dealloc(pairwiseobject *po)
  268. {
  269. PyTypeObject *tp = Py_TYPE(po);
  270. PyObject_GC_UnTrack(po);
  271. Py_XDECREF(po->it);
  272. Py_XDECREF(po->old);
  273. tp->tp_free(po);
  274. Py_DECREF(tp);
  275. }
  276. static int
  277. pairwise_traverse(pairwiseobject *po, visitproc visit, void *arg)
  278. {
  279. Py_VISIT(Py_TYPE(po));
  280. Py_VISIT(po->it);
  281. Py_VISIT(po->old);
  282. return 0;
  283. }
  284. static PyObject *
  285. pairwise_next(pairwiseobject *po)
  286. {
  287. PyObject *it = po->it;
  288. PyObject *old = po->old;
  289. PyObject *new, *result;
  290. if (it == NULL) {
  291. return NULL;
  292. }
  293. if (old == NULL) {
  294. old = (*Py_TYPE(it)->tp_iternext)(it);
  295. Py_XSETREF(po->old, old);
  296. if (old == NULL) {
  297. Py_CLEAR(po->it);
  298. return NULL;
  299. }
  300. it = po->it;
  301. if (it == NULL) {
  302. Py_CLEAR(po->old);
  303. return NULL;
  304. }
  305. }
  306. Py_INCREF(old);
  307. new = (*Py_TYPE(it)->tp_iternext)(it);
  308. if (new == NULL) {
  309. Py_CLEAR(po->it);
  310. Py_CLEAR(po->old);
  311. Py_DECREF(old);
  312. return NULL;
  313. }
  314. /* Future optimization: Reuse the result tuple as we do in enumerate() */
  315. result = PyTuple_Pack(2, old, new);
  316. Py_XSETREF(po->old, new);
  317. Py_DECREF(old);
  318. return result;
  319. }
  320. static PyType_Slot pairwise_slots[] = {
  321. {Py_tp_dealloc, pairwise_dealloc},
  322. {Py_tp_getattro, PyObject_GenericGetAttr},
  323. {Py_tp_doc, (void *)pairwise_new__doc__},
  324. {Py_tp_traverse, pairwise_traverse},
  325. {Py_tp_iter, PyObject_SelfIter},
  326. {Py_tp_iternext, pairwise_next},
  327. {Py_tp_alloc, PyType_GenericAlloc},
  328. {Py_tp_new, pairwise_new},
  329. {Py_tp_free, PyObject_GC_Del},
  330. {0, NULL},
  331. };
  332. static PyType_Spec pairwise_spec = {
  333. .name = "itertools.pairwise",
  334. .basicsize = sizeof(pairwiseobject),
  335. .flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_BASETYPE |
  336. Py_TPFLAGS_IMMUTABLETYPE),
  337. .slots = pairwise_slots,
  338. };
  339. /* groupby object ************************************************************/
  340. typedef struct {
  341. PyObject_HEAD
  342. PyObject *it;
  343. PyObject *keyfunc;
  344. PyObject *tgtkey;
  345. PyObject *currkey;
  346. PyObject *currvalue;
  347. const void *currgrouper; /* borrowed reference */
  348. itertools_state *state;
  349. } groupbyobject;
  350. static PyObject *_grouper_create(groupbyobject *, PyObject *);
  351. /*[clinic input]
  352. @classmethod
  353. itertools.groupby.__new__
  354. iterable as it: object
  355. Elements to divide into groups according to the key function.
  356. key as keyfunc: object = None
  357. A function for computing the group category for each element.
  358. If the key function is not specified or is None, the element itself
  359. is used for grouping.
  360. make an iterator that returns consecutive keys and groups from the iterable
  361. [clinic start generated code]*/
  362. static PyObject *
  363. itertools_groupby_impl(PyTypeObject *type, PyObject *it, PyObject *keyfunc)
  364. /*[clinic end generated code: output=cbb1ae3a90fd4141 input=6b3d123e87ff65a1]*/
  365. {
  366. groupbyobject *gbo;
  367. gbo = (groupbyobject *)type->tp_alloc(type, 0);
  368. if (gbo == NULL)
  369. return NULL;
  370. gbo->tgtkey = NULL;
  371. gbo->currkey = NULL;
  372. gbo->currvalue = NULL;
  373. gbo->keyfunc = Py_NewRef(keyfunc);
  374. gbo->it = PyObject_GetIter(it);
  375. if (gbo->it == NULL) {
  376. Py_DECREF(gbo);
  377. return NULL;
  378. }
  379. gbo->state = find_state_by_type(type);
  380. return (PyObject *)gbo;
  381. }
  382. static void
  383. groupby_dealloc(groupbyobject *gbo)
  384. {
  385. PyTypeObject *tp = Py_TYPE(gbo);
  386. PyObject_GC_UnTrack(gbo);
  387. Py_XDECREF(gbo->it);
  388. Py_XDECREF(gbo->keyfunc);
  389. Py_XDECREF(gbo->tgtkey);
  390. Py_XDECREF(gbo->currkey);
  391. Py_XDECREF(gbo->currvalue);
  392. tp->tp_free(gbo);
  393. Py_DECREF(tp);
  394. }
  395. static int
  396. groupby_traverse(groupbyobject *gbo, visitproc visit, void *arg)
  397. {
  398. Py_VISIT(Py_TYPE(gbo));
  399. Py_VISIT(gbo->it);
  400. Py_VISIT(gbo->keyfunc);
  401. Py_VISIT(gbo->tgtkey);
  402. Py_VISIT(gbo->currkey);
  403. Py_VISIT(gbo->currvalue);
  404. return 0;
  405. }
  406. Py_LOCAL_INLINE(int)
  407. groupby_step(groupbyobject *gbo)
  408. {
  409. PyObject *newvalue, *newkey, *oldvalue;
  410. newvalue = PyIter_Next(gbo->it);
  411. if (newvalue == NULL)
  412. return -1;
  413. if (gbo->keyfunc == Py_None) {
  414. newkey = Py_NewRef(newvalue);
  415. } else {
  416. newkey = PyObject_CallOneArg(gbo->keyfunc, newvalue);
  417. if (newkey == NULL) {
  418. Py_DECREF(newvalue);
  419. return -1;
  420. }
  421. }
  422. oldvalue = gbo->currvalue;
  423. gbo->currvalue = newvalue;
  424. Py_XSETREF(gbo->currkey, newkey);
  425. Py_XDECREF(oldvalue);
  426. return 0;
  427. }
  428. static PyObject *
  429. groupby_next(groupbyobject *gbo)
  430. {
  431. PyObject *r, *grouper;
  432. gbo->currgrouper = NULL;
  433. /* skip to next iteration group */
  434. for (;;) {
  435. if (gbo->currkey == NULL)
  436. /* pass */;
  437. else if (gbo->tgtkey == NULL)
  438. break;
  439. else {
  440. int rcmp;
  441. rcmp = PyObject_RichCompareBool(gbo->tgtkey, gbo->currkey, Py_EQ);
  442. if (rcmp == -1)
  443. return NULL;
  444. else if (rcmp == 0)
  445. break;
  446. }
  447. if (groupby_step(gbo) < 0)
  448. return NULL;
  449. }
  450. Py_INCREF(gbo->currkey);
  451. Py_XSETREF(gbo->tgtkey, gbo->currkey);
  452. grouper = _grouper_create(gbo, gbo->tgtkey);
  453. if (grouper == NULL)
  454. return NULL;
  455. r = PyTuple_Pack(2, gbo->currkey, grouper);
  456. Py_DECREF(grouper);
  457. return r;
  458. }
  459. static PyObject *
  460. groupby_reduce(groupbyobject *lz, PyObject *Py_UNUSED(ignored))
  461. {
  462. /* reduce as a 'new' call with an optional 'setstate' if groupby
  463. * has started
  464. */
  465. ITERTOOL_PICKLE_DEPRECATION;
  466. PyObject *value;
  467. if (lz->tgtkey && lz->currkey && lz->currvalue)
  468. value = Py_BuildValue("O(OO)(OOO)", Py_TYPE(lz),
  469. lz->it, lz->keyfunc, lz->currkey, lz->currvalue, lz->tgtkey);
  470. else
  471. value = Py_BuildValue("O(OO)", Py_TYPE(lz),
  472. lz->it, lz->keyfunc);
  473. return value;
  474. }
  475. PyDoc_STRVAR(reduce_doc, "Return state information for pickling.");
  476. static PyObject *
  477. groupby_setstate(groupbyobject *lz, PyObject *state)
  478. {
  479. ITERTOOL_PICKLE_DEPRECATION;
  480. PyObject *currkey, *currvalue, *tgtkey;
  481. if (!PyTuple_Check(state)) {
  482. PyErr_SetString(PyExc_TypeError, "state is not a tuple");
  483. return NULL;
  484. }
  485. if (!PyArg_ParseTuple(state, "OOO", &currkey, &currvalue, &tgtkey)) {
  486. return NULL;
  487. }
  488. Py_INCREF(currkey);
  489. Py_XSETREF(lz->currkey, currkey);
  490. Py_INCREF(currvalue);
  491. Py_XSETREF(lz->currvalue, currvalue);
  492. Py_INCREF(tgtkey);
  493. Py_XSETREF(lz->tgtkey, tgtkey);
  494. Py_RETURN_NONE;
  495. }
  496. PyDoc_STRVAR(setstate_doc, "Set state information for unpickling.");
  497. static PyMethodDef groupby_methods[] = {
  498. {"__reduce__", (PyCFunction)groupby_reduce, METH_NOARGS,
  499. reduce_doc},
  500. {"__setstate__", (PyCFunction)groupby_setstate, METH_O,
  501. setstate_doc},
  502. {NULL, NULL} /* sentinel */
  503. };
  504. static PyType_Slot groupby_slots[] = {
  505. {Py_tp_dealloc, groupby_dealloc},
  506. {Py_tp_getattro, PyObject_GenericGetAttr},
  507. {Py_tp_doc, (void *)itertools_groupby__doc__},
  508. {Py_tp_traverse, groupby_traverse},
  509. {Py_tp_iter, PyObject_SelfIter},
  510. {Py_tp_iternext, groupby_next},
  511. {Py_tp_methods, groupby_methods},
  512. {Py_tp_new, itertools_groupby},
  513. {Py_tp_free, PyObject_GC_Del},
  514. {0, NULL},
  515. };
  516. static PyType_Spec groupby_spec = {
  517. .name = "itertools.groupby",
  518. .basicsize= sizeof(groupbyobject),
  519. .flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_BASETYPE |
  520. Py_TPFLAGS_IMMUTABLETYPE),
  521. .slots = groupby_slots,
  522. };
  523. /* _grouper object (internal) ************************************************/
  524. typedef struct {
  525. PyObject_HEAD
  526. PyObject *parent;
  527. PyObject *tgtkey;
  528. } _grouperobject;
  529. /*[clinic input]
  530. @classmethod
  531. itertools._grouper.__new__
  532. parent: object(subclass_of='clinic_state_by_cls()->groupby_type')
  533. tgtkey: object
  534. /
  535. [clinic start generated code]*/
  536. static PyObject *
  537. itertools__grouper_impl(PyTypeObject *type, PyObject *parent,
  538. PyObject *tgtkey)
  539. /*[clinic end generated code: output=462efb1cdebb5914 input=afe05eb477118f12]*/
  540. {
  541. return _grouper_create((groupbyobject*) parent, tgtkey);
  542. }
  543. static PyObject *
  544. _grouper_create(groupbyobject *parent, PyObject *tgtkey)
  545. {
  546. itertools_state *state = parent->state;
  547. _grouperobject *igo = PyObject_GC_New(_grouperobject, state->_grouper_type);
  548. if (igo == NULL)
  549. return NULL;
  550. igo->parent = Py_NewRef(parent);
  551. igo->tgtkey = Py_NewRef(tgtkey);
  552. parent->currgrouper = igo; /* borrowed reference */
  553. PyObject_GC_Track(igo);
  554. return (PyObject *)igo;
  555. }
  556. static void
  557. _grouper_dealloc(_grouperobject *igo)
  558. {
  559. PyTypeObject *tp = Py_TYPE(igo);
  560. PyObject_GC_UnTrack(igo);
  561. Py_DECREF(igo->parent);
  562. Py_DECREF(igo->tgtkey);
  563. PyObject_GC_Del(igo);
  564. Py_DECREF(tp);
  565. }
  566. static int
  567. _grouper_traverse(_grouperobject *igo, visitproc visit, void *arg)
  568. {
  569. Py_VISIT(Py_TYPE(igo));
  570. Py_VISIT(igo->parent);
  571. Py_VISIT(igo->tgtkey);
  572. return 0;
  573. }
  574. static PyObject *
  575. _grouper_next(_grouperobject *igo)
  576. {
  577. groupbyobject *gbo = (groupbyobject *)igo->parent;
  578. PyObject *r;
  579. int rcmp;
  580. if (gbo->currgrouper != igo)
  581. return NULL;
  582. if (gbo->currvalue == NULL) {
  583. if (groupby_step(gbo) < 0)
  584. return NULL;
  585. }
  586. assert(gbo->currkey != NULL);
  587. rcmp = PyObject_RichCompareBool(igo->tgtkey, gbo->currkey, Py_EQ);
  588. if (rcmp <= 0)
  589. /* got any error or current group is end */
  590. return NULL;
  591. r = gbo->currvalue;
  592. gbo->currvalue = NULL;
  593. Py_CLEAR(gbo->currkey);
  594. return r;
  595. }
  596. static PyObject *
  597. _grouper_reduce(_grouperobject *lz, PyObject *Py_UNUSED(ignored))
  598. {
  599. ITERTOOL_PICKLE_DEPRECATION;
  600. if (((groupbyobject *)lz->parent)->currgrouper != lz) {
  601. return Py_BuildValue("N(())", _PyEval_GetBuiltin(&_Py_ID(iter)));
  602. }
  603. return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->parent, lz->tgtkey);
  604. }
  605. static PyMethodDef _grouper_methods[] = {
  606. {"__reduce__", (PyCFunction)_grouper_reduce, METH_NOARGS,
  607. reduce_doc},
  608. {NULL, NULL} /* sentinel */
  609. };
  610. static PyType_Slot _grouper_slots[] = {
  611. {Py_tp_dealloc, _grouper_dealloc},
  612. {Py_tp_getattro, PyObject_GenericGetAttr},
  613. {Py_tp_traverse, _grouper_traverse},
  614. {Py_tp_iter, PyObject_SelfIter},
  615. {Py_tp_iternext, _grouper_next},
  616. {Py_tp_methods, _grouper_methods},
  617. {Py_tp_new, itertools__grouper},
  618. {Py_tp_free, PyObject_GC_Del},
  619. {0, NULL},
  620. };
  621. static PyType_Spec _grouper_spec = {
  622. .name = "itertools._grouper",
  623. .basicsize = sizeof(_grouperobject),
  624. .flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
  625. Py_TPFLAGS_IMMUTABLETYPE),
  626. .slots = _grouper_slots,
  627. };
  628. /* tee object and with supporting function and objects ***********************/
  629. /* The teedataobject pre-allocates space for LINKCELLS number of objects.
  630. To help the object fit neatly inside cache lines (space for 16 to 32
  631. pointers), the value should be a multiple of 16 minus space for
  632. the other structure members including PyHEAD overhead. The larger the
  633. value, the less memory overhead per object and the less time spent
  634. allocating/deallocating new links. The smaller the number, the less
  635. wasted space and the more rapid freeing of older data.
  636. */
  637. #define LINKCELLS 57
  638. typedef struct {
  639. PyObject_HEAD
  640. PyObject *it;
  641. int numread; /* 0 <= numread <= LINKCELLS */
  642. int running;
  643. PyObject *nextlink;
  644. PyObject *(values[LINKCELLS]);
  645. } teedataobject;
  646. typedef struct {
  647. PyObject_HEAD
  648. teedataobject *dataobj;
  649. int index; /* 0 <= index <= LINKCELLS */
  650. PyObject *weakreflist;
  651. itertools_state *state;
  652. } teeobject;
  653. static PyObject *
  654. teedataobject_newinternal(itertools_state *state, PyObject *it)
  655. {
  656. teedataobject *tdo;
  657. tdo = PyObject_GC_New(teedataobject, state->teedataobject_type);
  658. if (tdo == NULL)
  659. return NULL;
  660. tdo->running = 0;
  661. tdo->numread = 0;
  662. tdo->nextlink = NULL;
  663. tdo->it = Py_NewRef(it);
  664. PyObject_GC_Track(tdo);
  665. return (PyObject *)tdo;
  666. }
  667. static PyObject *
  668. teedataobject_jumplink(itertools_state *state, teedataobject *tdo)
  669. {
  670. if (tdo->nextlink == NULL)
  671. tdo->nextlink = teedataobject_newinternal(state, tdo->it);
  672. return Py_XNewRef(tdo->nextlink);
  673. }
  674. static PyObject *
  675. teedataobject_getitem(teedataobject *tdo, int i)
  676. {
  677. PyObject *value;
  678. assert(i < LINKCELLS);
  679. if (i < tdo->numread)
  680. value = tdo->values[i];
  681. else {
  682. /* this is the lead iterator, so fetch more data */
  683. assert(i == tdo->numread);
  684. if (tdo->running) {
  685. PyErr_SetString(PyExc_RuntimeError,
  686. "cannot re-enter the tee iterator");
  687. return NULL;
  688. }
  689. tdo->running = 1;
  690. value = PyIter_Next(tdo->it);
  691. tdo->running = 0;
  692. if (value == NULL)
  693. return NULL;
  694. tdo->numread++;
  695. tdo->values[i] = value;
  696. }
  697. return Py_NewRef(value);
  698. }
  699. static int
  700. teedataobject_traverse(teedataobject *tdo, visitproc visit, void * arg)
  701. {
  702. int i;
  703. Py_VISIT(Py_TYPE(tdo));
  704. Py_VISIT(tdo->it);
  705. for (i = 0; i < tdo->numread; i++)
  706. Py_VISIT(tdo->values[i]);
  707. Py_VISIT(tdo->nextlink);
  708. return 0;
  709. }
  710. static void
  711. teedataobject_safe_decref(PyObject *obj)
  712. {
  713. while (obj && Py_REFCNT(obj) == 1) {
  714. PyObject *nextlink = ((teedataobject *)obj)->nextlink;
  715. ((teedataobject *)obj)->nextlink = NULL;
  716. Py_SETREF(obj, nextlink);
  717. }
  718. Py_XDECREF(obj);
  719. }
  720. static int
  721. teedataobject_clear(teedataobject *tdo)
  722. {
  723. int i;
  724. PyObject *tmp;
  725. Py_CLEAR(tdo->it);
  726. for (i=0 ; i<tdo->numread ; i++)
  727. Py_CLEAR(tdo->values[i]);
  728. tmp = tdo->nextlink;
  729. tdo->nextlink = NULL;
  730. teedataobject_safe_decref(tmp);
  731. return 0;
  732. }
  733. static void
  734. teedataobject_dealloc(teedataobject *tdo)
  735. {
  736. PyTypeObject *tp = Py_TYPE(tdo);
  737. PyObject_GC_UnTrack(tdo);
  738. teedataobject_clear(tdo);
  739. PyObject_GC_Del(tdo);
  740. Py_DECREF(tp);
  741. }
  742. static PyObject *
  743. teedataobject_reduce(teedataobject *tdo, PyObject *Py_UNUSED(ignored))
  744. {
  745. ITERTOOL_PICKLE_DEPRECATION;
  746. int i;
  747. /* create a temporary list of already iterated values */
  748. PyObject *values = PyList_New(tdo->numread);
  749. if (!values)
  750. return NULL;
  751. for (i=0 ; i<tdo->numread ; i++) {
  752. Py_INCREF(tdo->values[i]);
  753. PyList_SET_ITEM(values, i, tdo->values[i]);
  754. }
  755. return Py_BuildValue("O(ONO)", Py_TYPE(tdo), tdo->it,
  756. values,
  757. tdo->nextlink ? tdo->nextlink : Py_None);
  758. }
  759. /*[clinic input]
  760. @classmethod
  761. itertools.teedataobject.__new__
  762. iterable as it: object
  763. values: object(subclass_of='&PyList_Type')
  764. next: object
  765. /
  766. Data container common to multiple tee objects.
  767. [clinic start generated code]*/
  768. static PyObject *
  769. itertools_teedataobject_impl(PyTypeObject *type, PyObject *it,
  770. PyObject *values, PyObject *next)
  771. /*[clinic end generated code: output=3343ceb07e08df5e input=be60f2fabd2b72ba]*/
  772. {
  773. teedataobject *tdo;
  774. Py_ssize_t i, len;
  775. itertools_state *state = get_module_state_by_cls(type);
  776. assert(type == state->teedataobject_type);
  777. tdo = (teedataobject *)teedataobject_newinternal(state, it);
  778. if (!tdo)
  779. return NULL;
  780. len = PyList_GET_SIZE(values);
  781. if (len > LINKCELLS)
  782. goto err;
  783. for (i=0; i<len; i++) {
  784. tdo->values[i] = PyList_GET_ITEM(values, i);
  785. Py_INCREF(tdo->values[i]);
  786. }
  787. /* len <= LINKCELLS < INT_MAX */
  788. tdo->numread = Py_SAFE_DOWNCAST(len, Py_ssize_t, int);
  789. if (len == LINKCELLS) {
  790. if (next != Py_None) {
  791. if (!Py_IS_TYPE(next, state->teedataobject_type))
  792. goto err;
  793. assert(tdo->nextlink == NULL);
  794. tdo->nextlink = Py_NewRef(next);
  795. }
  796. } else {
  797. if (next != Py_None)
  798. goto err; /* shouldn't have a next if we are not full */
  799. }
  800. return (PyObject*)tdo;
  801. err:
  802. Py_XDECREF(tdo);
  803. PyErr_SetString(PyExc_ValueError, "Invalid arguments");
  804. return NULL;
  805. }
  806. static PyMethodDef teedataobject_methods[] = {
  807. {"__reduce__", (PyCFunction)teedataobject_reduce, METH_NOARGS,
  808. reduce_doc},
  809. {NULL, NULL} /* sentinel */
  810. };
  811. static PyType_Slot teedataobject_slots[] = {
  812. {Py_tp_dealloc, teedataobject_dealloc},
  813. {Py_tp_getattro, PyObject_GenericGetAttr},
  814. {Py_tp_doc, (void *)itertools_teedataobject__doc__},
  815. {Py_tp_traverse, teedataobject_traverse},
  816. {Py_tp_clear, teedataobject_clear},
  817. {Py_tp_methods, teedataobject_methods},
  818. {Py_tp_new, itertools_teedataobject},
  819. {Py_tp_free, PyObject_GC_Del},
  820. {0, NULL},
  821. };
  822. static PyType_Spec teedataobject_spec = {
  823. .name = "itertools._tee_dataobject",
  824. .basicsize = sizeof(teedataobject),
  825. .flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
  826. Py_TPFLAGS_IMMUTABLETYPE),
  827. .slots = teedataobject_slots,
  828. };
  829. static PyObject *
  830. tee_next(teeobject *to)
  831. {
  832. PyObject *value, *link;
  833. if (to->index >= LINKCELLS) {
  834. link = teedataobject_jumplink(to->state, to->dataobj);
  835. if (link == NULL)
  836. return NULL;
  837. Py_SETREF(to->dataobj, (teedataobject *)link);
  838. to->index = 0;
  839. }
  840. value = teedataobject_getitem(to->dataobj, to->index);
  841. if (value == NULL)
  842. return NULL;
  843. to->index++;
  844. return value;
  845. }
  846. static int
  847. tee_traverse(teeobject *to, visitproc visit, void *arg)
  848. {
  849. Py_VISIT(Py_TYPE(to));
  850. Py_VISIT((PyObject *)to->dataobj);
  851. return 0;
  852. }
  853. static PyObject *
  854. tee_copy(teeobject *to, PyObject *Py_UNUSED(ignored))
  855. {
  856. teeobject *newto;
  857. newto = PyObject_GC_New(teeobject, Py_TYPE(to));
  858. if (newto == NULL)
  859. return NULL;
  860. newto->dataobj = (teedataobject*)Py_NewRef(to->dataobj);
  861. newto->index = to->index;
  862. newto->weakreflist = NULL;
  863. newto->state = to->state;
  864. PyObject_GC_Track(newto);
  865. return (PyObject *)newto;
  866. }
  867. PyDoc_STRVAR(teecopy_doc, "Returns an independent iterator.");
  868. static PyObject *
  869. tee_fromiterable(itertools_state *state, PyObject *iterable)
  870. {
  871. teeobject *to;
  872. PyObject *it;
  873. it = PyObject_GetIter(iterable);
  874. if (it == NULL)
  875. return NULL;
  876. if (PyObject_TypeCheck(it, state->tee_type)) {
  877. to = (teeobject *)tee_copy((teeobject *)it, NULL);
  878. goto done;
  879. }
  880. PyObject *dataobj = teedataobject_newinternal(state, it);
  881. if (!dataobj) {
  882. to = NULL;
  883. goto done;
  884. }
  885. to = PyObject_GC_New(teeobject, state->tee_type);
  886. if (to == NULL) {
  887. Py_DECREF(dataobj);
  888. goto done;
  889. }
  890. to->dataobj = (teedataobject *)dataobj;
  891. to->index = 0;
  892. to->weakreflist = NULL;
  893. to->state = state;
  894. PyObject_GC_Track(to);
  895. done:
  896. Py_DECREF(it);
  897. return (PyObject *)to;
  898. }
  899. /*[clinic input]
  900. @classmethod
  901. itertools._tee.__new__
  902. iterable: object
  903. /
  904. Iterator wrapped to make it copyable.
  905. [clinic start generated code]*/
  906. static PyObject *
  907. itertools__tee_impl(PyTypeObject *type, PyObject *iterable)
  908. /*[clinic end generated code: output=b02d3fd26c810c3f input=adc0779d2afe37a2]*/
  909. {
  910. itertools_state *state = get_module_state_by_cls(type);
  911. return tee_fromiterable(state, iterable);
  912. }
  913. static int
  914. tee_clear(teeobject *to)
  915. {
  916. if (to->weakreflist != NULL)
  917. PyObject_ClearWeakRefs((PyObject *) to);
  918. Py_CLEAR(to->dataobj);
  919. return 0;
  920. }
  921. static void
  922. tee_dealloc(teeobject *to)
  923. {
  924. PyTypeObject *tp = Py_TYPE(to);
  925. PyObject_GC_UnTrack(to);
  926. tee_clear(to);
  927. PyObject_GC_Del(to);
  928. Py_DECREF(tp);
  929. }
  930. static PyObject *
  931. tee_reduce(teeobject *to, PyObject *Py_UNUSED(ignored))
  932. {
  933. ITERTOOL_PICKLE_DEPRECATION;
  934. return Py_BuildValue("O(())(Oi)", Py_TYPE(to), to->dataobj, to->index);
  935. }
  936. static PyObject *
  937. tee_setstate(teeobject *to, PyObject *state)
  938. {
  939. ITERTOOL_PICKLE_DEPRECATION;
  940. teedataobject *tdo;
  941. int index;
  942. if (!PyTuple_Check(state)) {
  943. PyErr_SetString(PyExc_TypeError, "state is not a tuple");
  944. return NULL;
  945. }
  946. PyTypeObject *tdo_type = to->state->teedataobject_type;
  947. if (!PyArg_ParseTuple(state, "O!i", tdo_type, &tdo, &index)) {
  948. return NULL;
  949. }
  950. if (index < 0 || index > LINKCELLS) {
  951. PyErr_SetString(PyExc_ValueError, "Index out of range");
  952. return NULL;
  953. }
  954. Py_INCREF(tdo);
  955. Py_XSETREF(to->dataobj, tdo);
  956. to->index = index;
  957. Py_RETURN_NONE;
  958. }
  959. static PyMethodDef tee_methods[] = {
  960. {"__copy__", (PyCFunction)tee_copy, METH_NOARGS, teecopy_doc},
  961. {"__reduce__", (PyCFunction)tee_reduce, METH_NOARGS, reduce_doc},
  962. {"__setstate__", (PyCFunction)tee_setstate, METH_O, setstate_doc},
  963. {NULL, NULL} /* sentinel */
  964. };
  965. static PyMemberDef tee_members[] = {
  966. {"__weaklistoffset__", T_PYSSIZET, offsetof(teeobject, weakreflist), READONLY},
  967. {NULL},
  968. };
  969. static PyType_Slot tee_slots[] = {
  970. {Py_tp_dealloc, tee_dealloc},
  971. {Py_tp_doc, (void *)itertools__tee__doc__},
  972. {Py_tp_traverse, tee_traverse},
  973. {Py_tp_clear, tee_clear},
  974. {Py_tp_iter, PyObject_SelfIter},
  975. {Py_tp_iternext, tee_next},
  976. {Py_tp_methods, tee_methods},
  977. {Py_tp_members, tee_members},
  978. {Py_tp_new, itertools__tee},
  979. {Py_tp_free, PyObject_GC_Del},
  980. {0, NULL},
  981. };
  982. static PyType_Spec tee_spec = {
  983. .name = "itertools._tee",
  984. .basicsize = sizeof(teeobject),
  985. .flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
  986. Py_TPFLAGS_IMMUTABLETYPE),
  987. .slots = tee_slots,
  988. };
  989. /*[clinic input]
  990. itertools.tee
  991. iterable: object
  992. n: Py_ssize_t = 2
  993. /
  994. Returns a tuple of n independent iterators.
  995. [clinic start generated code]*/
  996. static PyObject *
  997. itertools_tee_impl(PyObject *module, PyObject *iterable, Py_ssize_t n)
  998. /*[clinic end generated code: output=1c64519cd859c2f0 input=c99a1472c425d66d]*/
  999. {
  1000. Py_ssize_t i;
  1001. PyObject *it, *copyable, *copyfunc, *result;
  1002. if (n < 0) {
  1003. PyErr_SetString(PyExc_ValueError, "n must be >= 0");
  1004. return NULL;
  1005. }
  1006. result = PyTuple_New(n);
  1007. if (result == NULL)
  1008. return NULL;
  1009. if (n == 0)
  1010. return result;
  1011. it = PyObject_GetIter(iterable);
  1012. if (it == NULL) {
  1013. Py_DECREF(result);
  1014. return NULL;
  1015. }
  1016. if (_PyObject_LookupAttr(it, &_Py_ID(__copy__), &copyfunc) < 0) {
  1017. Py_DECREF(it);
  1018. Py_DECREF(result);
  1019. return NULL;
  1020. }
  1021. if (copyfunc != NULL) {
  1022. copyable = it;
  1023. }
  1024. else {
  1025. itertools_state *state = get_module_state(module);
  1026. copyable = tee_fromiterable(state, it);
  1027. Py_DECREF(it);
  1028. if (copyable == NULL) {
  1029. Py_DECREF(result);
  1030. return NULL;
  1031. }
  1032. copyfunc = PyObject_GetAttr(copyable, &_Py_ID(__copy__));
  1033. if (copyfunc == NULL) {
  1034. Py_DECREF(copyable);
  1035. Py_DECREF(result);
  1036. return NULL;
  1037. }
  1038. }
  1039. PyTuple_SET_ITEM(result, 0, copyable);
  1040. for (i = 1; i < n; i++) {
  1041. copyable = _PyObject_CallNoArgs(copyfunc);
  1042. if (copyable == NULL) {
  1043. Py_DECREF(copyfunc);
  1044. Py_DECREF(result);
  1045. return NULL;
  1046. }
  1047. PyTuple_SET_ITEM(result, i, copyable);
  1048. }
  1049. Py_DECREF(copyfunc);
  1050. return result;
  1051. }
  1052. /* cycle object **************************************************************/
  1053. typedef struct {
  1054. PyObject_HEAD
  1055. PyObject *it;
  1056. PyObject *saved;
  1057. Py_ssize_t index;
  1058. int firstpass;
  1059. } cycleobject;
  1060. /*[clinic input]
  1061. @classmethod
  1062. itertools.cycle.__new__
  1063. iterable: object
  1064. /
  1065. Return elements from the iterable until it is exhausted. Then repeat the sequence indefinitely.
  1066. [clinic start generated code]*/
  1067. static PyObject *
  1068. itertools_cycle_impl(PyTypeObject *type, PyObject *iterable)
  1069. /*[clinic end generated code: output=f60e5ec17a45b35c input=9d1d84bcf66e908b]*/
  1070. {
  1071. PyObject *it;
  1072. PyObject *saved;
  1073. cycleobject *lz;
  1074. /* Get iterator. */
  1075. it = PyObject_GetIter(iterable);
  1076. if (it == NULL)
  1077. return NULL;
  1078. saved = PyList_New(0);
  1079. if (saved == NULL) {
  1080. Py_DECREF(it);
  1081. return NULL;
  1082. }
  1083. /* create cycleobject structure */
  1084. lz = (cycleobject *)type->tp_alloc(type, 0);
  1085. if (lz == NULL) {
  1086. Py_DECREF(it);
  1087. Py_DECREF(saved);
  1088. return NULL;
  1089. }
  1090. lz->it = it;
  1091. lz->saved = saved;
  1092. lz->index = 0;
  1093. lz->firstpass = 0;
  1094. return (PyObject *)lz;
  1095. }
  1096. static void
  1097. cycle_dealloc(cycleobject *lz)
  1098. {
  1099. PyTypeObject *tp = Py_TYPE(lz);
  1100. PyObject_GC_UnTrack(lz);
  1101. Py_XDECREF(lz->it);
  1102. Py_XDECREF(lz->saved);
  1103. tp->tp_free(lz);
  1104. Py_DECREF(tp);
  1105. }
  1106. static int
  1107. cycle_traverse(cycleobject *lz, visitproc visit, void *arg)
  1108. {
  1109. Py_VISIT(Py_TYPE(lz));
  1110. Py_VISIT(lz->it);
  1111. Py_VISIT(lz->saved);
  1112. return 0;
  1113. }
  1114. static PyObject *
  1115. cycle_next(cycleobject *lz)
  1116. {
  1117. PyObject *item;
  1118. if (lz->it != NULL) {
  1119. item = PyIter_Next(lz->it);
  1120. if (item != NULL) {
  1121. if (lz->firstpass)
  1122. return item;
  1123. if (PyList_Append(lz->saved, item)) {
  1124. Py_DECREF(item);
  1125. return NULL;
  1126. }
  1127. return item;
  1128. }
  1129. /* Note: StopIteration is already cleared by PyIter_Next() */
  1130. if (PyErr_Occurred())
  1131. return NULL;
  1132. Py_CLEAR(lz->it);
  1133. }
  1134. if (PyList_GET_SIZE(lz->saved) == 0)
  1135. return NULL;
  1136. item = PyList_GET_ITEM(lz->saved, lz->index);
  1137. lz->index++;
  1138. if (lz->index >= PyList_GET_SIZE(lz->saved))
  1139. lz->index = 0;
  1140. return Py_NewRef(item);
  1141. }
  1142. static PyObject *
  1143. cycle_reduce(cycleobject *lz, PyObject *Py_UNUSED(ignored))
  1144. {
  1145. ITERTOOL_PICKLE_DEPRECATION;
  1146. /* Create a new cycle with the iterator tuple, then set the saved state */
  1147. if (lz->it == NULL) {
  1148. PyObject *it = PyObject_GetIter(lz->saved);
  1149. if (it == NULL)
  1150. return NULL;
  1151. if (lz->index != 0) {
  1152. PyObject *res = _PyObject_CallMethod(it, &_Py_ID(__setstate__),
  1153. "n", lz->index);
  1154. if (res == NULL) {
  1155. Py_DECREF(it);
  1156. return NULL;
  1157. }
  1158. Py_DECREF(res);
  1159. }
  1160. return Py_BuildValue("O(N)(OO)", Py_TYPE(lz), it, lz->saved, Py_True);
  1161. }
  1162. return Py_BuildValue("O(O)(OO)", Py_TYPE(lz), lz->it, lz->saved,
  1163. lz->firstpass ? Py_True : Py_False);
  1164. }
  1165. static PyObject *
  1166. cycle_setstate(cycleobject *lz, PyObject *state)
  1167. {
  1168. ITERTOOL_PICKLE_DEPRECATION;
  1169. PyObject *saved=NULL;
  1170. int firstpass;
  1171. if (!PyTuple_Check(state)) {
  1172. PyErr_SetString(PyExc_TypeError, "state is not a tuple");
  1173. return NULL;
  1174. }
  1175. // The second item can be 1/0 in old pickles and True/False in new pickles
  1176. if (!PyArg_ParseTuple(state, "O!i", &PyList_Type, &saved, &firstpass)) {
  1177. return NULL;
  1178. }
  1179. Py_INCREF(saved);
  1180. Py_XSETREF(lz->saved, saved);
  1181. lz->firstpass = firstpass != 0;
  1182. lz->index = 0;
  1183. Py_RETURN_NONE;
  1184. }
  1185. static PyMethodDef cycle_methods[] = {
  1186. {"__reduce__", (PyCFunction)cycle_reduce, METH_NOARGS,
  1187. reduce_doc},
  1188. {"__setstate__", (PyCFunction)cycle_setstate, METH_O,
  1189. setstate_doc},
  1190. {NULL, NULL} /* sentinel */
  1191. };
  1192. static PyType_Slot cycle_slots[] = {
  1193. {Py_tp_dealloc, cycle_dealloc},
  1194. {Py_tp_getattro, PyObject_GenericGetAttr},
  1195. {Py_tp_doc, (void *)itertools_cycle__doc__},
  1196. {Py_tp_traverse, cycle_traverse},
  1197. {Py_tp_iter, PyObject_SelfIter},
  1198. {Py_tp_iternext, cycle_next},
  1199. {Py_tp_methods, cycle_methods},
  1200. {Py_tp_new, itertools_cycle},
  1201. {Py_tp_free, PyObject_GC_Del},
  1202. {0, NULL},
  1203. };
  1204. static PyType_Spec cycle_spec = {
  1205. .name = "itertools.cycle",
  1206. .basicsize = sizeof(cycleobject),
  1207. .flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_BASETYPE |
  1208. Py_TPFLAGS_IMMUTABLETYPE),
  1209. .slots = cycle_slots,
  1210. };
  1211. /* dropwhile object **********************************************************/
  1212. typedef struct {
  1213. PyObject_HEAD
  1214. PyObject *func;
  1215. PyObject *it;
  1216. long start;
  1217. } dropwhileobject;
  1218. /*[clinic input]
  1219. @classmethod
  1220. itertools.dropwhile.__new__
  1221. predicate as func: object
  1222. iterable as seq: object
  1223. /
  1224. Drop items from the iterable while predicate(item) is true.
  1225. Afterwards, return every element until the iterable is exhausted.
  1226. [clinic start generated code]*/
  1227. static PyObject *
  1228. itertools_dropwhile_impl(PyTypeObject *type, PyObject *func, PyObject *seq)
  1229. /*[clinic end generated code: output=92f9d0d89af149e4 input=d39737147c9f0a26]*/
  1230. {
  1231. PyObject *it;
  1232. dropwhileobject *lz;
  1233. /* Get iterator. */
  1234. it = PyObject_GetIter(seq);
  1235. if (it == NULL)
  1236. return NULL;
  1237. /* create dropwhileobject structure */
  1238. lz = (dropwhileobject *)type->tp_alloc(type, 0);
  1239. if (lz == NULL) {
  1240. Py_DECREF(it);
  1241. return NULL;
  1242. }
  1243. lz->func = Py_NewRef(func);
  1244. lz->it = it;
  1245. lz->start = 0;
  1246. return (PyObject *)lz;
  1247. }
  1248. static void
  1249. dropwhile_dealloc(dropwhileobject *lz)
  1250. {
  1251. PyTypeObject *tp = Py_TYPE(lz);
  1252. PyObject_GC_UnTrack(lz);
  1253. Py_XDECREF(lz->func);
  1254. Py_XDECREF(lz->it);
  1255. tp->tp_free(lz);
  1256. Py_DECREF(tp);
  1257. }
  1258. static int
  1259. dropwhile_traverse(dropwhileobject *lz, visitproc visit, void *arg)
  1260. {
  1261. Py_VISIT(Py_TYPE(lz));
  1262. Py_VISIT(lz->it);
  1263. Py_VISIT(lz->func);
  1264. return 0;
  1265. }
  1266. static PyObject *
  1267. dropwhile_next(dropwhileobject *lz)
  1268. {
  1269. PyObject *item, *good;
  1270. PyObject *it = lz->it;
  1271. long ok;
  1272. PyObject *(*iternext)(PyObject *);
  1273. iternext = *Py_TYPE(it)->tp_iternext;
  1274. for (;;) {
  1275. item = iternext(it);
  1276. if (item == NULL)
  1277. return NULL;
  1278. if (lz->start == 1)
  1279. return item;
  1280. good = PyObject_CallOneArg(lz->func, item);
  1281. if (good == NULL) {
  1282. Py_DECREF(item);
  1283. return NULL;
  1284. }
  1285. ok = PyObject_IsTrue(good);
  1286. Py_DECREF(good);
  1287. if (ok == 0) {
  1288. lz->start = 1;
  1289. return item;
  1290. }
  1291. Py_DECREF(item);
  1292. if (ok < 0)
  1293. return NULL;
  1294. }
  1295. }
  1296. static PyObject *
  1297. dropwhile_reduce(dropwhileobject *lz, PyObject *Py_UNUSED(ignored))
  1298. {
  1299. ITERTOOL_PICKLE_DEPRECATION;
  1300. return Py_BuildValue("O(OO)l", Py_TYPE(lz), lz->func, lz->it, lz->start);
  1301. }
  1302. static PyObject *
  1303. dropwhile_setstate(dropwhileobject *lz, PyObject *state)
  1304. {
  1305. ITERTOOL_PICKLE_DEPRECATION;
  1306. int start = PyObject_IsTrue(state);
  1307. if (start < 0)
  1308. return NULL;
  1309. lz->start = start;
  1310. Py_RETURN_NONE;
  1311. }
  1312. static PyMethodDef dropwhile_methods[] = {
  1313. {"__reduce__", (PyCFunction)dropwhile_reduce, METH_NOARGS,
  1314. reduce_doc},
  1315. {"__setstate__", (PyCFunction)dropwhile_setstate, METH_O,
  1316. setstate_doc},
  1317. {NULL, NULL} /* sentinel */
  1318. };
  1319. static PyType_Slot dropwhile_slots[] = {
  1320. {Py_tp_dealloc, dropwhile_dealloc},
  1321. {Py_tp_getattro, PyObject_GenericGetAttr},
  1322. {Py_tp_doc, (void *)itertools_dropwhile__doc__},
  1323. {Py_tp_traverse, dropwhile_traverse},
  1324. {Py_tp_iter, PyObject_SelfIter},
  1325. {Py_tp_iternext, dropwhile_next},
  1326. {Py_tp_methods, dropwhile_methods},
  1327. {Py_tp_new, itertools_dropwhile},
  1328. {Py_tp_free, PyObject_GC_Del},
  1329. {0, NULL},
  1330. };
  1331. static PyType_Spec dropwhile_spec = {
  1332. .name = "itertools.dropwhile",
  1333. .basicsize = sizeof(dropwhileobject),
  1334. .flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_BASETYPE |
  1335. Py_TPFLAGS_IMMUTABLETYPE),
  1336. .slots = dropwhile_slots,
  1337. };
  1338. /* takewhile object **********************************************************/
  1339. typedef struct {
  1340. PyObject_HEAD
  1341. PyObject *func;
  1342. PyObject *it;
  1343. long stop;
  1344. } takewhileobject;
  1345. /*[clinic input]
  1346. @classmethod
  1347. itertools.takewhile.__new__
  1348. predicate as func: object
  1349. iterable as seq: object
  1350. /
  1351. Return successive entries from an iterable as long as the predicate evaluates to true for each entry.
  1352. [clinic start generated code]*/
  1353. static PyObject *
  1354. itertools_takewhile_impl(PyTypeObject *type, PyObject *func, PyObject *seq)
  1355. /*[clinic end generated code: output=bb179ea7864e2ef6 input=ba5255f7519aa119]*/
  1356. {
  1357. PyObject *it;
  1358. takewhileobject *lz;
  1359. /* Get iterator. */
  1360. it = PyObject_GetIter(seq);
  1361. if (it == NULL)
  1362. return NULL;
  1363. /* create takewhileobject structure */
  1364. lz = (takewhileobject *)type->tp_alloc(type, 0);
  1365. if (lz == NULL) {
  1366. Py_DECREF(it);
  1367. return NULL;
  1368. }
  1369. lz->func = Py_NewRef(func);
  1370. lz->it = it;
  1371. lz->stop = 0;
  1372. return (PyObject *)lz;
  1373. }
  1374. static void
  1375. takewhile_dealloc(takewhileobject *lz)
  1376. {
  1377. PyTypeObject *tp = Py_TYPE(lz);
  1378. PyObject_GC_UnTrack(lz);
  1379. Py_XDECREF(lz->func);
  1380. Py_XDECREF(lz->it);
  1381. tp->tp_free(lz);
  1382. Py_DECREF(tp);
  1383. }
  1384. static int
  1385. takewhile_traverse(takewhileobject *lz, visitproc visit, void *arg)
  1386. {
  1387. Py_VISIT(Py_TYPE(lz));
  1388. Py_VISIT(lz->it);
  1389. Py_VISIT(lz->func);
  1390. return 0;
  1391. }
  1392. static PyObject *
  1393. takewhile_next(takewhileobject *lz)
  1394. {
  1395. PyObject *item, *good;
  1396. PyObject *it = lz->it;
  1397. long ok;
  1398. if (lz->stop == 1)
  1399. return NULL;
  1400. item = (*Py_TYPE(it)->tp_iternext)(it);
  1401. if (item == NULL)
  1402. return NULL;
  1403. good = PyObject_CallOneArg(lz->func, item);
  1404. if (good == NULL) {
  1405. Py_DECREF(item);
  1406. return NULL;
  1407. }
  1408. ok = PyObject_IsTrue(good);
  1409. Py_DECREF(good);
  1410. if (ok > 0)
  1411. return item;
  1412. Py_DECREF(item);
  1413. if (ok == 0)
  1414. lz->stop = 1;
  1415. return NULL;
  1416. }
  1417. static PyObject *
  1418. takewhile_reduce(takewhileobject *lz, PyObject *Py_UNUSED(ignored))
  1419. {
  1420. ITERTOOL_PICKLE_DEPRECATION;
  1421. return Py_BuildValue("O(OO)l", Py_TYPE(lz), lz->func, lz->it, lz->stop);
  1422. }
  1423. static PyObject *
  1424. takewhile_reduce_setstate(takewhileobject *lz, PyObject *state)
  1425. {
  1426. ITERTOOL_PICKLE_DEPRECATION;
  1427. int stop = PyObject_IsTrue(state);
  1428. if (stop < 0)
  1429. return NULL;
  1430. lz->stop = stop;
  1431. Py_RETURN_NONE;
  1432. }
  1433. static PyMethodDef takewhile_reduce_methods[] = {
  1434. {"__reduce__", (PyCFunction)takewhile_reduce, METH_NOARGS,
  1435. reduce_doc},
  1436. {"__setstate__", (PyCFunction)takewhile_reduce_setstate, METH_O,
  1437. setstate_doc},
  1438. {NULL, NULL} /* sentinel */
  1439. };
  1440. static PyType_Slot takewhile_slots[] = {
  1441. {Py_tp_dealloc, takewhile_dealloc},
  1442. {Py_tp_getattro, PyObject_GenericGetAttr},
  1443. {Py_tp_doc, (void *)itertools_takewhile__doc__},
  1444. {Py_tp_traverse, takewhile_traverse},
  1445. {Py_tp_iter, PyObject_SelfIter},
  1446. {Py_tp_iternext, takewhile_next},
  1447. {Py_tp_methods, takewhile_reduce_methods},
  1448. {Py_tp_new, itertools_takewhile},
  1449. {Py_tp_free, PyObject_GC_Del},
  1450. {0, NULL},
  1451. };
  1452. static PyType_Spec takewhile_spec = {
  1453. .name = "itertools.takewhile",
  1454. .basicsize = sizeof(takewhileobject),
  1455. .flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_BASETYPE |
  1456. Py_TPFLAGS_IMMUTABLETYPE),
  1457. .slots = takewhile_slots,
  1458. };
  1459. /* islice object *************************************************************/
  1460. typedef struct {
  1461. PyObject_HEAD
  1462. PyObject *it;
  1463. Py_ssize_t next;
  1464. Py_ssize_t stop;
  1465. Py_ssize_t step;
  1466. Py_ssize_t cnt;
  1467. } isliceobject;
  1468. static PyObject *
  1469. islice_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
  1470. {
  1471. PyObject *seq;
  1472. Py_ssize_t start=0, stop=-1, step=1;
  1473. PyObject *it, *a1=NULL, *a2=NULL, *a3=NULL;
  1474. Py_ssize_t numargs;
  1475. isliceobject *lz;
  1476. itertools_state *st = find_state_by_type(type);
  1477. PyTypeObject *islice_type = st->islice_type;
  1478. if ((type == islice_type || type->tp_init == islice_type->tp_init) &&
  1479. !_PyArg_NoKeywords("islice", kwds))
  1480. return NULL;
  1481. if (!PyArg_UnpackTuple(args, "islice", 2, 4, &seq, &a1, &a2, &a3))
  1482. return NULL;
  1483. numargs = PyTuple_Size(args);
  1484. if (numargs == 2) {
  1485. if (a1 != Py_None) {
  1486. stop = PyNumber_AsSsize_t(a1, PyExc_OverflowError);
  1487. if (stop == -1) {
  1488. if (PyErr_Occurred())
  1489. PyErr_Clear();
  1490. PyErr_SetString(PyExc_ValueError,
  1491. "Stop argument for islice() must be None or "
  1492. "an integer: 0 <= x <= sys.maxsize.");
  1493. return NULL;
  1494. }
  1495. }
  1496. } else {
  1497. if (a1 != Py_None)
  1498. start = PyNumber_AsSsize_t(a1, PyExc_OverflowError);
  1499. if (start == -1 && PyErr_Occurred())
  1500. PyErr_Clear();
  1501. if (a2 != Py_None) {
  1502. stop = PyNumber_AsSsize_t(a2, PyExc_OverflowError);
  1503. if (stop == -1) {
  1504. if (PyErr_Occurred())
  1505. PyErr_Clear();
  1506. PyErr_SetString(PyExc_ValueError,
  1507. "Stop argument for islice() must be None or "
  1508. "an integer: 0 <= x <= sys.maxsize.");
  1509. return NULL;
  1510. }
  1511. }
  1512. }
  1513. if (start<0 || stop<-1) {
  1514. PyErr_SetString(PyExc_ValueError,
  1515. "Indices for islice() must be None or "
  1516. "an integer: 0 <= x <= sys.maxsize.");
  1517. return NULL;
  1518. }
  1519. if (a3 != NULL) {
  1520. if (a3 != Py_None)
  1521. step = PyNumber_AsSsize_t(a3, PyExc_OverflowError);
  1522. if (step == -1 && PyErr_Occurred())
  1523. PyErr_Clear();
  1524. }
  1525. if (step<1) {
  1526. PyErr_SetString(PyExc_ValueError,
  1527. "Step for islice() must be a positive integer or None.");
  1528. return NULL;
  1529. }
  1530. /* Get iterator. */
  1531. it = PyObject_GetIter(seq);
  1532. if (it == NULL)
  1533. return NULL;
  1534. /* create isliceobject structure */
  1535. lz = (isliceobject *)type->tp_alloc(type, 0);
  1536. if (lz == NULL) {
  1537. Py_DECREF(it);
  1538. return NULL;
  1539. }
  1540. lz->it = it;
  1541. lz->next = start;
  1542. lz->stop = stop;
  1543. lz->step = step;
  1544. lz->cnt = 0L;
  1545. return (PyObject *)lz;
  1546. }
  1547. static void
  1548. islice_dealloc(isliceobject *lz)
  1549. {
  1550. PyTypeObject *tp = Py_TYPE(lz);
  1551. PyObject_GC_UnTrack(lz);
  1552. Py_XDECREF(lz->it);
  1553. tp->tp_free(lz);
  1554. Py_DECREF(tp);
  1555. }
  1556. static int
  1557. islice_traverse(isliceobject *lz, visitproc visit, void *arg)
  1558. {
  1559. Py_VISIT(Py_TYPE(lz));
  1560. Py_VISIT(lz->it);
  1561. return 0;
  1562. }
  1563. static PyObject *
  1564. islice_next(isliceobject *lz)
  1565. {
  1566. PyObject *item;
  1567. PyObject *it = lz->it;
  1568. Py_ssize_t stop = lz->stop;
  1569. Py_ssize_t oldnext;
  1570. PyObject *(*iternext)(PyObject *);
  1571. if (it == NULL)
  1572. return NULL;
  1573. iternext = *Py_TYPE(it)->tp_iternext;
  1574. while (lz->cnt < lz->next) {
  1575. item = iternext(it);
  1576. if (item == NULL)
  1577. goto empty;
  1578. Py_DECREF(item);
  1579. lz->cnt++;
  1580. }
  1581. if (stop != -1 && lz->cnt >= stop)
  1582. goto empty;
  1583. item = iternext(it);
  1584. if (item == NULL)
  1585. goto empty;
  1586. lz->cnt++;
  1587. oldnext = lz->next;
  1588. /* The (size_t) cast below avoids the danger of undefined
  1589. behaviour from signed integer overflow. */
  1590. lz->next += (size_t)lz->step;
  1591. if (lz->next < oldnext || (stop != -1 && lz->next > stop))
  1592. lz->next = stop;
  1593. return item;
  1594. empty:
  1595. Py_CLEAR(lz->it);
  1596. return NULL;
  1597. }
  1598. static PyObject *
  1599. islice_reduce(isliceobject *lz, PyObject *Py_UNUSED(ignored))
  1600. {
  1601. ITERTOOL_PICKLE_DEPRECATION;
  1602. /* When unpickled, generate a new object with the same bounds,
  1603. * then 'setstate' with the next and count
  1604. */
  1605. PyObject *stop;
  1606. if (lz->it == NULL) {
  1607. PyObject *empty_list;
  1608. PyObject *empty_it;
  1609. empty_list = PyList_New(0);
  1610. if (empty_list == NULL)
  1611. return NULL;
  1612. empty_it = PyObject_GetIter(empty_list);
  1613. Py_DECREF(empty_list);
  1614. if (empty_it == NULL)
  1615. return NULL;
  1616. return Py_BuildValue("O(Nn)n", Py_TYPE(lz), empty_it, 0, 0);
  1617. }
  1618. if (lz->stop == -1) {
  1619. stop = Py_NewRef(Py_None);
  1620. } else {
  1621. stop = PyLong_FromSsize_t(lz->stop);
  1622. if (stop == NULL)
  1623. return NULL;
  1624. }
  1625. return Py_BuildValue("O(OnNn)n", Py_TYPE(lz),
  1626. lz->it, lz->next, stop, lz->step,
  1627. lz->cnt);
  1628. }
  1629. static PyObject *
  1630. islice_setstate(isliceobject *lz, PyObject *state)
  1631. {
  1632. ITERTOOL_PICKLE_DEPRECATION;
  1633. Py_ssize_t cnt = PyLong_AsSsize_t(state);
  1634. if (cnt == -1 && PyErr_Occurred())
  1635. return NULL;
  1636. lz->cnt = cnt;
  1637. Py_RETURN_NONE;
  1638. }
  1639. static PyMethodDef islice_methods[] = {
  1640. {"__reduce__", (PyCFunction)islice_reduce, METH_NOARGS,
  1641. reduce_doc},
  1642. {"__setstate__", (PyCFunction)islice_setstate, METH_O,
  1643. setstate_doc},
  1644. {NULL, NULL} /* sentinel */
  1645. };
  1646. PyDoc_STRVAR(islice_doc,
  1647. "islice(iterable, stop) --> islice object\n\
  1648. islice(iterable, start, stop[, step]) --> islice object\n\
  1649. \n\
  1650. Return an iterator whose next() method returns selected values from an\n\
  1651. iterable. If start is specified, will skip all preceding elements;\n\
  1652. otherwise, start defaults to zero. Step defaults to one. If\n\
  1653. specified as another value, step determines how many values are\n\
  1654. skipped between successive calls. Works like a slice() on a list\n\
  1655. but returns an iterator.");
  1656. static PyType_Slot islice_slots[] = {
  1657. {Py_tp_dealloc, islice_dealloc},
  1658. {Py_tp_getattro, PyObject_GenericGetAttr},
  1659. {Py_tp_doc, (void *)islice_doc},
  1660. {Py_tp_traverse, islice_traverse},
  1661. {Py_tp_iter, PyObject_SelfIter},
  1662. {Py_tp_iternext, islice_next},
  1663. {Py_tp_methods, islice_methods},
  1664. {Py_tp_new, islice_new},
  1665. {Py_tp_free, PyObject_GC_Del},
  1666. {0, NULL},
  1667. };
  1668. static PyType_Spec islice_spec = {
  1669. .name = "itertools.islice",
  1670. .basicsize = sizeof(isliceobject),
  1671. .flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_BASETYPE |
  1672. Py_TPFLAGS_IMMUTABLETYPE),
  1673. .slots = islice_slots,
  1674. };
  1675. /* starmap object ************************************************************/
  1676. typedef struct {
  1677. PyObject_HEAD
  1678. PyObject *func;
  1679. PyObject *it;
  1680. } starmapobject;
  1681. /*[clinic input]
  1682. @classmethod
  1683. itertools.starmap.__new__
  1684. function as func: object
  1685. iterable as seq: object
  1686. /
  1687. Return an iterator whose values are returned from the function evaluated with an argument tuple taken from the given sequence.
  1688. [clinic start generated code]*/
  1689. static PyObject *
  1690. itertools_starmap_impl(PyTypeObject *type, PyObject *func, PyObject *seq)
  1691. /*[clinic end generated code: output=79eeb81d452c6e8d input=844766df6a0d4dad]*/
  1692. {
  1693. PyObject *it;
  1694. starmapobject *lz;
  1695. /* Get iterator. */
  1696. it = PyObject_GetIter(seq);
  1697. if (it == NULL)
  1698. return NULL;
  1699. /* create starmapobject structure */
  1700. lz = (starmapobject *)type->tp_alloc(type, 0);
  1701. if (lz == NULL) {
  1702. Py_DECREF(it);
  1703. return NULL;
  1704. }
  1705. lz->func = Py_NewRef(func);
  1706. lz->it = it;
  1707. return (PyObject *)lz;
  1708. }
  1709. static void
  1710. starmap_dealloc(starmapobject *lz)
  1711. {
  1712. PyTypeObject *tp = Py_TYPE(lz);
  1713. PyObject_GC_UnTrack(lz);
  1714. Py_XDECREF(lz->func);
  1715. Py_XDECREF(lz->it);
  1716. tp->tp_free(lz);
  1717. Py_DECREF(tp);
  1718. }
  1719. static int
  1720. starmap_traverse(starmapobject *lz, visitproc visit, void *arg)
  1721. {
  1722. Py_VISIT(Py_TYPE(lz));
  1723. Py_VISIT(lz->it);
  1724. Py_VISIT(lz->func);
  1725. return 0;
  1726. }
  1727. static PyObject *
  1728. starmap_next(starmapobject *lz)
  1729. {
  1730. PyObject *args;
  1731. PyObject *result;
  1732. PyObject *it = lz->it;
  1733. args = (*Py_TYPE(it)->tp_iternext)(it);
  1734. if (args == NULL)
  1735. return NULL;
  1736. if (!PyTuple_CheckExact(args)) {
  1737. PyObject *newargs = PySequence_Tuple(args);
  1738. Py_DECREF(args);
  1739. if (newargs == NULL)
  1740. return NULL;
  1741. args = newargs;
  1742. }
  1743. result = PyObject_Call(lz->func, args, NULL);
  1744. Py_DECREF(args);
  1745. return result;
  1746. }
  1747. static PyObject *
  1748. starmap_reduce(starmapobject *lz, PyObject *Py_UNUSED(ignored))
  1749. {
  1750. ITERTOOL_PICKLE_DEPRECATION;
  1751. /* Just pickle the iterator */
  1752. return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->func, lz->it);
  1753. }
  1754. static PyMethodDef starmap_methods[] = {
  1755. {"__reduce__", (PyCFunction)starmap_reduce, METH_NOARGS,
  1756. reduce_doc},
  1757. {NULL, NULL} /* sentinel */
  1758. };
  1759. static PyType_Slot starmap_slots[] = {
  1760. {Py_tp_dealloc, starmap_dealloc},
  1761. {Py_tp_getattro, PyObject_GenericGetAttr},
  1762. {Py_tp_doc, (void *)itertools_starmap__doc__},
  1763. {Py_tp_traverse, starmap_traverse},
  1764. {Py_tp_iter, PyObject_SelfIter},
  1765. {Py_tp_iternext, starmap_next},
  1766. {Py_tp_methods, starmap_methods},
  1767. {Py_tp_new, itertools_starmap},
  1768. {Py_tp_free, PyObject_GC_Del},
  1769. {0, NULL},
  1770. };
  1771. static PyType_Spec starmap_spec = {
  1772. .name = "itertools.starmap",
  1773. .basicsize = sizeof(starmapobject),
  1774. .flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_BASETYPE |
  1775. Py_TPFLAGS_IMMUTABLETYPE),
  1776. .slots = starmap_slots,
  1777. };
  1778. /* chain object **************************************************************/
  1779. typedef struct {
  1780. PyObject_HEAD
  1781. PyObject *source; /* Iterator over input iterables */
  1782. PyObject *active; /* Currently running input iterator */
  1783. } chainobject;
  1784. static PyObject *
  1785. chain_new_internal(PyTypeObject *type, PyObject *source)
  1786. {
  1787. chainobject *lz;
  1788. lz = (chainobject *)type->tp_alloc(type, 0);
  1789. if (lz == NULL) {
  1790. Py_DECREF(source);
  1791. return NULL;
  1792. }
  1793. lz->source = source;
  1794. lz->active = NULL;
  1795. return (PyObject *)lz;
  1796. }
  1797. static PyObject *
  1798. chain_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
  1799. {
  1800. PyObject *source;
  1801. itertools_state *state = find_state_by_type(type);
  1802. PyTypeObject *chain_type = state->chain_type;
  1803. if ((type == chain_type || type->tp_init == chain_type->tp_init) &&
  1804. !_PyArg_NoKeywords("chain", kwds))
  1805. return NULL;
  1806. source = PyObject_GetIter(args);
  1807. if (source == NULL)
  1808. return NULL;
  1809. return chain_new_internal(type, source);
  1810. }
  1811. /*[clinic input]
  1812. @classmethod
  1813. itertools.chain.from_iterable
  1814. iterable as arg: object
  1815. /
  1816. Alternative chain() constructor taking a single iterable argument that evaluates lazily.
  1817. [clinic start generated code]*/
  1818. static PyObject *
  1819. itertools_chain_from_iterable(PyTypeObject *type, PyObject *arg)
  1820. /*[clinic end generated code: output=667ae7a7f7b68654 input=72c39e3a2ca3be85]*/
  1821. {
  1822. PyObject *source;
  1823. source = PyObject_GetIter(arg);
  1824. if (source == NULL)
  1825. return NULL;
  1826. return chain_new_internal(type, source);
  1827. }
  1828. static void
  1829. chain_dealloc(chainobject *lz)
  1830. {
  1831. PyTypeObject *tp = Py_TYPE(lz);
  1832. PyObject_GC_UnTrack(lz);
  1833. Py_XDECREF(lz->active);
  1834. Py_XDECREF(lz->source);
  1835. tp->tp_free(lz);
  1836. Py_DECREF(tp);
  1837. }
  1838. static int
  1839. chain_traverse(chainobject *lz, visitproc visit, void *arg)
  1840. {
  1841. Py_VISIT(Py_TYPE(lz));
  1842. Py_VISIT(lz->source);
  1843. Py_VISIT(lz->active);
  1844. return 0;
  1845. }
  1846. static PyObject *
  1847. chain_next(chainobject *lz)
  1848. {
  1849. PyObject *item;
  1850. /* lz->source is the iterator of iterables. If it's NULL, we've already
  1851. * consumed them all. lz->active is the current iterator. If it's NULL,
  1852. * we should grab a new one from lz->source. */
  1853. while (lz->source != NULL) {
  1854. if (lz->active == NULL) {
  1855. PyObject *iterable = PyIter_Next(lz->source);
  1856. if (iterable == NULL) {
  1857. Py_CLEAR(lz->source);
  1858. return NULL; /* no more input sources */
  1859. }
  1860. lz->active = PyObject_GetIter(iterable);
  1861. Py_DECREF(iterable);
  1862. if (lz->active == NULL) {
  1863. Py_CLEAR(lz->source);
  1864. return NULL; /* input not iterable */
  1865. }
  1866. }
  1867. item = (*Py_TYPE(lz->active)->tp_iternext)(lz->active);
  1868. if (item != NULL)
  1869. return item;
  1870. if (PyErr_Occurred()) {
  1871. if (PyErr_ExceptionMatches(PyExc_StopIteration))
  1872. PyErr_Clear();
  1873. else
  1874. return NULL; /* input raised an exception */
  1875. }
  1876. /* lz->active is consumed, try with the next iterable. */
  1877. Py_CLEAR(lz->active);
  1878. }
  1879. /* Everything had been consumed already. */
  1880. return NULL;
  1881. }
  1882. static PyObject *
  1883. chain_reduce(chainobject *lz, PyObject *Py_UNUSED(ignored))
  1884. {
  1885. ITERTOOL_PICKLE_DEPRECATION;
  1886. if (lz->source) {
  1887. /* we can't pickle function objects (itertools.from_iterable) so
  1888. * we must use setstate to replace the iterable. One day we
  1889. * will fix pickling of functions
  1890. */
  1891. if (lz->active) {
  1892. return Py_BuildValue("O()(OO)", Py_TYPE(lz), lz->source, lz->active);
  1893. } else {
  1894. return Py_BuildValue("O()(O)", Py_TYPE(lz), lz->source);
  1895. }
  1896. } else {
  1897. return Py_BuildValue("O()", Py_TYPE(lz)); /* exhausted */
  1898. }
  1899. return NULL;
  1900. }
  1901. static PyObject *
  1902. chain_setstate(chainobject *lz, PyObject *state)
  1903. {
  1904. ITERTOOL_PICKLE_DEPRECATION;
  1905. PyObject *source, *active=NULL;
  1906. if (!PyTuple_Check(state)) {
  1907. PyErr_SetString(PyExc_TypeError, "state is not a tuple");
  1908. return NULL;
  1909. }
  1910. if (!PyArg_ParseTuple(state, "O|O", &source, &active)) {
  1911. return NULL;
  1912. }
  1913. if (!PyIter_Check(source) || (active != NULL && !PyIter_Check(active))) {
  1914. PyErr_SetString(PyExc_TypeError, "Arguments must be iterators.");
  1915. return NULL;
  1916. }
  1917. Py_INCREF(source);
  1918. Py_XSETREF(lz->source, source);
  1919. Py_XINCREF(active);
  1920. Py_XSETREF(lz->active, active);
  1921. Py_RETURN_NONE;
  1922. }
  1923. PyDoc_STRVAR(chain_doc,
  1924. "chain(*iterables) --> chain object\n\
  1925. \n\
  1926. Return a chain object whose .__next__() method returns elements from the\n\
  1927. first iterable until it is exhausted, then elements from the next\n\
  1928. iterable, until all of the iterables are exhausted.");
  1929. static PyMethodDef chain_methods[] = {
  1930. ITERTOOLS_CHAIN_FROM_ITERABLE_METHODDEF
  1931. {"__reduce__", (PyCFunction)chain_reduce, METH_NOARGS,
  1932. reduce_doc},
  1933. {"__setstate__", (PyCFunction)chain_setstate, METH_O,
  1934. setstate_doc},
  1935. {"__class_getitem__", Py_GenericAlias,
  1936. METH_O|METH_CLASS, PyDoc_STR("See PEP 585")},
  1937. {NULL, NULL} /* sentinel */
  1938. };
  1939. static PyType_Slot chain_slots[] = {
  1940. {Py_tp_dealloc, chain_dealloc},
  1941. {Py_tp_getattro, PyObject_GenericGetAttr},
  1942. {Py_tp_doc, (void *)chain_doc},
  1943. {Py_tp_traverse, chain_traverse},
  1944. {Py_tp_iter, PyObject_SelfIter},
  1945. {Py_tp_iternext, chain_next},
  1946. {Py_tp_methods, chain_methods},
  1947. {Py_tp_new, chain_new},
  1948. {Py_tp_free, PyObject_GC_Del},
  1949. {0, NULL},
  1950. };
  1951. static PyType_Spec chain_spec = {
  1952. .name = "itertools.chain",
  1953. .basicsize = sizeof(chainobject),
  1954. .flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_BASETYPE |
  1955. Py_TPFLAGS_IMMUTABLETYPE),
  1956. .slots = chain_slots,
  1957. };
  1958. /* product object ************************************************************/
  1959. typedef struct {
  1960. PyObject_HEAD
  1961. PyObject *pools; /* tuple of pool tuples */
  1962. Py_ssize_t *indices; /* one index per pool */
  1963. PyObject *result; /* most recently returned result tuple */
  1964. int stopped; /* set to 1 when the iterator is exhausted */
  1965. } productobject;
  1966. static PyObject *
  1967. product_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
  1968. {
  1969. productobject *lz;
  1970. Py_ssize_t nargs, npools, repeat=1;
  1971. PyObject *pools = NULL;
  1972. Py_ssize_t *indices = NULL;
  1973. Py_ssize_t i;
  1974. if (kwds != NULL) {
  1975. char *kwlist[] = {"repeat", 0};
  1976. PyObject *tmpargs = PyTuple_New(0);
  1977. if (tmpargs == NULL)
  1978. return NULL;
  1979. if (!PyArg_ParseTupleAndKeywords(tmpargs, kwds, "|n:product",
  1980. kwlist, &repeat)) {
  1981. Py_DECREF(tmpargs);
  1982. return NULL;
  1983. }
  1984. Py_DECREF(tmpargs);
  1985. if (repeat < 0) {
  1986. PyErr_SetString(PyExc_ValueError,
  1987. "repeat argument cannot be negative");
  1988. return NULL;
  1989. }
  1990. }
  1991. assert(PyTuple_CheckExact(args));
  1992. if (repeat == 0) {
  1993. nargs = 0;
  1994. } else {
  1995. nargs = PyTuple_GET_SIZE(args);
  1996. if ((size_t)nargs > PY_SSIZE_T_MAX/sizeof(Py_ssize_t)/repeat) {
  1997. PyErr_SetString(PyExc_OverflowError, "repeat argument too large");
  1998. return NULL;
  1999. }
  2000. }
  2001. npools = nargs * repeat;
  2002. indices = PyMem_New(Py_ssize_t, npools);
  2003. if (indices == NULL) {
  2004. PyErr_NoMemory();
  2005. goto error;
  2006. }
  2007. pools = PyTuple_New(npools);
  2008. if (pools == NULL)
  2009. goto error;
  2010. for (i=0; i < nargs ; ++i) {
  2011. PyObject *item = PyTuple_GET_ITEM(args, i);
  2012. PyObject *pool = PySequence_Tuple(item);
  2013. if (pool == NULL)
  2014. goto error;
  2015. PyTuple_SET_ITEM(pools, i, pool);
  2016. indices[i] = 0;
  2017. }
  2018. for ( ; i < npools; ++i) {
  2019. PyObject *pool = PyTuple_GET_ITEM(pools, i - nargs);
  2020. Py_INCREF(pool);
  2021. PyTuple_SET_ITEM(pools, i, pool);
  2022. indices[i] = 0;
  2023. }
  2024. /* create productobject structure */
  2025. lz = (productobject *)type->tp_alloc(type, 0);
  2026. if (lz == NULL)
  2027. goto error;
  2028. lz->pools = pools;
  2029. lz->indices = indices;
  2030. lz->result = NULL;
  2031. lz->stopped = 0;
  2032. return (PyObject *)lz;
  2033. error:
  2034. if (indices != NULL)
  2035. PyMem_Free(indices);
  2036. Py_XDECREF(pools);
  2037. return NULL;
  2038. }
  2039. static void
  2040. product_dealloc(productobject *lz)
  2041. {
  2042. PyTypeObject *tp = Py_TYPE(lz);
  2043. PyObject_GC_UnTrack(lz);
  2044. Py_XDECREF(lz->pools);
  2045. Py_XDECREF(lz->result);
  2046. if (lz->indices != NULL)
  2047. PyMem_Free(lz->indices);
  2048. tp->tp_free(lz);
  2049. Py_DECREF(tp);
  2050. }
  2051. static PyObject *
  2052. product_sizeof(productobject *lz, void *unused)
  2053. {
  2054. size_t res = _PyObject_SIZE(Py_TYPE(lz));
  2055. res += (size_t)PyTuple_GET_SIZE(lz->pools) * sizeof(Py_ssize_t);
  2056. return PyLong_FromSize_t(res);
  2057. }
  2058. PyDoc_STRVAR(sizeof_doc, "Returns size in memory, in bytes.");
  2059. static int
  2060. product_traverse(productobject *lz, visitproc visit, void *arg)
  2061. {
  2062. Py_VISIT(Py_TYPE(lz));
  2063. Py_VISIT(lz->pools);
  2064. Py_VISIT(lz->result);
  2065. return 0;
  2066. }
  2067. static PyObject *
  2068. product_next(productobject *lz)
  2069. {
  2070. PyObject *pool;
  2071. PyObject *elem;
  2072. PyObject *oldelem;
  2073. PyObject *pools = lz->pools;
  2074. PyObject *result = lz->result;
  2075. Py_ssize_t npools = PyTuple_GET_SIZE(pools);
  2076. Py_ssize_t i;
  2077. if (lz->stopped)
  2078. return NULL;
  2079. if (result == NULL) {
  2080. /* On the first pass, return an initial tuple filled with the
  2081. first element from each pool. */
  2082. result = PyTuple_New(npools);
  2083. if (result == NULL)
  2084. goto empty;
  2085. lz->result = result;
  2086. for (i=0; i < npools; i++) {
  2087. pool = PyTuple_GET_ITEM(pools, i);
  2088. if (PyTuple_GET_SIZE(pool) == 0)
  2089. goto empty;
  2090. elem = PyTuple_GET_ITEM(pool, 0);
  2091. Py_INCREF(elem);
  2092. PyTuple_SET_ITEM(result, i, elem);
  2093. }
  2094. } else {
  2095. Py_ssize_t *indices = lz->indices;
  2096. /* Copy the previous result tuple or re-use it if available */
  2097. if (Py_REFCNT(result) > 1) {
  2098. PyObject *old_result = result;
  2099. result = _PyTuple_FromArray(_PyTuple_ITEMS(old_result), npools);
  2100. if (result == NULL)
  2101. goto empty;
  2102. lz->result = result;
  2103. Py_DECREF(old_result);
  2104. }
  2105. // bpo-42536: The GC may have untracked this result tuple. Since we're
  2106. // recycling it, make sure it's tracked again:
  2107. else if (!_PyObject_GC_IS_TRACKED(result)) {
  2108. _PyObject_GC_TRACK(result);
  2109. }
  2110. /* Now, we've got the only copy so we can update it in-place */
  2111. assert (npools==0 || Py_REFCNT(result) == 1);
  2112. /* Update the pool indices right-to-left. Only advance to the
  2113. next pool when the previous one rolls-over */
  2114. for (i=npools-1 ; i >= 0 ; i--) {
  2115. pool = PyTuple_GET_ITEM(pools, i);
  2116. indices[i]++;
  2117. if (indices[i] == PyTuple_GET_SIZE(pool)) {
  2118. /* Roll-over and advance to next pool */
  2119. indices[i] = 0;
  2120. elem = PyTuple_GET_ITEM(pool, 0);
  2121. Py_INCREF(elem);
  2122. oldelem = PyTuple_GET_ITEM(result, i);
  2123. PyTuple_SET_ITEM(result, i, elem);
  2124. Py_DECREF(oldelem);
  2125. } else {
  2126. /* No rollover. Just increment and stop here. */
  2127. elem = PyTuple_GET_ITEM(pool, indices[i]);
  2128. Py_INCREF(elem);
  2129. oldelem = PyTuple_GET_ITEM(result, i);
  2130. PyTuple_SET_ITEM(result, i, elem);
  2131. Py_DECREF(oldelem);
  2132. break;
  2133. }
  2134. }
  2135. /* If i is negative, then the indices have all rolled-over
  2136. and we're done. */
  2137. if (i < 0)
  2138. goto empty;
  2139. }
  2140. return Py_NewRef(result);
  2141. empty:
  2142. lz->stopped = 1;
  2143. return NULL;
  2144. }
  2145. static PyObject *
  2146. product_reduce(productobject *lz, PyObject *Py_UNUSED(ignored))
  2147. {
  2148. ITERTOOL_PICKLE_DEPRECATION;
  2149. if (lz->stopped) {
  2150. return Py_BuildValue("O(())", Py_TYPE(lz));
  2151. } else if (lz->result == NULL) {
  2152. return Py_BuildValue("OO", Py_TYPE(lz), lz->pools);
  2153. } else {
  2154. PyObject *indices;
  2155. Py_ssize_t n, i;
  2156. /* we must pickle the indices use them for setstate, and
  2157. * additionally indicate that the iterator has started
  2158. */
  2159. n = PyTuple_GET_SIZE(lz->pools);
  2160. indices = PyTuple_New(n);
  2161. if (indices == NULL)
  2162. return NULL;
  2163. for (i=0; i<n; i++){
  2164. PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
  2165. if (!index) {
  2166. Py_DECREF(indices);
  2167. return NULL;
  2168. }
  2169. PyTuple_SET_ITEM(indices, i, index);
  2170. }
  2171. return Py_BuildValue("OON", Py_TYPE(lz), lz->pools, indices);
  2172. }
  2173. }
  2174. static PyObject *
  2175. product_setstate(productobject *lz, PyObject *state)
  2176. {
  2177. ITERTOOL_PICKLE_DEPRECATION;
  2178. PyObject *result;
  2179. Py_ssize_t n, i;
  2180. n = PyTuple_GET_SIZE(lz->pools);
  2181. if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != n) {
  2182. PyErr_SetString(PyExc_ValueError, "invalid arguments");
  2183. return NULL;
  2184. }
  2185. for (i=0; i<n; i++)
  2186. {
  2187. PyObject* indexObject = PyTuple_GET_ITEM(state, i);
  2188. Py_ssize_t index = PyLong_AsSsize_t(indexObject);
  2189. PyObject* pool;
  2190. Py_ssize_t poolsize;
  2191. if (index < 0 && PyErr_Occurred())
  2192. return NULL; /* not an integer */
  2193. pool = PyTuple_GET_ITEM(lz->pools, i);
  2194. poolsize = PyTuple_GET_SIZE(pool);
  2195. if (poolsize == 0) {
  2196. lz->stopped = 1;
  2197. Py_RETURN_NONE;
  2198. }
  2199. /* clamp the index */
  2200. if (index < 0)
  2201. index = 0;
  2202. else if (index > poolsize-1)
  2203. index = poolsize-1;
  2204. lz->indices[i] = index;
  2205. }
  2206. result = PyTuple_New(n);
  2207. if (!result)
  2208. return NULL;
  2209. for (i=0; i<n; i++) {
  2210. PyObject *pool = PyTuple_GET_ITEM(lz->pools, i);
  2211. PyObject *element = PyTuple_GET_ITEM(pool, lz->indices[i]);
  2212. Py_INCREF(element);
  2213. PyTuple_SET_ITEM(result, i, element);
  2214. }
  2215. Py_XSETREF(lz->result, result);
  2216. Py_RETURN_NONE;
  2217. }
  2218. static PyMethodDef product_methods[] = {
  2219. {"__reduce__", (PyCFunction)product_reduce, METH_NOARGS,
  2220. reduce_doc},
  2221. {"__setstate__", (PyCFunction)product_setstate, METH_O,
  2222. setstate_doc},
  2223. {"__sizeof__", (PyCFunction)product_sizeof, METH_NOARGS,
  2224. sizeof_doc},
  2225. {NULL, NULL} /* sentinel */
  2226. };
  2227. PyDoc_STRVAR(product_doc,
  2228. "product(*iterables, repeat=1) --> product object\n\
  2229. \n\
  2230. Cartesian product of input iterables. Equivalent to nested for-loops.\n\n\
  2231. For example, product(A, B) returns the same as: ((x,y) for x in A for y in B).\n\
  2232. The leftmost iterators are in the outermost for-loop, so the output tuples\n\
  2233. cycle in a manner similar to an odometer (with the rightmost element changing\n\
  2234. on every iteration).\n\n\
  2235. To compute the product of an iterable with itself, specify the number\n\
  2236. of repetitions with the optional repeat keyword argument. For example,\n\
  2237. product(A, repeat=4) means the same as product(A, A, A, A).\n\n\
  2238. product('ab', range(3)) --> ('a',0) ('a',1) ('a',2) ('b',0) ('b',1) ('b',2)\n\
  2239. product((0,1), (0,1), (0,1)) --> (0,0,0) (0,0,1) (0,1,0) (0,1,1) (1,0,0) ...");
  2240. static PyType_Slot product_slots[] = {
  2241. {Py_tp_dealloc, product_dealloc},
  2242. {Py_tp_getattro, PyObject_GenericGetAttr},
  2243. {Py_tp_doc, (void *)product_doc},
  2244. {Py_tp_traverse, product_traverse},
  2245. {Py_tp_iter, PyObject_SelfIter},
  2246. {Py_tp_iternext, product_next},
  2247. {Py_tp_methods, product_methods},
  2248. {Py_tp_new, product_new},
  2249. {Py_tp_free, PyObject_GC_Del},
  2250. {0, NULL},
  2251. };
  2252. static PyType_Spec product_spec = {
  2253. .name = "itertools.product",
  2254. .basicsize = sizeof(productobject),
  2255. .flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_BASETYPE |
  2256. Py_TPFLAGS_IMMUTABLETYPE),
  2257. .slots = product_slots,
  2258. };
  2259. /* combinations object *******************************************************/
  2260. typedef struct {
  2261. PyObject_HEAD
  2262. PyObject *pool; /* input converted to a tuple */
  2263. Py_ssize_t *indices; /* one index per result element */
  2264. PyObject *result; /* most recently returned result tuple */
  2265. Py_ssize_t r; /* size of result tuple */
  2266. int stopped; /* set to 1 when the iterator is exhausted */
  2267. } combinationsobject;
  2268. /*[clinic input]
  2269. @classmethod
  2270. itertools.combinations.__new__
  2271. iterable: object
  2272. r: Py_ssize_t
  2273. Return successive r-length combinations of elements in the iterable.
  2274. combinations(range(4), 3) --> (0,1,2), (0,1,3), (0,2,3), (1,2,3)
  2275. [clinic start generated code]*/
  2276. static PyObject *
  2277. itertools_combinations_impl(PyTypeObject *type, PyObject *iterable,
  2278. Py_ssize_t r)
  2279. /*[clinic end generated code: output=87a689b39c40039c input=06bede09e3da20f8]*/
  2280. {
  2281. combinationsobject *co;
  2282. Py_ssize_t n;
  2283. PyObject *pool = NULL;
  2284. Py_ssize_t *indices = NULL;
  2285. Py_ssize_t i;
  2286. pool = PySequence_Tuple(iterable);
  2287. if (pool == NULL)
  2288. goto error;
  2289. n = PyTuple_GET_SIZE(pool);
  2290. if (r < 0) {
  2291. PyErr_SetString(PyExc_ValueError, "r must be non-negative");
  2292. goto error;
  2293. }
  2294. indices = PyMem_New(Py_ssize_t, r);
  2295. if (indices == NULL) {
  2296. PyErr_NoMemory();
  2297. goto error;
  2298. }
  2299. for (i=0 ; i<r ; i++)
  2300. indices[i] = i;
  2301. /* create combinationsobject structure */
  2302. co = (combinationsobject *)type->tp_alloc(type, 0);
  2303. if (co == NULL)
  2304. goto error;
  2305. co->pool = pool;
  2306. co->indices = indices;
  2307. co->result = NULL;
  2308. co->r = r;
  2309. co->stopped = r > n ? 1 : 0;
  2310. return (PyObject *)co;
  2311. error:
  2312. if (indices != NULL)
  2313. PyMem_Free(indices);
  2314. Py_XDECREF(pool);
  2315. return NULL;
  2316. }
  2317. static void
  2318. combinations_dealloc(combinationsobject *co)
  2319. {
  2320. PyTypeObject *tp = Py_TYPE(co);
  2321. PyObject_GC_UnTrack(co);
  2322. Py_XDECREF(co->pool);
  2323. Py_XDECREF(co->result);
  2324. if (co->indices != NULL)
  2325. PyMem_Free(co->indices);
  2326. tp->tp_free(co);
  2327. Py_DECREF(tp);
  2328. }
  2329. static PyObject *
  2330. combinations_sizeof(combinationsobject *co, void *unused)
  2331. {
  2332. size_t res = _PyObject_SIZE(Py_TYPE(co));
  2333. res += (size_t)co->r * sizeof(Py_ssize_t);
  2334. return PyLong_FromSize_t(res);
  2335. }
  2336. static int
  2337. combinations_traverse(combinationsobject *co, visitproc visit, void *arg)
  2338. {
  2339. Py_VISIT(Py_TYPE(co));
  2340. Py_VISIT(co->pool);
  2341. Py_VISIT(co->result);
  2342. return 0;
  2343. }
  2344. static PyObject *
  2345. combinations_next(combinationsobject *co)
  2346. {
  2347. PyObject *elem;
  2348. PyObject *oldelem;
  2349. PyObject *pool = co->pool;
  2350. Py_ssize_t *indices = co->indices;
  2351. PyObject *result = co->result;
  2352. Py_ssize_t n = PyTuple_GET_SIZE(pool);
  2353. Py_ssize_t r = co->r;
  2354. Py_ssize_t i, j, index;
  2355. if (co->stopped)
  2356. return NULL;
  2357. if (result == NULL) {
  2358. /* On the first pass, initialize result tuple using the indices */
  2359. result = PyTuple_New(r);
  2360. if (result == NULL)
  2361. goto empty;
  2362. co->result = result;
  2363. for (i=0; i<r ; i++) {
  2364. index = indices[i];
  2365. elem = PyTuple_GET_ITEM(pool, index);
  2366. Py_INCREF(elem);
  2367. PyTuple_SET_ITEM(result, i, elem);
  2368. }
  2369. } else {
  2370. /* Copy the previous result tuple or re-use it if available */
  2371. if (Py_REFCNT(result) > 1) {
  2372. PyObject *old_result = result;
  2373. result = _PyTuple_FromArray(_PyTuple_ITEMS(old_result), r);
  2374. if (result == NULL)
  2375. goto empty;
  2376. co->result = result;
  2377. Py_DECREF(old_result);
  2378. }
  2379. // bpo-42536: The GC may have untracked this result tuple. Since we're
  2380. // recycling it, make sure it's tracked again:
  2381. else if (!_PyObject_GC_IS_TRACKED(result)) {
  2382. _PyObject_GC_TRACK(result);
  2383. }
  2384. /* Now, we've got the only copy so we can update it in-place
  2385. * CPython's empty tuple is a singleton and cached in
  2386. * PyTuple's freelist.
  2387. */
  2388. assert(r == 0 || Py_REFCNT(result) == 1);
  2389. /* Scan indices right-to-left until finding one that is not
  2390. at its maximum (i + n - r). */
  2391. for (i=r-1 ; i >= 0 && indices[i] == i+n-r ; i--)
  2392. ;
  2393. /* If i is negative, then the indices are all at
  2394. their maximum value and we're done. */
  2395. if (i < 0)
  2396. goto empty;
  2397. /* Increment the current index which we know is not at its
  2398. maximum. Then move back to the right setting each index
  2399. to its lowest possible value (one higher than the index
  2400. to its left -- this maintains the sort order invariant). */
  2401. indices[i]++;
  2402. for (j=i+1 ; j<r ; j++)
  2403. indices[j] = indices[j-1] + 1;
  2404. /* Update the result tuple for the new indices
  2405. starting with i, the leftmost index that changed */
  2406. for ( ; i<r ; i++) {
  2407. index = indices[i];
  2408. elem = PyTuple_GET_ITEM(pool, index);
  2409. Py_INCREF(elem);
  2410. oldelem = PyTuple_GET_ITEM(result, i);
  2411. PyTuple_SET_ITEM(result, i, elem);
  2412. Py_DECREF(oldelem);
  2413. }
  2414. }
  2415. return Py_NewRef(result);
  2416. empty:
  2417. co->stopped = 1;
  2418. return NULL;
  2419. }
  2420. static PyObject *
  2421. combinations_reduce(combinationsobject *lz, PyObject *Py_UNUSED(ignored))
  2422. {
  2423. ITERTOOL_PICKLE_DEPRECATION;
  2424. if (lz->result == NULL) {
  2425. return Py_BuildValue("O(On)", Py_TYPE(lz), lz->pool, lz->r);
  2426. } else if (lz->stopped) {
  2427. return Py_BuildValue("O(()n)", Py_TYPE(lz), lz->r);
  2428. } else {
  2429. PyObject *indices;
  2430. Py_ssize_t i;
  2431. /* we must pickle the indices and use them for setstate */
  2432. indices = PyTuple_New(lz->r);
  2433. if (!indices)
  2434. return NULL;
  2435. for (i=0; i<lz->r; i++)
  2436. {
  2437. PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
  2438. if (!index) {
  2439. Py_DECREF(indices);
  2440. return NULL;
  2441. }
  2442. PyTuple_SET_ITEM(indices, i, index);
  2443. }
  2444. return Py_BuildValue("O(On)N", Py_TYPE(lz), lz->pool, lz->r, indices);
  2445. }
  2446. }
  2447. static PyObject *
  2448. combinations_setstate(combinationsobject *lz, PyObject *state)
  2449. {
  2450. ITERTOOL_PICKLE_DEPRECATION;
  2451. PyObject *result;
  2452. Py_ssize_t i;
  2453. Py_ssize_t n = PyTuple_GET_SIZE(lz->pool);
  2454. if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != lz->r) {
  2455. PyErr_SetString(PyExc_ValueError, "invalid arguments");
  2456. return NULL;
  2457. }
  2458. for (i=0; i<lz->r; i++) {
  2459. Py_ssize_t max;
  2460. PyObject* indexObject = PyTuple_GET_ITEM(state, i);
  2461. Py_ssize_t index = PyLong_AsSsize_t(indexObject);
  2462. if (index == -1 && PyErr_Occurred())
  2463. return NULL; /* not an integer */
  2464. max = i + n - lz->r;
  2465. /* clamp the index (beware of negative max) */
  2466. if (index > max)
  2467. index = max;
  2468. if (index < 0)
  2469. index = 0;
  2470. lz->indices[i] = index;
  2471. }
  2472. result = PyTuple_New(lz->r);
  2473. if (result == NULL)
  2474. return NULL;
  2475. for (i=0; i<lz->r; i++) {
  2476. PyObject *element = PyTuple_GET_ITEM(lz->pool, lz->indices[i]);
  2477. Py_INCREF(element);
  2478. PyTuple_SET_ITEM(result, i, element);
  2479. }
  2480. Py_XSETREF(lz->result, result);
  2481. Py_RETURN_NONE;
  2482. }
  2483. static PyMethodDef combinations_methods[] = {
  2484. {"__reduce__", (PyCFunction)combinations_reduce, METH_NOARGS,
  2485. reduce_doc},
  2486. {"__setstate__", (PyCFunction)combinations_setstate, METH_O,
  2487. setstate_doc},
  2488. {"__sizeof__", (PyCFunction)combinations_sizeof, METH_NOARGS,
  2489. sizeof_doc},
  2490. {NULL, NULL} /* sentinel */
  2491. };
  2492. static PyType_Slot combinations_slots[] = {
  2493. {Py_tp_dealloc, combinations_dealloc},
  2494. {Py_tp_getattro, PyObject_GenericGetAttr},
  2495. {Py_tp_doc, (void *)itertools_combinations__doc__},
  2496. {Py_tp_traverse, combinations_traverse},
  2497. {Py_tp_iter, PyObject_SelfIter},
  2498. {Py_tp_iternext, combinations_next},
  2499. {Py_tp_methods, combinations_methods},
  2500. {Py_tp_new, itertools_combinations},
  2501. {Py_tp_free, PyObject_GC_Del},
  2502. {0, NULL},
  2503. };
  2504. static PyType_Spec combinations_spec = {
  2505. .name = "itertools.combinations",
  2506. .basicsize = sizeof(combinationsobject),
  2507. .flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_BASETYPE |
  2508. Py_TPFLAGS_IMMUTABLETYPE),
  2509. .slots = combinations_slots,
  2510. };
  2511. /* combinations with replacement object **************************************/
  2512. /* Equivalent to:
  2513. def combinations_with_replacement(iterable, r):
  2514. "combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC"
  2515. # number items returned: (n+r-1)! / r! / (n-1)!
  2516. pool = tuple(iterable)
  2517. n = len(pool)
  2518. indices = [0] * r
  2519. yield tuple(pool[i] for i in indices)
  2520. while 1:
  2521. for i in reversed(range(r)):
  2522. if indices[i] != n - 1:
  2523. break
  2524. else:
  2525. return
  2526. indices[i:] = [indices[i] + 1] * (r - i)
  2527. yield tuple(pool[i] for i in indices)
  2528. def combinations_with_replacement2(iterable, r):
  2529. 'Alternate version that filters from product()'
  2530. pool = tuple(iterable)
  2531. n = len(pool)
  2532. for indices in product(range(n), repeat=r):
  2533. if sorted(indices) == list(indices):
  2534. yield tuple(pool[i] for i in indices)
  2535. */
  2536. typedef struct {
  2537. PyObject_HEAD
  2538. PyObject *pool; /* input converted to a tuple */
  2539. Py_ssize_t *indices; /* one index per result element */
  2540. PyObject *result; /* most recently returned result tuple */
  2541. Py_ssize_t r; /* size of result tuple */
  2542. int stopped; /* set to 1 when the cwr iterator is exhausted */
  2543. } cwrobject;
  2544. /*[clinic input]
  2545. @classmethod
  2546. itertools.combinations_with_replacement.__new__
  2547. iterable: object
  2548. r: Py_ssize_t
  2549. Return successive r-length combinations of elements in the iterable allowing individual elements to have successive repeats.
  2550. combinations_with_replacement('ABC', 2) --> ('A','A'), ('A','B'), ('A','C'), ('B','B'), ('B','C'), ('C','C')
  2551. [clinic start generated code]*/
  2552. static PyObject *
  2553. itertools_combinations_with_replacement_impl(PyTypeObject *type,
  2554. PyObject *iterable,
  2555. Py_ssize_t r)
  2556. /*[clinic end generated code: output=48b26856d4e659ca input=1dc58e82a0878fdc]*/
  2557. {
  2558. cwrobject *co;
  2559. Py_ssize_t n;
  2560. PyObject *pool = NULL;
  2561. Py_ssize_t *indices = NULL;
  2562. Py_ssize_t i;
  2563. pool = PySequence_Tuple(iterable);
  2564. if (pool == NULL)
  2565. goto error;
  2566. n = PyTuple_GET_SIZE(pool);
  2567. if (r < 0) {
  2568. PyErr_SetString(PyExc_ValueError, "r must be non-negative");
  2569. goto error;
  2570. }
  2571. indices = PyMem_New(Py_ssize_t, r);
  2572. if (indices == NULL) {
  2573. PyErr_NoMemory();
  2574. goto error;
  2575. }
  2576. for (i=0 ; i<r ; i++)
  2577. indices[i] = 0;
  2578. /* create cwrobject structure */
  2579. co = (cwrobject *)type->tp_alloc(type, 0);
  2580. if (co == NULL)
  2581. goto error;
  2582. co->pool = pool;
  2583. co->indices = indices;
  2584. co->result = NULL;
  2585. co->r = r;
  2586. co->stopped = !n && r;
  2587. return (PyObject *)co;
  2588. error:
  2589. if (indices != NULL)
  2590. PyMem_Free(indices);
  2591. Py_XDECREF(pool);
  2592. return NULL;
  2593. }
  2594. static void
  2595. cwr_dealloc(cwrobject *co)
  2596. {
  2597. PyTypeObject *tp = Py_TYPE(co);
  2598. PyObject_GC_UnTrack(co);
  2599. Py_XDECREF(co->pool);
  2600. Py_XDECREF(co->result);
  2601. if (co->indices != NULL)
  2602. PyMem_Free(co->indices);
  2603. tp->tp_free(co);
  2604. Py_DECREF(tp);
  2605. }
  2606. static PyObject *
  2607. cwr_sizeof(cwrobject *co, void *unused)
  2608. {
  2609. size_t res = _PyObject_SIZE(Py_TYPE(co));
  2610. res += (size_t)co->r * sizeof(Py_ssize_t);
  2611. return PyLong_FromSize_t(res);
  2612. }
  2613. static int
  2614. cwr_traverse(cwrobject *co, visitproc visit, void *arg)
  2615. {
  2616. Py_VISIT(Py_TYPE(co));
  2617. Py_VISIT(co->pool);
  2618. Py_VISIT(co->result);
  2619. return 0;
  2620. }
  2621. static PyObject *
  2622. cwr_next(cwrobject *co)
  2623. {
  2624. PyObject *elem;
  2625. PyObject *oldelem;
  2626. PyObject *pool = co->pool;
  2627. Py_ssize_t *indices = co->indices;
  2628. PyObject *result = co->result;
  2629. Py_ssize_t n = PyTuple_GET_SIZE(pool);
  2630. Py_ssize_t r = co->r;
  2631. Py_ssize_t i, index;
  2632. if (co->stopped)
  2633. return NULL;
  2634. if (result == NULL) {
  2635. /* On the first pass, initialize result tuple with pool[0] */
  2636. result = PyTuple_New(r);
  2637. if (result == NULL)
  2638. goto empty;
  2639. co->result = result;
  2640. if (n > 0) {
  2641. elem = PyTuple_GET_ITEM(pool, 0);
  2642. for (i=0; i<r ; i++) {
  2643. assert(indices[i] == 0);
  2644. Py_INCREF(elem);
  2645. PyTuple_SET_ITEM(result, i, elem);
  2646. }
  2647. }
  2648. } else {
  2649. /* Copy the previous result tuple or re-use it if available */
  2650. if (Py_REFCNT(result) > 1) {
  2651. PyObject *old_result = result;
  2652. result = _PyTuple_FromArray(_PyTuple_ITEMS(old_result), r);
  2653. if (result == NULL)
  2654. goto empty;
  2655. co->result = result;
  2656. Py_DECREF(old_result);
  2657. }
  2658. // bpo-42536: The GC may have untracked this result tuple. Since we're
  2659. // recycling it, make sure it's tracked again:
  2660. else if (!_PyObject_GC_IS_TRACKED(result)) {
  2661. _PyObject_GC_TRACK(result);
  2662. }
  2663. /* Now, we've got the only copy so we can update it in-place CPython's
  2664. empty tuple is a singleton and cached in PyTuple's freelist. */
  2665. assert(r == 0 || Py_REFCNT(result) == 1);
  2666. /* Scan indices right-to-left until finding one that is not
  2667. * at its maximum (n-1). */
  2668. for (i=r-1 ; i >= 0 && indices[i] == n-1; i--)
  2669. ;
  2670. /* If i is negative, then the indices are all at
  2671. their maximum value and we're done. */
  2672. if (i < 0)
  2673. goto empty;
  2674. /* Increment the current index which we know is not at its
  2675. maximum. Then set all to the right to the same value. */
  2676. index = indices[i] + 1;
  2677. assert(index < n);
  2678. elem = PyTuple_GET_ITEM(pool, index);
  2679. for ( ; i<r ; i++) {
  2680. indices[i] = index;
  2681. Py_INCREF(elem);
  2682. oldelem = PyTuple_GET_ITEM(result, i);
  2683. PyTuple_SET_ITEM(result, i, elem);
  2684. Py_DECREF(oldelem);
  2685. }
  2686. }
  2687. return Py_NewRef(result);
  2688. empty:
  2689. co->stopped = 1;
  2690. return NULL;
  2691. }
  2692. static PyObject *
  2693. cwr_reduce(cwrobject *lz, PyObject *Py_UNUSED(ignored))
  2694. {
  2695. ITERTOOL_PICKLE_DEPRECATION;
  2696. if (lz->result == NULL) {
  2697. return Py_BuildValue("O(On)", Py_TYPE(lz), lz->pool, lz->r);
  2698. } else if (lz->stopped) {
  2699. return Py_BuildValue("O(()n)", Py_TYPE(lz), lz->r);
  2700. } else {
  2701. PyObject *indices;
  2702. Py_ssize_t i;
  2703. /* we must pickle the indices and use them for setstate */
  2704. indices = PyTuple_New(lz->r);
  2705. if (!indices)
  2706. return NULL;
  2707. for (i=0; i<lz->r; i++) {
  2708. PyObject* index = PyLong_FromSsize_t(lz->indices[i]);
  2709. if (!index) {
  2710. Py_DECREF(indices);
  2711. return NULL;
  2712. }
  2713. PyTuple_SET_ITEM(indices, i, index);
  2714. }
  2715. return Py_BuildValue("O(On)N", Py_TYPE(lz), lz->pool, lz->r, indices);
  2716. }
  2717. }
  2718. static PyObject *
  2719. cwr_setstate(cwrobject *lz, PyObject *state)
  2720. {
  2721. ITERTOOL_PICKLE_DEPRECATION;
  2722. PyObject *result;
  2723. Py_ssize_t n, i;
  2724. if (!PyTuple_Check(state) || PyTuple_GET_SIZE(state) != lz->r)
  2725. {
  2726. PyErr_SetString(PyExc_ValueError, "invalid arguments");
  2727. return NULL;
  2728. }
  2729. n = PyTuple_GET_SIZE(lz->pool);
  2730. for (i=0; i<lz->r; i++) {
  2731. PyObject* indexObject = PyTuple_GET_ITEM(state, i);
  2732. Py_ssize_t index = PyLong_AsSsize_t(indexObject);
  2733. if (index < 0 && PyErr_Occurred())
  2734. return NULL; /* not an integer */
  2735. /* clamp the index */
  2736. if (index < 0)
  2737. index = 0;
  2738. else if (index > n-1)
  2739. index = n-1;
  2740. lz->indices[i] = index;
  2741. }
  2742. result = PyTuple_New(lz->r);
  2743. if (result == NULL)
  2744. return NULL;
  2745. for (i=0; i<lz->r; i++) {
  2746. PyObject *element = PyTuple_GET_ITEM(lz->pool, lz->indices[i]);
  2747. Py_INCREF(element);
  2748. PyTuple_SET_ITEM(result, i, element);
  2749. }
  2750. Py_XSETREF(lz->result, result);
  2751. Py_RETURN_NONE;
  2752. }
  2753. static PyMethodDef cwr_methods[] = {
  2754. {"__reduce__", (PyCFunction)cwr_reduce, METH_NOARGS,
  2755. reduce_doc},
  2756. {"__setstate__", (PyCFunction)cwr_setstate, METH_O,
  2757. setstate_doc},
  2758. {"__sizeof__", (PyCFunction)cwr_sizeof, METH_NOARGS,
  2759. sizeof_doc},
  2760. {NULL, NULL} /* sentinel */
  2761. };
  2762. static PyType_Slot cwr_slots[] = {
  2763. {Py_tp_dealloc, cwr_dealloc},
  2764. {Py_tp_getattro, PyObject_GenericGetAttr},
  2765. {Py_tp_doc, (void *)itertools_combinations_with_replacement__doc__},
  2766. {Py_tp_traverse, cwr_traverse},
  2767. {Py_tp_iter, PyObject_SelfIter},
  2768. {Py_tp_iternext, cwr_next},
  2769. {Py_tp_methods, cwr_methods},
  2770. {Py_tp_new, itertools_combinations_with_replacement},
  2771. {Py_tp_free, PyObject_GC_Del},
  2772. {0, NULL},
  2773. };
  2774. static PyType_Spec cwr_spec = {
  2775. .name = "itertools.combinations_with_replacement",
  2776. .basicsize = sizeof(cwrobject),
  2777. .flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_BASETYPE |
  2778. Py_TPFLAGS_IMMUTABLETYPE),
  2779. .slots = cwr_slots,
  2780. };
  2781. /* permutations object ********************************************************
  2782. def permutations(iterable, r=None):
  2783. # permutations('ABCD', 2) --> AB AC AD BA BC BD CA CB CD DA DB DC
  2784. # permutations(range(3)) --> 012 021 102 120 201 210
  2785. pool = tuple(iterable)
  2786. n = len(pool)
  2787. r = n if r is None else r
  2788. if r > n:
  2789. return
  2790. indices = list(range(n))
  2791. cycles = list(range(n, n-r, -1))
  2792. yield tuple(pool[i] for i in indices[:r])
  2793. while n:
  2794. for i in reversed(range(r)):
  2795. cycles[i] -= 1
  2796. if cycles[i] == 0:
  2797. indices[i:] = indices[i+1:] + indices[i:i+1]
  2798. cycles[i] = n - i
  2799. else:
  2800. j = cycles[i]
  2801. indices[i], indices[-j] = indices[-j], indices[i]
  2802. yield tuple(pool[i] for i in indices[:r])
  2803. break
  2804. else:
  2805. return
  2806. */
  2807. typedef struct {
  2808. PyObject_HEAD
  2809. PyObject *pool; /* input converted to a tuple */
  2810. Py_ssize_t *indices; /* one index per element in the pool */
  2811. Py_ssize_t *cycles; /* one rollover counter per element in the result */
  2812. PyObject *result; /* most recently returned result tuple */
  2813. Py_ssize_t r; /* size of result tuple */
  2814. int stopped; /* set to 1 when the iterator is exhausted */
  2815. } permutationsobject;
  2816. /*[clinic input]
  2817. @classmethod
  2818. itertools.permutations.__new__
  2819. iterable: object
  2820. r as robj: object = None
  2821. Return successive r-length permutations of elements in the iterable.
  2822. permutations(range(3), 2) --> (0,1), (0,2), (1,0), (1,2), (2,0), (2,1)
  2823. [clinic start generated code]*/
  2824. static PyObject *
  2825. itertools_permutations_impl(PyTypeObject *type, PyObject *iterable,
  2826. PyObject *robj)
  2827. /*[clinic end generated code: output=296a72fa76d620ea input=57d0170a4ac0ec7a]*/
  2828. {
  2829. permutationsobject *po;
  2830. Py_ssize_t n;
  2831. Py_ssize_t r;
  2832. PyObject *pool = NULL;
  2833. Py_ssize_t *indices = NULL;
  2834. Py_ssize_t *cycles = NULL;
  2835. Py_ssize_t i;
  2836. pool = PySequence_Tuple(iterable);
  2837. if (pool == NULL)
  2838. goto error;
  2839. n = PyTuple_GET_SIZE(pool);
  2840. r = n;
  2841. if (robj != Py_None) {
  2842. if (!PyLong_Check(robj)) {
  2843. PyErr_SetString(PyExc_TypeError, "Expected int as r");
  2844. goto error;
  2845. }
  2846. r = PyLong_AsSsize_t(robj);
  2847. if (r == -1 && PyErr_Occurred())
  2848. goto error;
  2849. }
  2850. if (r < 0) {
  2851. PyErr_SetString(PyExc_ValueError, "r must be non-negative");
  2852. goto error;
  2853. }
  2854. indices = PyMem_New(Py_ssize_t, n);
  2855. cycles = PyMem_New(Py_ssize_t, r);
  2856. if (indices == NULL || cycles == NULL) {
  2857. PyErr_NoMemory();
  2858. goto error;
  2859. }
  2860. for (i=0 ; i<n ; i++)
  2861. indices[i] = i;
  2862. for (i=0 ; i<r ; i++)
  2863. cycles[i] = n - i;
  2864. /* create permutationsobject structure */
  2865. po = (permutationsobject *)type->tp_alloc(type, 0);
  2866. if (po == NULL)
  2867. goto error;
  2868. po->pool = pool;
  2869. po->indices = indices;
  2870. po->cycles = cycles;
  2871. po->result = NULL;
  2872. po->r = r;
  2873. po->stopped = r > n ? 1 : 0;
  2874. return (PyObject *)po;
  2875. error:
  2876. if (indices != NULL)
  2877. PyMem_Free(indices);
  2878. if (cycles != NULL)
  2879. PyMem_Free(cycles);
  2880. Py_XDECREF(pool);
  2881. return NULL;
  2882. }
  2883. static void
  2884. permutations_dealloc(permutationsobject *po)
  2885. {
  2886. PyTypeObject *tp = Py_TYPE(po);
  2887. PyObject_GC_UnTrack(po);
  2888. Py_XDECREF(po->pool);
  2889. Py_XDECREF(po->result);
  2890. PyMem_Free(po->indices);
  2891. PyMem_Free(po->cycles);
  2892. tp->tp_free(po);
  2893. Py_DECREF(tp);
  2894. }
  2895. static PyObject *
  2896. permutations_sizeof(permutationsobject *po, void *unused)
  2897. {
  2898. size_t res = _PyObject_SIZE(Py_TYPE(po));
  2899. res += (size_t)PyTuple_GET_SIZE(po->pool) * sizeof(Py_ssize_t);
  2900. res += (size_t)po->r * sizeof(Py_ssize_t);
  2901. return PyLong_FromSize_t(res);
  2902. }
  2903. static int
  2904. permutations_traverse(permutationsobject *po, visitproc visit, void *arg)
  2905. {
  2906. Py_VISIT(Py_TYPE(po));
  2907. Py_VISIT(po->pool);
  2908. Py_VISIT(po->result);
  2909. return 0;
  2910. }
  2911. static PyObject *
  2912. permutations_next(permutationsobject *po)
  2913. {
  2914. PyObject *elem;
  2915. PyObject *oldelem;
  2916. PyObject *pool = po->pool;
  2917. Py_ssize_t *indices = po->indices;
  2918. Py_ssize_t *cycles = po->cycles;
  2919. PyObject *result = po->result;
  2920. Py_ssize_t n = PyTuple_GET_SIZE(pool);
  2921. Py_ssize_t r = po->r;
  2922. Py_ssize_t i, j, k, index;
  2923. if (po->stopped)
  2924. return NULL;
  2925. if (result == NULL) {
  2926. /* On the first pass, initialize result tuple using the indices */
  2927. result = PyTuple_New(r);
  2928. if (result == NULL)
  2929. goto empty;
  2930. po->result = result;
  2931. for (i=0; i<r ; i++) {
  2932. index = indices[i];
  2933. elem = PyTuple_GET_ITEM(pool, index);
  2934. Py_INCREF(elem);
  2935. PyTuple_SET_ITEM(result, i, elem);
  2936. }
  2937. } else {
  2938. if (n == 0)
  2939. goto empty;
  2940. /* Copy the previous result tuple or re-use it if available */
  2941. if (Py_REFCNT(result) > 1) {
  2942. PyObject *old_result = result;
  2943. result = _PyTuple_FromArray(_PyTuple_ITEMS(old_result), r);
  2944. if (result == NULL)
  2945. goto empty;
  2946. po->result = result;
  2947. Py_DECREF(old_result);
  2948. }
  2949. // bpo-42536: The GC may have untracked this result tuple. Since we're
  2950. // recycling it, make sure it's tracked again:
  2951. else if (!_PyObject_GC_IS_TRACKED(result)) {
  2952. _PyObject_GC_TRACK(result);
  2953. }
  2954. /* Now, we've got the only copy so we can update it in-place */
  2955. assert(r == 0 || Py_REFCNT(result) == 1);
  2956. /* Decrement rightmost cycle, moving leftward upon zero rollover */
  2957. for (i=r-1 ; i>=0 ; i--) {
  2958. cycles[i] -= 1;
  2959. if (cycles[i] == 0) {
  2960. /* rotatation: indices[i:] = indices[i+1:] + indices[i:i+1] */
  2961. index = indices[i];
  2962. for (j=i ; j<n-1 ; j++)
  2963. indices[j] = indices[j+1];
  2964. indices[n-1] = index;
  2965. cycles[i] = n - i;
  2966. } else {
  2967. j = cycles[i];
  2968. index = indices[i];
  2969. indices[i] = indices[n-j];
  2970. indices[n-j] = index;
  2971. for (k=i; k<r ; k++) {
  2972. /* start with i, the leftmost element that changed */
  2973. /* yield tuple(pool[k] for k in indices[:r]) */
  2974. index = indices[k];
  2975. elem = PyTuple_GET_ITEM(pool, index);
  2976. Py_INCREF(elem);
  2977. oldelem = PyTuple_GET_ITEM(result, k);
  2978. PyTuple_SET_ITEM(result, k, elem);
  2979. Py_DECREF(oldelem);
  2980. }
  2981. break;
  2982. }
  2983. }
  2984. /* If i is negative, then the cycles have all
  2985. rolled-over and we're done. */
  2986. if (i < 0)
  2987. goto empty;
  2988. }
  2989. return Py_NewRef(result);
  2990. empty:
  2991. po->stopped = 1;
  2992. return NULL;
  2993. }
  2994. static PyObject *
  2995. permutations_reduce(permutationsobject *po, PyObject *Py_UNUSED(ignored))
  2996. {
  2997. ITERTOOL_PICKLE_DEPRECATION;
  2998. if (po->result == NULL) {
  2999. return Py_BuildValue("O(On)", Py_TYPE(po), po->pool, po->r);
  3000. } else if (po->stopped) {
  3001. return Py_BuildValue("O(()n)", Py_TYPE(po), po->r);
  3002. } else {
  3003. PyObject *indices=NULL, *cycles=NULL;
  3004. Py_ssize_t n, i;
  3005. /* we must pickle the indices and cycles and use them for setstate */
  3006. n = PyTuple_GET_SIZE(po->pool);
  3007. indices = PyTuple_New(n);
  3008. if (indices == NULL)
  3009. goto err;
  3010. for (i=0; i<n; i++) {
  3011. PyObject* index = PyLong_FromSsize_t(po->indices[i]);
  3012. if (!index)
  3013. goto err;
  3014. PyTuple_SET_ITEM(indices, i, index);
  3015. }
  3016. cycles = PyTuple_New(po->r);
  3017. if (cycles == NULL)
  3018. goto err;
  3019. for (i=0 ; i<po->r ; i++) {
  3020. PyObject* index = PyLong_FromSsize_t(po->cycles[i]);
  3021. if (!index)
  3022. goto err;
  3023. PyTuple_SET_ITEM(cycles, i, index);
  3024. }
  3025. return Py_BuildValue("O(On)(NN)", Py_TYPE(po),
  3026. po->pool, po->r,
  3027. indices, cycles);
  3028. err:
  3029. Py_XDECREF(indices);
  3030. Py_XDECREF(cycles);
  3031. return NULL;
  3032. }
  3033. }
  3034. static PyObject *
  3035. permutations_setstate(permutationsobject *po, PyObject *state)
  3036. {
  3037. ITERTOOL_PICKLE_DEPRECATION;
  3038. PyObject *indices, *cycles, *result;
  3039. Py_ssize_t n, i;
  3040. if (!PyTuple_Check(state)) {
  3041. PyErr_SetString(PyExc_TypeError, "state is not a tuple");
  3042. return NULL;
  3043. }
  3044. if (!PyArg_ParseTuple(state, "O!O!",
  3045. &PyTuple_Type, &indices,
  3046. &PyTuple_Type, &cycles)) {
  3047. return NULL;
  3048. }
  3049. n = PyTuple_GET_SIZE(po->pool);
  3050. if (PyTuple_GET_SIZE(indices) != n || PyTuple_GET_SIZE(cycles) != po->r) {
  3051. PyErr_SetString(PyExc_ValueError, "invalid arguments");
  3052. return NULL;
  3053. }
  3054. for (i=0; i<n; i++) {
  3055. PyObject* indexObject = PyTuple_GET_ITEM(indices, i);
  3056. Py_ssize_t index = PyLong_AsSsize_t(indexObject);
  3057. if (index < 0 && PyErr_Occurred())
  3058. return NULL; /* not an integer */
  3059. /* clamp the index */
  3060. if (index < 0)
  3061. index = 0;
  3062. else if (index > n-1)
  3063. index = n-1;
  3064. po->indices[i] = index;
  3065. }
  3066. for (i=0; i<po->r; i++) {
  3067. PyObject* indexObject = PyTuple_GET_ITEM(cycles, i);
  3068. Py_ssize_t index = PyLong_AsSsize_t(indexObject);
  3069. if (index < 0 && PyErr_Occurred())
  3070. return NULL; /* not an integer */
  3071. if (index < 1)
  3072. index = 1;
  3073. else if (index > n-i)
  3074. index = n-i;
  3075. po->cycles[i] = index;
  3076. }
  3077. result = PyTuple_New(po->r);
  3078. if (result == NULL)
  3079. return NULL;
  3080. for (i=0; i<po->r; i++) {
  3081. PyObject *element = PyTuple_GET_ITEM(po->pool, po->indices[i]);
  3082. Py_INCREF(element);
  3083. PyTuple_SET_ITEM(result, i, element);
  3084. }
  3085. Py_XSETREF(po->result, result);
  3086. Py_RETURN_NONE;
  3087. }
  3088. static PyMethodDef permuations_methods[] = {
  3089. {"__reduce__", (PyCFunction)permutations_reduce, METH_NOARGS,
  3090. reduce_doc},
  3091. {"__setstate__", (PyCFunction)permutations_setstate, METH_O,
  3092. setstate_doc},
  3093. {"__sizeof__", (PyCFunction)permutations_sizeof, METH_NOARGS,
  3094. sizeof_doc},
  3095. {NULL, NULL} /* sentinel */
  3096. };
  3097. static PyType_Slot permutations_slots[] = {
  3098. {Py_tp_dealloc, permutations_dealloc},
  3099. {Py_tp_getattro, PyObject_GenericGetAttr},
  3100. {Py_tp_doc, (void *)itertools_permutations__doc__},
  3101. {Py_tp_traverse, permutations_traverse},
  3102. {Py_tp_iter, PyObject_SelfIter},
  3103. {Py_tp_iternext, permutations_next},
  3104. {Py_tp_methods, permuations_methods},
  3105. {Py_tp_new, itertools_permutations},
  3106. {Py_tp_free, PyObject_GC_Del},
  3107. {0, NULL},
  3108. };
  3109. static PyType_Spec permutations_spec = {
  3110. .name = "itertools.permutations",
  3111. .basicsize = sizeof(permutationsobject),
  3112. .flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_BASETYPE |
  3113. Py_TPFLAGS_IMMUTABLETYPE),
  3114. .slots = permutations_slots,
  3115. };
  3116. /* accumulate object ********************************************************/
  3117. typedef struct {
  3118. PyObject_HEAD
  3119. PyObject *total;
  3120. PyObject *it;
  3121. PyObject *binop;
  3122. PyObject *initial;
  3123. itertools_state *state;
  3124. } accumulateobject;
  3125. /*[clinic input]
  3126. @classmethod
  3127. itertools.accumulate.__new__
  3128. iterable: object
  3129. func as binop: object = None
  3130. *
  3131. initial: object = None
  3132. Return series of accumulated sums (or other binary function results).
  3133. [clinic start generated code]*/
  3134. static PyObject *
  3135. itertools_accumulate_impl(PyTypeObject *type, PyObject *iterable,
  3136. PyObject *binop, PyObject *initial)
  3137. /*[clinic end generated code: output=66da2650627128f8 input=c4ce20ac59bf7ffd]*/
  3138. {
  3139. PyObject *it;
  3140. accumulateobject *lz;
  3141. /* Get iterator. */
  3142. it = PyObject_GetIter(iterable);
  3143. if (it == NULL)
  3144. return NULL;
  3145. /* create accumulateobject structure */
  3146. lz = (accumulateobject *)type->tp_alloc(type, 0);
  3147. if (lz == NULL) {
  3148. Py_DECREF(it);
  3149. return NULL;
  3150. }
  3151. if (binop != Py_None) {
  3152. lz->binop = Py_XNewRef(binop);
  3153. }
  3154. lz->total = NULL;
  3155. lz->it = it;
  3156. lz->initial = Py_XNewRef(initial);
  3157. lz->state = find_state_by_type(type);
  3158. return (PyObject *)lz;
  3159. }
  3160. static void
  3161. accumulate_dealloc(accumulateobject *lz)
  3162. {
  3163. PyTypeObject *tp = Py_TYPE(lz);
  3164. PyObject_GC_UnTrack(lz);
  3165. Py_XDECREF(lz->binop);
  3166. Py_XDECREF(lz->total);
  3167. Py_XDECREF(lz->it);
  3168. Py_XDECREF(lz->initial);
  3169. tp->tp_free(lz);
  3170. Py_DECREF(tp);
  3171. }
  3172. static int
  3173. accumulate_traverse(accumulateobject *lz, visitproc visit, void *arg)
  3174. {
  3175. Py_VISIT(Py_TYPE(lz));
  3176. Py_VISIT(lz->binop);
  3177. Py_VISIT(lz->it);
  3178. Py_VISIT(lz->total);
  3179. Py_VISIT(lz->initial);
  3180. return 0;
  3181. }
  3182. static PyObject *
  3183. accumulate_next(accumulateobject *lz)
  3184. {
  3185. PyObject *val, *newtotal;
  3186. if (lz->initial != Py_None) {
  3187. lz->total = lz->initial;
  3188. lz->initial = Py_NewRef(Py_None);
  3189. return Py_NewRef(lz->total);
  3190. }
  3191. val = (*Py_TYPE(lz->it)->tp_iternext)(lz->it);
  3192. if (val == NULL)
  3193. return NULL;
  3194. if (lz->total == NULL) {
  3195. lz->total = Py_NewRef(val);
  3196. return lz->total;
  3197. }
  3198. if (lz->binop == NULL)
  3199. newtotal = PyNumber_Add(lz->total, val);
  3200. else
  3201. newtotal = PyObject_CallFunctionObjArgs(lz->binop, lz->total, val, NULL);
  3202. Py_DECREF(val);
  3203. if (newtotal == NULL)
  3204. return NULL;
  3205. Py_INCREF(newtotal);
  3206. Py_SETREF(lz->total, newtotal);
  3207. return newtotal;
  3208. }
  3209. static PyObject *
  3210. accumulate_reduce(accumulateobject *lz, PyObject *Py_UNUSED(ignored))
  3211. {
  3212. ITERTOOL_PICKLE_DEPRECATION;
  3213. itertools_state *state = lz->state;
  3214. if (lz->initial != Py_None) {
  3215. PyObject *it;
  3216. assert(lz->total == NULL);
  3217. it = PyObject_CallFunction((PyObject *)(state->chain_type), "(O)O",
  3218. lz->initial, lz->it);
  3219. if (it == NULL)
  3220. return NULL;
  3221. return Py_BuildValue("O(NO)O", Py_TYPE(lz),
  3222. it, lz->binop?lz->binop:Py_None, Py_None);
  3223. }
  3224. if (lz->total == Py_None) {
  3225. PyObject *it;
  3226. it = PyObject_CallFunction((PyObject *)(state->chain_type), "(O)O",
  3227. lz->total, lz->it);
  3228. if (it == NULL)
  3229. return NULL;
  3230. it = PyObject_CallFunction((PyObject *)Py_TYPE(lz), "NO",
  3231. it, lz->binop ? lz->binop : Py_None);
  3232. if (it == NULL)
  3233. return NULL;
  3234. return Py_BuildValue("O(NiO)", state->islice_type, it, 1, Py_None);
  3235. }
  3236. return Py_BuildValue("O(OO)O", Py_TYPE(lz),
  3237. lz->it, lz->binop?lz->binop:Py_None,
  3238. lz->total?lz->total:Py_None);
  3239. }
  3240. static PyObject *
  3241. accumulate_setstate(accumulateobject *lz, PyObject *state)
  3242. {
  3243. ITERTOOL_PICKLE_DEPRECATION;
  3244. Py_INCREF(state);
  3245. Py_XSETREF(lz->total, state);
  3246. Py_RETURN_NONE;
  3247. }
  3248. static PyMethodDef accumulate_methods[] = {
  3249. {"__reduce__", (PyCFunction)accumulate_reduce, METH_NOARGS,
  3250. reduce_doc},
  3251. {"__setstate__", (PyCFunction)accumulate_setstate, METH_O,
  3252. setstate_doc},
  3253. {NULL, NULL} /* sentinel */
  3254. };
  3255. static PyType_Slot accumulate_slots[] = {
  3256. {Py_tp_dealloc, accumulate_dealloc},
  3257. {Py_tp_getattro, PyObject_GenericGetAttr},
  3258. {Py_tp_doc, (void *)itertools_accumulate__doc__},
  3259. {Py_tp_traverse, accumulate_traverse},
  3260. {Py_tp_iter, PyObject_SelfIter},
  3261. {Py_tp_iternext, accumulate_next},
  3262. {Py_tp_methods, accumulate_methods},
  3263. {Py_tp_new, itertools_accumulate},
  3264. {Py_tp_free, PyObject_GC_Del},
  3265. {0, NULL},
  3266. };
  3267. static PyType_Spec accumulate_spec = {
  3268. .name = "itertools.accumulate",
  3269. .basicsize = sizeof(accumulateobject),
  3270. .flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_BASETYPE |
  3271. Py_TPFLAGS_IMMUTABLETYPE),
  3272. .slots = accumulate_slots,
  3273. };
  3274. /* compress object ************************************************************/
  3275. /* Equivalent to:
  3276. def compress(data, selectors):
  3277. "compress('ABCDEF', [1,0,1,0,1,1]) --> A C E F"
  3278. return (d for d, s in zip(data, selectors) if s)
  3279. */
  3280. typedef struct {
  3281. PyObject_HEAD
  3282. PyObject *data;
  3283. PyObject *selectors;
  3284. } compressobject;
  3285. /*[clinic input]
  3286. @classmethod
  3287. itertools.compress.__new__
  3288. data as seq1: object
  3289. selectors as seq2: object
  3290. Return data elements corresponding to true selector elements.
  3291. Forms a shorter iterator from selected data elements using the selectors to
  3292. choose the data elements.
  3293. [clinic start generated code]*/
  3294. static PyObject *
  3295. itertools_compress_impl(PyTypeObject *type, PyObject *seq1, PyObject *seq2)
  3296. /*[clinic end generated code: output=7e67157212ed09e0 input=79596d7cd20c77e5]*/
  3297. {
  3298. PyObject *data=NULL, *selectors=NULL;
  3299. compressobject *lz;
  3300. data = PyObject_GetIter(seq1);
  3301. if (data == NULL)
  3302. goto fail;
  3303. selectors = PyObject_GetIter(seq2);
  3304. if (selectors == NULL)
  3305. goto fail;
  3306. /* create compressobject structure */
  3307. lz = (compressobject *)type->tp_alloc(type, 0);
  3308. if (lz == NULL)
  3309. goto fail;
  3310. lz->data = data;
  3311. lz->selectors = selectors;
  3312. return (PyObject *)lz;
  3313. fail:
  3314. Py_XDECREF(data);
  3315. Py_XDECREF(selectors);
  3316. return NULL;
  3317. }
  3318. static void
  3319. compress_dealloc(compressobject *lz)
  3320. {
  3321. PyTypeObject *tp = Py_TYPE(lz);
  3322. PyObject_GC_UnTrack(lz);
  3323. Py_XDECREF(lz->data);
  3324. Py_XDECREF(lz->selectors);
  3325. tp->tp_free(lz);
  3326. Py_DECREF(tp);
  3327. }
  3328. static int
  3329. compress_traverse(compressobject *lz, visitproc visit, void *arg)
  3330. {
  3331. Py_VISIT(Py_TYPE(lz));
  3332. Py_VISIT(lz->data);
  3333. Py_VISIT(lz->selectors);
  3334. return 0;
  3335. }
  3336. static PyObject *
  3337. compress_next(compressobject *lz)
  3338. {
  3339. PyObject *data = lz->data, *selectors = lz->selectors;
  3340. PyObject *datum, *selector;
  3341. PyObject *(*datanext)(PyObject *) = *Py_TYPE(data)->tp_iternext;
  3342. PyObject *(*selectornext)(PyObject *) = *Py_TYPE(selectors)->tp_iternext;
  3343. int ok;
  3344. while (1) {
  3345. /* Steps: get datum, get selector, evaluate selector.
  3346. Order is important (to match the pure python version
  3347. in terms of which input gets a chance to raise an
  3348. exception first).
  3349. */
  3350. datum = datanext(data);
  3351. if (datum == NULL)
  3352. return NULL;
  3353. selector = selectornext(selectors);
  3354. if (selector == NULL) {
  3355. Py_DECREF(datum);
  3356. return NULL;
  3357. }
  3358. ok = PyObject_IsTrue(selector);
  3359. Py_DECREF(selector);
  3360. if (ok > 0)
  3361. return datum;
  3362. Py_DECREF(datum);
  3363. if (ok < 0)
  3364. return NULL;
  3365. }
  3366. }
  3367. static PyObject *
  3368. compress_reduce(compressobject *lz, PyObject *Py_UNUSED(ignored))
  3369. {
  3370. ITERTOOL_PICKLE_DEPRECATION;
  3371. return Py_BuildValue("O(OO)", Py_TYPE(lz),
  3372. lz->data, lz->selectors);
  3373. }
  3374. static PyMethodDef compress_methods[] = {
  3375. {"__reduce__", (PyCFunction)compress_reduce, METH_NOARGS,
  3376. reduce_doc},
  3377. {NULL, NULL} /* sentinel */
  3378. };
  3379. static PyType_Slot compress_slots[] = {
  3380. {Py_tp_dealloc, compress_dealloc},
  3381. {Py_tp_getattro, PyObject_GenericGetAttr},
  3382. {Py_tp_doc, (void *)itertools_compress__doc__},
  3383. {Py_tp_traverse, compress_traverse},
  3384. {Py_tp_iter, PyObject_SelfIter},
  3385. {Py_tp_iternext, compress_next},
  3386. {Py_tp_methods, compress_methods},
  3387. {Py_tp_new, itertools_compress},
  3388. {Py_tp_free, PyObject_GC_Del},
  3389. {0, NULL},
  3390. };
  3391. static PyType_Spec compress_spec = {
  3392. .name = "itertools.compress",
  3393. .basicsize = sizeof(compressobject),
  3394. .flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_BASETYPE |
  3395. Py_TPFLAGS_IMMUTABLETYPE),
  3396. .slots = compress_slots,
  3397. };
  3398. /* filterfalse object ************************************************************/
  3399. typedef struct {
  3400. PyObject_HEAD
  3401. PyObject *func;
  3402. PyObject *it;
  3403. } filterfalseobject;
  3404. /*[clinic input]
  3405. @classmethod
  3406. itertools.filterfalse.__new__
  3407. function as func: object
  3408. iterable as seq: object
  3409. /
  3410. Return those items of iterable for which function(item) is false.
  3411. If function is None, return the items that are false.
  3412. [clinic start generated code]*/
  3413. static PyObject *
  3414. itertools_filterfalse_impl(PyTypeObject *type, PyObject *func, PyObject *seq)
  3415. /*[clinic end generated code: output=55f87eab9fc0484e input=2d684a2c66f99cde]*/
  3416. {
  3417. PyObject *it;
  3418. filterfalseobject *lz;
  3419. /* Get iterator. */
  3420. it = PyObject_GetIter(seq);
  3421. if (it == NULL)
  3422. return NULL;
  3423. /* create filterfalseobject structure */
  3424. lz = (filterfalseobject *)type->tp_alloc(type, 0);
  3425. if (lz == NULL) {
  3426. Py_DECREF(it);
  3427. return NULL;
  3428. }
  3429. lz->func = Py_NewRef(func);
  3430. lz->it = it;
  3431. return (PyObject *)lz;
  3432. }
  3433. static void
  3434. filterfalse_dealloc(filterfalseobject *lz)
  3435. {
  3436. PyTypeObject *tp = Py_TYPE(lz);
  3437. PyObject_GC_UnTrack(lz);
  3438. Py_XDECREF(lz->func);
  3439. Py_XDECREF(lz->it);
  3440. tp->tp_free(lz);
  3441. Py_DECREF(tp);
  3442. }
  3443. static int
  3444. filterfalse_traverse(filterfalseobject *lz, visitproc visit, void *arg)
  3445. {
  3446. Py_VISIT(Py_TYPE(lz));
  3447. Py_VISIT(lz->it);
  3448. Py_VISIT(lz->func);
  3449. return 0;
  3450. }
  3451. static PyObject *
  3452. filterfalse_next(filterfalseobject *lz)
  3453. {
  3454. PyObject *item;
  3455. PyObject *it = lz->it;
  3456. long ok;
  3457. PyObject *(*iternext)(PyObject *);
  3458. iternext = *Py_TYPE(it)->tp_iternext;
  3459. for (;;) {
  3460. item = iternext(it);
  3461. if (item == NULL)
  3462. return NULL;
  3463. if (lz->func == Py_None || lz->func == (PyObject *)&PyBool_Type) {
  3464. ok = PyObject_IsTrue(item);
  3465. } else {
  3466. PyObject *good;
  3467. good = PyObject_CallOneArg(lz->func, item);
  3468. if (good == NULL) {
  3469. Py_DECREF(item);
  3470. return NULL;
  3471. }
  3472. ok = PyObject_IsTrue(good);
  3473. Py_DECREF(good);
  3474. }
  3475. if (ok == 0)
  3476. return item;
  3477. Py_DECREF(item);
  3478. if (ok < 0)
  3479. return NULL;
  3480. }
  3481. }
  3482. static PyObject *
  3483. filterfalse_reduce(filterfalseobject *lz, PyObject *Py_UNUSED(ignored))
  3484. {
  3485. ITERTOOL_PICKLE_DEPRECATION;
  3486. return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->func, lz->it);
  3487. }
  3488. static PyMethodDef filterfalse_methods[] = {
  3489. {"__reduce__", (PyCFunction)filterfalse_reduce, METH_NOARGS,
  3490. reduce_doc},
  3491. {NULL, NULL} /* sentinel */
  3492. };
  3493. static PyType_Slot filterfalse_slots[] = {
  3494. {Py_tp_dealloc, filterfalse_dealloc},
  3495. {Py_tp_getattro, PyObject_GenericGetAttr},
  3496. {Py_tp_doc, (void *)itertools_filterfalse__doc__},
  3497. {Py_tp_traverse, filterfalse_traverse},
  3498. {Py_tp_iter, PyObject_SelfIter},
  3499. {Py_tp_iternext, filterfalse_next},
  3500. {Py_tp_methods, filterfalse_methods},
  3501. {Py_tp_new, itertools_filterfalse},
  3502. {Py_tp_free, PyObject_GC_Del},
  3503. {0, NULL},
  3504. };
  3505. static PyType_Spec filterfalse_spec = {
  3506. .name = "itertools.filterfalse",
  3507. .basicsize = sizeof(filterfalseobject),
  3508. .flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_BASETYPE |
  3509. Py_TPFLAGS_IMMUTABLETYPE),
  3510. .slots = filterfalse_slots,
  3511. };
  3512. /* count object ************************************************************/
  3513. typedef struct {
  3514. PyObject_HEAD
  3515. Py_ssize_t cnt;
  3516. PyObject *long_cnt;
  3517. PyObject *long_step;
  3518. } countobject;
  3519. /* Counting logic and invariants:
  3520. fast_mode: when cnt an integer < PY_SSIZE_T_MAX and no step is specified.
  3521. assert(cnt != PY_SSIZE_T_MAX && long_cnt == NULL && long_step==PyLong(1));
  3522. Advances with: cnt += 1
  3523. When count hits Y_SSIZE_T_MAX, switch to slow_mode.
  3524. slow_mode: when cnt == PY_SSIZE_T_MAX, step is not int(1), or cnt is a float.
  3525. assert(cnt == PY_SSIZE_T_MAX && long_cnt != NULL && long_step != NULL);
  3526. All counting is done with python objects (no overflows or underflows).
  3527. Advances with: long_cnt += long_step
  3528. Step may be zero -- effectively a slow version of repeat(cnt).
  3529. Either long_cnt or long_step may be a float, Fraction, or Decimal.
  3530. */
  3531. /*[clinic input]
  3532. @classmethod
  3533. itertools.count.__new__
  3534. start as long_cnt: object(c_default="NULL") = 0
  3535. step as long_step: object(c_default="NULL") = 1
  3536. Return a count object whose .__next__() method returns consecutive values.
  3537. Equivalent to:
  3538. def count(firstval=0, step=1):
  3539. x = firstval
  3540. while 1:
  3541. yield x
  3542. x += step
  3543. [clinic start generated code]*/
  3544. static PyObject *
  3545. itertools_count_impl(PyTypeObject *type, PyObject *long_cnt,
  3546. PyObject *long_step)
  3547. /*[clinic end generated code: output=09a9250aebd00b1c input=d7a85eec18bfcd94]*/
  3548. {
  3549. countobject *lz;
  3550. int fast_mode;
  3551. Py_ssize_t cnt = 0;
  3552. long step;
  3553. if ((long_cnt != NULL && !PyNumber_Check(long_cnt)) ||
  3554. (long_step != NULL && !PyNumber_Check(long_step))) {
  3555. PyErr_SetString(PyExc_TypeError, "a number is required");
  3556. return NULL;
  3557. }
  3558. fast_mode = (long_cnt == NULL || PyLong_Check(long_cnt)) &&
  3559. (long_step == NULL || PyLong_Check(long_step));
  3560. /* If not specified, start defaults to 0 */
  3561. if (long_cnt != NULL) {
  3562. if (fast_mode) {
  3563. assert(PyLong_Check(long_cnt));
  3564. cnt = PyLong_AsSsize_t(long_cnt);
  3565. if (cnt == -1 && PyErr_Occurred()) {
  3566. PyErr_Clear();
  3567. fast_mode = 0;
  3568. }
  3569. }
  3570. } else {
  3571. cnt = 0;
  3572. long_cnt = _PyLong_GetZero();
  3573. }
  3574. Py_INCREF(long_cnt);
  3575. /* If not specified, step defaults to 1 */
  3576. if (long_step == NULL) {
  3577. long_step = _PyLong_GetOne();
  3578. }
  3579. Py_INCREF(long_step);
  3580. assert(long_cnt != NULL && long_step != NULL);
  3581. /* Fast mode only works when the step is 1 */
  3582. if (fast_mode) {
  3583. assert(PyLong_Check(long_step));
  3584. step = PyLong_AsLong(long_step);
  3585. if (step != 1) {
  3586. fast_mode = 0;
  3587. if (step == -1 && PyErr_Occurred())
  3588. PyErr_Clear();
  3589. }
  3590. }
  3591. if (fast_mode)
  3592. Py_CLEAR(long_cnt);
  3593. else
  3594. cnt = PY_SSIZE_T_MAX;
  3595. assert((cnt != PY_SSIZE_T_MAX && long_cnt == NULL && fast_mode) ||
  3596. (cnt == PY_SSIZE_T_MAX && long_cnt != NULL && !fast_mode));
  3597. assert(!fast_mode ||
  3598. (PyLong_Check(long_step) && PyLong_AS_LONG(long_step) == 1));
  3599. /* create countobject structure */
  3600. lz = (countobject *)type->tp_alloc(type, 0);
  3601. if (lz == NULL) {
  3602. Py_XDECREF(long_cnt);
  3603. Py_DECREF(long_step);
  3604. return NULL;
  3605. }
  3606. lz->cnt = cnt;
  3607. lz->long_cnt = long_cnt;
  3608. lz->long_step = long_step;
  3609. return (PyObject *)lz;
  3610. }
  3611. static void
  3612. count_dealloc(countobject *lz)
  3613. {
  3614. PyTypeObject *tp = Py_TYPE(lz);
  3615. PyObject_GC_UnTrack(lz);
  3616. Py_XDECREF(lz->long_cnt);
  3617. Py_XDECREF(lz->long_step);
  3618. tp->tp_free(lz);
  3619. Py_DECREF(tp);
  3620. }
  3621. static int
  3622. count_traverse(countobject *lz, visitproc visit, void *arg)
  3623. {
  3624. Py_VISIT(Py_TYPE(lz));
  3625. Py_VISIT(lz->long_cnt);
  3626. Py_VISIT(lz->long_step);
  3627. return 0;
  3628. }
  3629. static PyObject *
  3630. count_nextlong(countobject *lz)
  3631. {
  3632. PyObject *long_cnt;
  3633. PyObject *stepped_up;
  3634. long_cnt = lz->long_cnt;
  3635. if (long_cnt == NULL) {
  3636. /* Switch to slow_mode */
  3637. long_cnt = PyLong_FromSsize_t(PY_SSIZE_T_MAX);
  3638. if (long_cnt == NULL)
  3639. return NULL;
  3640. }
  3641. assert(lz->cnt == PY_SSIZE_T_MAX && long_cnt != NULL);
  3642. stepped_up = PyNumber_Add(long_cnt, lz->long_step);
  3643. if (stepped_up == NULL)
  3644. return NULL;
  3645. lz->long_cnt = stepped_up;
  3646. return long_cnt;
  3647. }
  3648. static PyObject *
  3649. count_next(countobject *lz)
  3650. {
  3651. if (lz->cnt == PY_SSIZE_T_MAX)
  3652. return count_nextlong(lz);
  3653. return PyLong_FromSsize_t(lz->cnt++);
  3654. }
  3655. static PyObject *
  3656. count_repr(countobject *lz)
  3657. {
  3658. if (lz->cnt != PY_SSIZE_T_MAX)
  3659. return PyUnicode_FromFormat("%s(%zd)",
  3660. _PyType_Name(Py_TYPE(lz)), lz->cnt);
  3661. if (PyLong_Check(lz->long_step)) {
  3662. long step = PyLong_AsLong(lz->long_step);
  3663. if (step == -1 && PyErr_Occurred()) {
  3664. PyErr_Clear();
  3665. }
  3666. if (step == 1) {
  3667. /* Don't display step when it is an integer equal to 1 */
  3668. return PyUnicode_FromFormat("%s(%R)",
  3669. _PyType_Name(Py_TYPE(lz)),
  3670. lz->long_cnt);
  3671. }
  3672. }
  3673. return PyUnicode_FromFormat("%s(%R, %R)",
  3674. _PyType_Name(Py_TYPE(lz)),
  3675. lz->long_cnt, lz->long_step);
  3676. }
  3677. static PyObject *
  3678. count_reduce(countobject *lz, PyObject *Py_UNUSED(ignored))
  3679. {
  3680. ITERTOOL_PICKLE_DEPRECATION;
  3681. if (lz->cnt == PY_SSIZE_T_MAX)
  3682. return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->long_cnt, lz->long_step);
  3683. return Py_BuildValue("O(n)", Py_TYPE(lz), lz->cnt);
  3684. }
  3685. static PyMethodDef count_methods[] = {
  3686. {"__reduce__", (PyCFunction)count_reduce, METH_NOARGS,
  3687. reduce_doc},
  3688. {NULL, NULL} /* sentinel */
  3689. };
  3690. static PyType_Slot count_slots[] = {
  3691. {Py_tp_dealloc, count_dealloc},
  3692. {Py_tp_repr, count_repr},
  3693. {Py_tp_getattro, PyObject_GenericGetAttr},
  3694. {Py_tp_doc, (void *)itertools_count__doc__},
  3695. {Py_tp_traverse, count_traverse},
  3696. {Py_tp_iter, PyObject_SelfIter},
  3697. {Py_tp_iternext, count_next},
  3698. {Py_tp_methods, count_methods},
  3699. {Py_tp_new, itertools_count},
  3700. {Py_tp_free, PyObject_GC_Del},
  3701. {0, NULL},
  3702. };
  3703. static PyType_Spec count_spec = {
  3704. .name = "itertools.count",
  3705. .basicsize = sizeof(countobject),
  3706. .flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_BASETYPE |
  3707. Py_TPFLAGS_IMMUTABLETYPE),
  3708. .slots = count_slots,
  3709. };
  3710. /* repeat object ************************************************************/
  3711. typedef struct {
  3712. PyObject_HEAD
  3713. PyObject *element;
  3714. Py_ssize_t cnt;
  3715. } repeatobject;
  3716. static PyObject *
  3717. repeat_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
  3718. {
  3719. repeatobject *ro;
  3720. PyObject *element;
  3721. Py_ssize_t cnt = -1, n_args;
  3722. static char *kwargs[] = {"object", "times", NULL};
  3723. n_args = PyTuple_GET_SIZE(args);
  3724. if (kwds != NULL)
  3725. n_args += PyDict_GET_SIZE(kwds);
  3726. if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|n:repeat", kwargs,
  3727. &element, &cnt))
  3728. return NULL;
  3729. /* Does user supply times argument? */
  3730. if (n_args == 2 && cnt < 0)
  3731. cnt = 0;
  3732. ro = (repeatobject *)type->tp_alloc(type, 0);
  3733. if (ro == NULL)
  3734. return NULL;
  3735. ro->element = Py_NewRef(element);
  3736. ro->cnt = cnt;
  3737. return (PyObject *)ro;
  3738. }
  3739. static void
  3740. repeat_dealloc(repeatobject *ro)
  3741. {
  3742. PyTypeObject *tp = Py_TYPE(ro);
  3743. PyObject_GC_UnTrack(ro);
  3744. Py_XDECREF(ro->element);
  3745. tp->tp_free(ro);
  3746. Py_DECREF(tp);
  3747. }
  3748. static int
  3749. repeat_traverse(repeatobject *ro, visitproc visit, void *arg)
  3750. {
  3751. Py_VISIT(Py_TYPE(ro));
  3752. Py_VISIT(ro->element);
  3753. return 0;
  3754. }
  3755. static PyObject *
  3756. repeat_next(repeatobject *ro)
  3757. {
  3758. if (ro->cnt == 0)
  3759. return NULL;
  3760. if (ro->cnt > 0)
  3761. ro->cnt--;
  3762. return Py_NewRef(ro->element);
  3763. }
  3764. static PyObject *
  3765. repeat_repr(repeatobject *ro)
  3766. {
  3767. if (ro->cnt == -1)
  3768. return PyUnicode_FromFormat("%s(%R)",
  3769. _PyType_Name(Py_TYPE(ro)), ro->element);
  3770. else
  3771. return PyUnicode_FromFormat("%s(%R, %zd)",
  3772. _PyType_Name(Py_TYPE(ro)), ro->element,
  3773. ro->cnt);
  3774. }
  3775. static PyObject *
  3776. repeat_len(repeatobject *ro, PyObject *Py_UNUSED(ignored))
  3777. {
  3778. if (ro->cnt == -1) {
  3779. PyErr_SetString(PyExc_TypeError, "len() of unsized object");
  3780. return NULL;
  3781. }
  3782. return PyLong_FromSize_t(ro->cnt);
  3783. }
  3784. PyDoc_STRVAR(length_hint_doc, "Private method returning an estimate of len(list(it)).");
  3785. static PyObject *
  3786. repeat_reduce(repeatobject *ro, PyObject *Py_UNUSED(ignored))
  3787. {
  3788. ITERTOOL_PICKLE_DEPRECATION;
  3789. /* unpickle this so that a new repeat iterator is constructed with an
  3790. * object, then call __setstate__ on it to set cnt
  3791. */
  3792. if (ro->cnt >= 0)
  3793. return Py_BuildValue("O(On)", Py_TYPE(ro), ro->element, ro->cnt);
  3794. else
  3795. return Py_BuildValue("O(O)", Py_TYPE(ro), ro->element);
  3796. }
  3797. static PyMethodDef repeat_methods[] = {
  3798. {"__length_hint__", (PyCFunction)repeat_len, METH_NOARGS, length_hint_doc},
  3799. {"__reduce__", (PyCFunction)repeat_reduce, METH_NOARGS, reduce_doc},
  3800. {NULL, NULL} /* sentinel */
  3801. };
  3802. PyDoc_STRVAR(repeat_doc,
  3803. "repeat(object [,times]) -> create an iterator which returns the object\n\
  3804. for the specified number of times. If not specified, returns the object\n\
  3805. endlessly.");
  3806. static PyType_Slot repeat_slots[] = {
  3807. {Py_tp_dealloc, repeat_dealloc},
  3808. {Py_tp_repr, repeat_repr},
  3809. {Py_tp_getattro, PyObject_GenericGetAttr},
  3810. {Py_tp_doc, (void *)repeat_doc},
  3811. {Py_tp_traverse, repeat_traverse},
  3812. {Py_tp_iter, PyObject_SelfIter},
  3813. {Py_tp_iternext, repeat_next},
  3814. {Py_tp_methods, repeat_methods},
  3815. {Py_tp_new, repeat_new},
  3816. {Py_tp_free, PyObject_GC_Del},
  3817. {0, NULL},
  3818. };
  3819. static PyType_Spec repeat_spec = {
  3820. .name = "itertools.repeat",
  3821. .basicsize = sizeof(repeatobject),
  3822. .flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_BASETYPE |
  3823. Py_TPFLAGS_IMMUTABLETYPE),
  3824. .slots = repeat_slots,
  3825. };
  3826. /* ziplongest object *********************************************************/
  3827. typedef struct {
  3828. PyObject_HEAD
  3829. Py_ssize_t tuplesize;
  3830. Py_ssize_t numactive;
  3831. PyObject *ittuple; /* tuple of iterators */
  3832. PyObject *result;
  3833. PyObject *fillvalue;
  3834. } ziplongestobject;
  3835. static PyObject *
  3836. zip_longest_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
  3837. {
  3838. ziplongestobject *lz;
  3839. Py_ssize_t i;
  3840. PyObject *ittuple; /* tuple of iterators */
  3841. PyObject *result;
  3842. PyObject *fillvalue = Py_None;
  3843. Py_ssize_t tuplesize;
  3844. if (kwds != NULL && PyDict_CheckExact(kwds) && PyDict_GET_SIZE(kwds) > 0) {
  3845. fillvalue = NULL;
  3846. if (PyDict_GET_SIZE(kwds) == 1) {
  3847. fillvalue = PyDict_GetItemWithError(kwds, &_Py_ID(fillvalue));
  3848. }
  3849. if (fillvalue == NULL) {
  3850. if (!PyErr_Occurred()) {
  3851. PyErr_SetString(PyExc_TypeError,
  3852. "zip_longest() got an unexpected keyword argument");
  3853. }
  3854. return NULL;
  3855. }
  3856. }
  3857. /* args must be a tuple */
  3858. assert(PyTuple_Check(args));
  3859. tuplesize = PyTuple_GET_SIZE(args);
  3860. /* obtain iterators */
  3861. ittuple = PyTuple_New(tuplesize);
  3862. if (ittuple == NULL)
  3863. return NULL;
  3864. for (i=0; i < tuplesize; i++) {
  3865. PyObject *item = PyTuple_GET_ITEM(args, i);
  3866. PyObject *it = PyObject_GetIter(item);
  3867. if (it == NULL) {
  3868. Py_DECREF(ittuple);
  3869. return NULL;
  3870. }
  3871. PyTuple_SET_ITEM(ittuple, i, it);
  3872. }
  3873. /* create a result holder */
  3874. result = PyTuple_New(tuplesize);
  3875. if (result == NULL) {
  3876. Py_DECREF(ittuple);
  3877. return NULL;
  3878. }
  3879. for (i=0 ; i < tuplesize ; i++) {
  3880. Py_INCREF(Py_None);
  3881. PyTuple_SET_ITEM(result, i, Py_None);
  3882. }
  3883. /* create ziplongestobject structure */
  3884. lz = (ziplongestobject *)type->tp_alloc(type, 0);
  3885. if (lz == NULL) {
  3886. Py_DECREF(ittuple);
  3887. Py_DECREF(result);
  3888. return NULL;
  3889. }
  3890. lz->ittuple = ittuple;
  3891. lz->tuplesize = tuplesize;
  3892. lz->numactive = tuplesize;
  3893. lz->result = result;
  3894. lz->fillvalue = Py_NewRef(fillvalue);
  3895. return (PyObject *)lz;
  3896. }
  3897. static void
  3898. zip_longest_dealloc(ziplongestobject *lz)
  3899. {
  3900. PyTypeObject *tp = Py_TYPE(lz);
  3901. PyObject_GC_UnTrack(lz);
  3902. Py_XDECREF(lz->ittuple);
  3903. Py_XDECREF(lz->result);
  3904. Py_XDECREF(lz->fillvalue);
  3905. tp->tp_free(lz);
  3906. Py_DECREF(tp);
  3907. }
  3908. static int
  3909. zip_longest_traverse(ziplongestobject *lz, visitproc visit, void *arg)
  3910. {
  3911. Py_VISIT(Py_TYPE(lz));
  3912. Py_VISIT(lz->ittuple);
  3913. Py_VISIT(lz->result);
  3914. Py_VISIT(lz->fillvalue);
  3915. return 0;
  3916. }
  3917. static PyObject *
  3918. zip_longest_next(ziplongestobject *lz)
  3919. {
  3920. Py_ssize_t i;
  3921. Py_ssize_t tuplesize = lz->tuplesize;
  3922. PyObject *result = lz->result;
  3923. PyObject *it;
  3924. PyObject *item;
  3925. PyObject *olditem;
  3926. if (tuplesize == 0)
  3927. return NULL;
  3928. if (lz->numactive == 0)
  3929. return NULL;
  3930. if (Py_REFCNT(result) == 1) {
  3931. Py_INCREF(result);
  3932. for (i=0 ; i < tuplesize ; i++) {
  3933. it = PyTuple_GET_ITEM(lz->ittuple, i);
  3934. if (it == NULL) {
  3935. item = Py_NewRef(lz->fillvalue);
  3936. } else {
  3937. item = PyIter_Next(it);
  3938. if (item == NULL) {
  3939. lz->numactive -= 1;
  3940. if (lz->numactive == 0 || PyErr_Occurred()) {
  3941. lz->numactive = 0;
  3942. Py_DECREF(result);
  3943. return NULL;
  3944. } else {
  3945. item = Py_NewRef(lz->fillvalue);
  3946. PyTuple_SET_ITEM(lz->ittuple, i, NULL);
  3947. Py_DECREF(it);
  3948. }
  3949. }
  3950. }
  3951. olditem = PyTuple_GET_ITEM(result, i);
  3952. PyTuple_SET_ITEM(result, i, item);
  3953. Py_DECREF(olditem);
  3954. }
  3955. // bpo-42536: The GC may have untracked this result tuple. Since we're
  3956. // recycling it, make sure it's tracked again:
  3957. if (!_PyObject_GC_IS_TRACKED(result)) {
  3958. _PyObject_GC_TRACK(result);
  3959. }
  3960. } else {
  3961. result = PyTuple_New(tuplesize);
  3962. if (result == NULL)
  3963. return NULL;
  3964. for (i=0 ; i < tuplesize ; i++) {
  3965. it = PyTuple_GET_ITEM(lz->ittuple, i);
  3966. if (it == NULL) {
  3967. item = Py_NewRef(lz->fillvalue);
  3968. } else {
  3969. item = PyIter_Next(it);
  3970. if (item == NULL) {
  3971. lz->numactive -= 1;
  3972. if (lz->numactive == 0 || PyErr_Occurred()) {
  3973. lz->numactive = 0;
  3974. Py_DECREF(result);
  3975. return NULL;
  3976. } else {
  3977. item = Py_NewRef(lz->fillvalue);
  3978. PyTuple_SET_ITEM(lz->ittuple, i, NULL);
  3979. Py_DECREF(it);
  3980. }
  3981. }
  3982. }
  3983. PyTuple_SET_ITEM(result, i, item);
  3984. }
  3985. }
  3986. return result;
  3987. }
  3988. static PyObject *
  3989. zip_longest_reduce(ziplongestobject *lz, PyObject *Py_UNUSED(ignored))
  3990. {
  3991. ITERTOOL_PICKLE_DEPRECATION;
  3992. /* Create a new tuple with empty sequences where appropriate to pickle.
  3993. * Then use setstate to set the fillvalue
  3994. */
  3995. int i;
  3996. PyObject *args = PyTuple_New(PyTuple_GET_SIZE(lz->ittuple));
  3997. if (args == NULL)
  3998. return NULL;
  3999. for (i=0; i<PyTuple_GET_SIZE(lz->ittuple); i++) {
  4000. PyObject *elem = PyTuple_GET_ITEM(lz->ittuple, i);
  4001. if (elem == NULL) {
  4002. elem = PyTuple_New(0);
  4003. if (elem == NULL) {
  4004. Py_DECREF(args);
  4005. return NULL;
  4006. }
  4007. } else
  4008. Py_INCREF(elem);
  4009. PyTuple_SET_ITEM(args, i, elem);
  4010. }
  4011. return Py_BuildValue("ONO", Py_TYPE(lz), args, lz->fillvalue);
  4012. }
  4013. static PyObject *
  4014. zip_longest_setstate(ziplongestobject *lz, PyObject *state)
  4015. {
  4016. ITERTOOL_PICKLE_DEPRECATION;
  4017. Py_INCREF(state);
  4018. Py_XSETREF(lz->fillvalue, state);
  4019. Py_RETURN_NONE;
  4020. }
  4021. static PyMethodDef zip_longest_methods[] = {
  4022. {"__reduce__", (PyCFunction)zip_longest_reduce, METH_NOARGS,
  4023. reduce_doc},
  4024. {"__setstate__", (PyCFunction)zip_longest_setstate, METH_O,
  4025. setstate_doc},
  4026. {NULL, NULL} /* sentinel */
  4027. };
  4028. PyDoc_STRVAR(zip_longest_doc,
  4029. "zip_longest(iter1 [,iter2 [...]], [fillvalue=None]) --> zip_longest object\n\
  4030. \n\
  4031. Return a zip_longest object whose .__next__() method returns a tuple where\n\
  4032. the i-th element comes from the i-th iterable argument. The .__next__()\n\
  4033. method continues until the longest iterable in the argument sequence\n\
  4034. is exhausted and then it raises StopIteration. When the shorter iterables\n\
  4035. are exhausted, the fillvalue is substituted in their place. The fillvalue\n\
  4036. defaults to None or can be specified by a keyword argument.\n\
  4037. ");
  4038. static PyType_Slot ziplongest_slots[] = {
  4039. {Py_tp_dealloc, zip_longest_dealloc},
  4040. {Py_tp_getattro, PyObject_GenericGetAttr},
  4041. {Py_tp_doc, (void *)zip_longest_doc},
  4042. {Py_tp_traverse, zip_longest_traverse},
  4043. {Py_tp_iter, PyObject_SelfIter},
  4044. {Py_tp_iternext, zip_longest_next},
  4045. {Py_tp_methods, zip_longest_methods},
  4046. {Py_tp_new, zip_longest_new},
  4047. {Py_tp_free, PyObject_GC_Del},
  4048. {0, NULL},
  4049. };
  4050. static PyType_Spec ziplongest_spec = {
  4051. .name = "itertools.zip_longest",
  4052. .basicsize = sizeof(ziplongestobject),
  4053. .flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_BASETYPE |
  4054. Py_TPFLAGS_IMMUTABLETYPE),
  4055. .slots = ziplongest_slots,
  4056. };
  4057. /* module level code ********************************************************/
  4058. PyDoc_STRVAR(module_doc,
  4059. "Functional tools for creating and using iterators.\n\
  4060. \n\
  4061. Infinite iterators:\n\
  4062. count(start=0, step=1) --> start, start+step, start+2*step, ...\n\
  4063. cycle(p) --> p0, p1, ... plast, p0, p1, ...\n\
  4064. repeat(elem [,n]) --> elem, elem, elem, ... endlessly or up to n times\n\
  4065. \n\
  4066. Iterators terminating on the shortest input sequence:\n\
  4067. accumulate(p[, func]) --> p0, p0+p1, p0+p1+p2\n\
  4068. batched(p, n) --> [p0, p1, ..., p_n-1], [p_n, p_n+1, ..., p_2n-1], ...\n\
  4069. chain(p, q, ...) --> p0, p1, ... plast, q0, q1, ...\n\
  4070. chain.from_iterable([p, q, ...]) --> p0, p1, ... plast, q0, q1, ...\n\
  4071. compress(data, selectors) --> (d[0] if s[0]), (d[1] if s[1]), ...\n\
  4072. dropwhile(predicate, seq) --> seq[n], seq[n+1], starting when predicate fails\n\
  4073. groupby(iterable[, keyfunc]) --> sub-iterators grouped by value of keyfunc(v)\n\
  4074. filterfalse(predicate, seq) --> elements of seq where predicate(elem) is False\n\
  4075. islice(seq, [start,] stop [, step]) --> elements from\n\
  4076. seq[start:stop:step]\n\
  4077. pairwise(s) --> (s[0],s[1]), (s[1],s[2]), (s[2], s[3]), ...\n\
  4078. starmap(fun, seq) --> fun(*seq[0]), fun(*seq[1]), ...\n\
  4079. tee(it, n=2) --> (it1, it2 , ... itn) splits one iterator into n\n\
  4080. takewhile(predicate, seq) --> seq[0], seq[1], until predicate fails\n\
  4081. zip_longest(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ...\n\
  4082. \n\
  4083. Combinatoric generators:\n\
  4084. product(p, q, ... [repeat=1]) --> cartesian product\n\
  4085. permutations(p[, r])\n\
  4086. combinations(p, r)\n\
  4087. combinations_with_replacement(p, r)\n\
  4088. ");
  4089. static int
  4090. itertoolsmodule_traverse(PyObject *mod, visitproc visit, void *arg)
  4091. {
  4092. itertools_state *state = get_module_state(mod);
  4093. Py_VISIT(state->accumulate_type);
  4094. Py_VISIT(state->batched_type);
  4095. Py_VISIT(state->chain_type);
  4096. Py_VISIT(state->combinations_type);
  4097. Py_VISIT(state->compress_type);
  4098. Py_VISIT(state->count_type);
  4099. Py_VISIT(state->cwr_type);
  4100. Py_VISIT(state->cycle_type);
  4101. Py_VISIT(state->dropwhile_type);
  4102. Py_VISIT(state->filterfalse_type);
  4103. Py_VISIT(state->groupby_type);
  4104. Py_VISIT(state->_grouper_type);
  4105. Py_VISIT(state->islice_type);
  4106. Py_VISIT(state->pairwise_type);
  4107. Py_VISIT(state->permutations_type);
  4108. Py_VISIT(state->product_type);
  4109. Py_VISIT(state->repeat_type);
  4110. Py_VISIT(state->starmap_type);
  4111. Py_VISIT(state->takewhile_type);
  4112. Py_VISIT(state->tee_type);
  4113. Py_VISIT(state->teedataobject_type);
  4114. Py_VISIT(state->ziplongest_type);
  4115. return 0;
  4116. }
  4117. static int
  4118. itertoolsmodule_clear(PyObject *mod)
  4119. {
  4120. itertools_state *state = get_module_state(mod);
  4121. Py_CLEAR(state->accumulate_type);
  4122. Py_CLEAR(state->batched_type);
  4123. Py_CLEAR(state->chain_type);
  4124. Py_CLEAR(state->combinations_type);
  4125. Py_CLEAR(state->compress_type);
  4126. Py_CLEAR(state->count_type);
  4127. Py_CLEAR(state->cwr_type);
  4128. Py_CLEAR(state->cycle_type);
  4129. Py_CLEAR(state->dropwhile_type);
  4130. Py_CLEAR(state->filterfalse_type);
  4131. Py_CLEAR(state->groupby_type);
  4132. Py_CLEAR(state->_grouper_type);
  4133. Py_CLEAR(state->islice_type);
  4134. Py_CLEAR(state->pairwise_type);
  4135. Py_CLEAR(state->permutations_type);
  4136. Py_CLEAR(state->product_type);
  4137. Py_CLEAR(state->repeat_type);
  4138. Py_CLEAR(state->starmap_type);
  4139. Py_CLEAR(state->takewhile_type);
  4140. Py_CLEAR(state->tee_type);
  4141. Py_CLEAR(state->teedataobject_type);
  4142. Py_CLEAR(state->ziplongest_type);
  4143. return 0;
  4144. }
  4145. static void
  4146. itertoolsmodule_free(void *mod)
  4147. {
  4148. (void)itertoolsmodule_clear((PyObject *)mod);
  4149. }
  4150. #define ADD_TYPE(module, type, spec) \
  4151. do { \
  4152. type = (PyTypeObject *)PyType_FromModuleAndSpec(module, spec, NULL); \
  4153. if (type == NULL) { \
  4154. return -1; \
  4155. } \
  4156. if (PyModule_AddType(module, type) < 0) { \
  4157. return -1; \
  4158. } \
  4159. } while (0)
  4160. static int
  4161. itertoolsmodule_exec(PyObject *mod)
  4162. {
  4163. itertools_state *state = get_module_state(mod);
  4164. ADD_TYPE(mod, state->accumulate_type, &accumulate_spec);
  4165. ADD_TYPE(mod, state->batched_type, &batched_spec);
  4166. ADD_TYPE(mod, state->chain_type, &chain_spec);
  4167. ADD_TYPE(mod, state->combinations_type, &combinations_spec);
  4168. ADD_TYPE(mod, state->compress_type, &compress_spec);
  4169. ADD_TYPE(mod, state->count_type, &count_spec);
  4170. ADD_TYPE(mod, state->cwr_type, &cwr_spec);
  4171. ADD_TYPE(mod, state->cycle_type, &cycle_spec);
  4172. ADD_TYPE(mod, state->dropwhile_type, &dropwhile_spec);
  4173. ADD_TYPE(mod, state->filterfalse_type, &filterfalse_spec);
  4174. ADD_TYPE(mod, state->groupby_type, &groupby_spec);
  4175. ADD_TYPE(mod, state->_grouper_type, &_grouper_spec);
  4176. ADD_TYPE(mod, state->islice_type, &islice_spec);
  4177. ADD_TYPE(mod, state->pairwise_type, &pairwise_spec);
  4178. ADD_TYPE(mod, state->permutations_type, &permutations_spec);
  4179. ADD_TYPE(mod, state->product_type, &product_spec);
  4180. ADD_TYPE(mod, state->repeat_type, &repeat_spec);
  4181. ADD_TYPE(mod, state->starmap_type, &starmap_spec);
  4182. ADD_TYPE(mod, state->takewhile_type, &takewhile_spec);
  4183. ADD_TYPE(mod, state->tee_type, &tee_spec);
  4184. ADD_TYPE(mod, state->teedataobject_type, &teedataobject_spec);
  4185. ADD_TYPE(mod, state->ziplongest_type, &ziplongest_spec);
  4186. Py_SET_TYPE(state->teedataobject_type, &PyType_Type);
  4187. return 0;
  4188. }
  4189. static struct PyModuleDef_Slot itertoolsmodule_slots[] = {
  4190. {Py_mod_exec, itertoolsmodule_exec},
  4191. {Py_mod_multiple_interpreters, Py_MOD_PER_INTERPRETER_GIL_SUPPORTED},
  4192. {0, NULL}
  4193. };
  4194. static PyMethodDef module_methods[] = {
  4195. ITERTOOLS_TEE_METHODDEF
  4196. {NULL, NULL} /* sentinel */
  4197. };
  4198. static struct PyModuleDef itertoolsmodule = {
  4199. .m_base = PyModuleDef_HEAD_INIT,
  4200. .m_name = "itertools",
  4201. .m_doc = module_doc,
  4202. .m_size = sizeof(itertools_state),
  4203. .m_methods = module_methods,
  4204. .m_slots = itertoolsmodule_slots,
  4205. .m_traverse = itertoolsmodule_traverse,
  4206. .m_clear = itertoolsmodule_clear,
  4207. .m_free = itertoolsmodule_free,
  4208. };
  4209. PyMODINIT_FUNC
  4210. PyInit_itertools(void)
  4211. {
  4212. return PyModuleDef_Init(&itertoolsmodule);
  4213. }