test_more.py 198 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718818919019119219319419519619719819920020120220320420520620720820921021121221321421521621721821922022122222322422522622722822923023123223323423523623723823924024124224324424524624724824925025125225325425525625725825926026126226326426526626726826927027127227327427527627727827928028128228328428528628728828929029129229329429529629729829930030130230330430530630730830931031131231331431531631731831932032132232332432532632732832933033133233333433533633733833934034134234334434534634734834935035135235335435535635735835936036136236336436536636736836937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140240340440540640740840941041141241341441541641741841942042142242342442542642742842943043143243343443543643743843944044144244344444544644744844945045145245345445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751851952052152252352452552652752852953053153253353453553653753853954054154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258358458558658758858959059159259359459559659759859960060160260360460560660760860961061161261361461561661761861962062162262362462562662762862963063163263363463563663763863964064164264364464564664764864965065165265365465565665765865966066166266366466566666766866967067167267367467567667767867968068168268368468568668768868969069169269369469569669769869970070170270370470570670770870971071171271371471571671771871972072172272372472572672772872973073173273373473573673773873974074174274374474574674774874975075175275375475575675775875976076176276376476576676776876977077177277377477577677777877978078178278378478578678778878979079179279379479579679779879980080180280380480580680780880981081181281381481581681781881982082182282382482582682782882983083183283383483583683783883984084184284384484584684784884985085185285385485585685785885986086186286386486586686786886987087187287387487587687787887988088188288388488588688788888989089189289389489589689789889990090190290390490590690790890991091191291391491591691791891992092192292392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001100210031004100510061007100810091010101110121013101410151016101710181019102010211022102310241025102610271028102910301031103210331034103510361037103810391040104110421043104410451046104710481049105010511052105310541055105610571058105910601061106210631064106510661067106810691070107110721073107410751076107710781079108010811082108310841085108610871088108910901091109210931094109510961097109810991100110111021103110411051106110711081109111011111112111311141115111611171118111911201121112211231124112511261127112811291130113111321133113411351136113711381139114011411142114311441145114611471148114911501151115211531154115511561157115811591160116111621163116411651166116711681169117011711172117311741175117611771178117911801181118211831184118511861187118811891190119111921193119411951196119711981199120012011202120312041205120612071208120912101211121212131214121512161217121812191220122112221223122412251226122712281229123012311232123312341235123612371238123912401241124212431244124512461247124812491250125112521253125412551256125712581259126012611262126312641265126612671268126912701271127212731274127512761277127812791280128112821283128412851286128712881289129012911292129312941295129612971298129913001301130213031304130513061307130813091310131113121313131413151316131713181319132013211322132313241325132613271328132913301331133213331334133513361337133813391340134113421343134413451346134713481349135013511352135313541355135613571358135913601361136213631364136513661367136813691370137113721373137413751376137713781379138013811382138313841385138613871388138913901391139213931394139513961397139813991400140114021403140414051406140714081409141014111412141314141415141614171418141914201421142214231424142514261427142814291430143114321433143414351436143714381439144014411442144314441445144614471448144914501451145214531454145514561457145814591460146114621463146414651466146714681469147014711472147314741475147614771478147914801481148214831484148514861487148814891490149114921493149414951496149714981499150015011502150315041505150615071508150915101511151215131514151515161517151815191520152115221523152415251526152715281529153015311532153315341535153615371538153915401541154215431544154515461547154815491550155115521553155415551556155715581559156015611562156315641565156615671568156915701571157215731574157515761577157815791580158115821583158415851586158715881589159015911592159315941595159615971598159916001601160216031604160516061607160816091610161116121613161416151616161716181619162016211622162316241625162616271628162916301631163216331634163516361637163816391640164116421643164416451646164716481649165016511652165316541655165616571658165916601661166216631664166516661667166816691670167116721673167416751676167716781679168016811682168316841685168616871688168916901691169216931694169516961697169816991700170117021703170417051706170717081709171017111712171317141715171617171718171917201721172217231724172517261727172817291730173117321733173417351736173717381739174017411742174317441745174617471748174917501751175217531754175517561757175817591760176117621763176417651766176717681769177017711772177317741775177617771778177917801781178217831784178517861787178817891790179117921793179417951796179717981799180018011802180318041805180618071808180918101811181218131814181518161817181818191820182118221823182418251826182718281829183018311832183318341835183618371838183918401841184218431844184518461847184818491850185118521853185418551856185718581859186018611862186318641865186618671868186918701871187218731874187518761877187818791880188118821883188418851886188718881889189018911892189318941895189618971898189919001901190219031904190519061907190819091910191119121913191419151916191719181919192019211922192319241925192619271928192919301931193219331934193519361937193819391940194119421943194419451946194719481949195019511952195319541955195619571958195919601961196219631964196519661967196819691970197119721973197419751976197719781979198019811982198319841985198619871988198919901991199219931994199519961997199819992000200120022003200420052006200720082009201020112012201320142015201620172018201920202021202220232024202520262027202820292030203120322033203420352036203720382039204020412042204320442045204620472048204920502051205220532054205520562057205820592060206120622063206420652066206720682069207020712072207320742075207620772078207920802081208220832084208520862087208820892090209120922093209420952096209720982099210021012102210321042105210621072108210921102111211221132114211521162117211821192120212121222123212421252126212721282129213021312132213321342135213621372138213921402141214221432144214521462147214821492150215121522153215421552156215721582159216021612162216321642165216621672168216921702171217221732174217521762177217821792180218121822183218421852186218721882189219021912192219321942195219621972198219922002201220222032204220522062207220822092210221122122213221422152216221722182219222022212222222322242225222622272228222922302231223222332234223522362237223822392240224122422243224422452246224722482249225022512252225322542255225622572258225922602261226222632264226522662267226822692270227122722273227422752276227722782279228022812282228322842285228622872288228922902291229222932294229522962297229822992300230123022303230423052306230723082309231023112312231323142315231623172318231923202321232223232324232523262327232823292330233123322333233423352336233723382339234023412342234323442345234623472348234923502351235223532354235523562357235823592360236123622363236423652366236723682369237023712372237323742375237623772378237923802381238223832384238523862387238823892390239123922393239423952396239723982399240024012402240324042405240624072408240924102411241224132414241524162417241824192420242124222423242424252426242724282429243024312432243324342435243624372438243924402441244224432444244524462447244824492450245124522453245424552456245724582459246024612462246324642465246624672468246924702471247224732474247524762477247824792480248124822483248424852486248724882489249024912492249324942495249624972498249925002501250225032504250525062507250825092510251125122513251425152516251725182519252025212522252325242525252625272528252925302531253225332534253525362537253825392540254125422543254425452546254725482549255025512552255325542555255625572558255925602561256225632564256525662567256825692570257125722573257425752576257725782579258025812582258325842585258625872588258925902591259225932594259525962597259825992600260126022603260426052606260726082609261026112612261326142615261626172618261926202621262226232624262526262627262826292630263126322633263426352636263726382639264026412642264326442645264626472648264926502651265226532654265526562657265826592660266126622663266426652666266726682669267026712672267326742675267626772678267926802681268226832684268526862687268826892690269126922693269426952696269726982699270027012702270327042705270627072708270927102711271227132714271527162717271827192720272127222723272427252726272727282729273027312732273327342735273627372738273927402741274227432744274527462747274827492750275127522753275427552756275727582759276027612762276327642765276627672768276927702771277227732774277527762777277827792780278127822783278427852786278727882789279027912792279327942795279627972798279928002801280228032804280528062807280828092810281128122813281428152816281728182819282028212822282328242825282628272828282928302831283228332834283528362837283828392840284128422843284428452846284728482849285028512852285328542855285628572858285928602861286228632864286528662867286828692870287128722873287428752876287728782879288028812882288328842885288628872888288928902891289228932894289528962897289828992900290129022903290429052906290729082909291029112912291329142915291629172918291929202921292229232924292529262927292829292930293129322933293429352936293729382939294029412942294329442945294629472948294929502951295229532954295529562957295829592960296129622963296429652966296729682969297029712972297329742975297629772978297929802981298229832984298529862987298829892990299129922993299429952996299729982999300030013002300330043005300630073008300930103011301230133014301530163017301830193020302130223023302430253026302730283029303030313032303330343035303630373038303930403041304230433044304530463047304830493050305130523053305430553056305730583059306030613062306330643065306630673068306930703071307230733074307530763077307830793080308130823083308430853086308730883089309030913092309330943095309630973098309931003101310231033104310531063107310831093110311131123113311431153116311731183119312031213122312331243125312631273128312931303131313231333134313531363137313831393140314131423143314431453146314731483149315031513152315331543155315631573158315931603161316231633164316531663167316831693170317131723173317431753176317731783179318031813182318331843185318631873188318931903191319231933194319531963197319831993200320132023203320432053206320732083209321032113212321332143215321632173218321932203221322232233224322532263227322832293230323132323233323432353236323732383239324032413242324332443245324632473248324932503251325232533254325532563257325832593260326132623263326432653266326732683269327032713272327332743275327632773278327932803281328232833284328532863287328832893290329132923293329432953296329732983299330033013302330333043305330633073308330933103311331233133314331533163317331833193320332133223323332433253326332733283329333033313332333333343335333633373338333933403341334233433344334533463347334833493350335133523353335433553356335733583359336033613362336333643365336633673368336933703371337233733374337533763377337833793380338133823383338433853386338733883389339033913392339333943395339633973398339934003401340234033404340534063407340834093410341134123413341434153416341734183419342034213422342334243425342634273428342934303431343234333434343534363437343834393440344134423443344434453446344734483449345034513452345334543455345634573458345934603461346234633464346534663467346834693470347134723473347434753476347734783479348034813482348334843485348634873488348934903491349234933494349534963497349834993500350135023503350435053506350735083509351035113512351335143515351635173518351935203521352235233524352535263527352835293530353135323533353435353536353735383539354035413542354335443545354635473548354935503551355235533554355535563557355835593560356135623563356435653566356735683569357035713572357335743575357635773578357935803581358235833584358535863587358835893590359135923593359435953596359735983599360036013602360336043605360636073608360936103611361236133614361536163617361836193620362136223623362436253626362736283629363036313632363336343635363636373638363936403641364236433644364536463647364836493650365136523653365436553656365736583659366036613662366336643665366636673668366936703671367236733674367536763677367836793680368136823683368436853686368736883689369036913692369336943695369636973698369937003701370237033704370537063707370837093710371137123713371437153716371737183719372037213722372337243725372637273728372937303731373237333734373537363737373837393740374137423743374437453746374737483749375037513752375337543755375637573758375937603761376237633764376537663767376837693770377137723773377437753776377737783779378037813782378337843785378637873788378937903791379237933794379537963797379837993800380138023803380438053806380738083809381038113812381338143815381638173818381938203821382238233824382538263827382838293830383138323833383438353836383738383839384038413842384338443845384638473848384938503851385238533854385538563857385838593860386138623863386438653866386738683869387038713872387338743875387638773878387938803881388238833884388538863887388838893890389138923893389438953896389738983899390039013902390339043905390639073908390939103911391239133914391539163917391839193920392139223923392439253926392739283929393039313932393339343935393639373938393939403941394239433944394539463947394839493950395139523953395439553956395739583959396039613962396339643965396639673968396939703971397239733974397539763977397839793980398139823983398439853986398739883989399039913992399339943995399639973998399940004001400240034004400540064007400840094010401140124013401440154016401740184019402040214022402340244025402640274028402940304031403240334034403540364037403840394040404140424043404440454046404740484049405040514052405340544055405640574058405940604061406240634064406540664067406840694070407140724073407440754076407740784079408040814082408340844085408640874088408940904091409240934094409540964097409840994100410141024103410441054106410741084109411041114112411341144115411641174118411941204121412241234124412541264127412841294130413141324133413441354136413741384139414041414142414341444145414641474148414941504151415241534154415541564157415841594160416141624163416441654166416741684169417041714172417341744175417641774178417941804181418241834184418541864187418841894190419141924193419441954196419741984199420042014202420342044205420642074208420942104211421242134214421542164217421842194220422142224223422442254226422742284229423042314232423342344235423642374238423942404241424242434244424542464247424842494250425142524253425442554256425742584259426042614262426342644265426642674268426942704271427242734274427542764277427842794280428142824283428442854286428742884289429042914292429342944295429642974298429943004301430243034304430543064307430843094310431143124313431443154316431743184319432043214322432343244325432643274328432943304331433243334334433543364337433843394340434143424343434443454346434743484349435043514352435343544355435643574358435943604361436243634364436543664367436843694370437143724373437443754376437743784379438043814382438343844385438643874388438943904391439243934394439543964397439843994400440144024403440444054406440744084409441044114412441344144415441644174418441944204421442244234424442544264427442844294430443144324433443444354436443744384439444044414442444344444445444644474448444944504451445244534454445544564457445844594460446144624463446444654466446744684469447044714472447344744475447644774478447944804481448244834484448544864487448844894490449144924493449444954496449744984499450045014502450345044505450645074508450945104511451245134514451545164517451845194520452145224523452445254526452745284529453045314532453345344535453645374538453945404541454245434544454545464547454845494550455145524553455445554556455745584559456045614562456345644565456645674568456945704571457245734574457545764577457845794580458145824583458445854586458745884589459045914592459345944595459645974598459946004601460246034604460546064607460846094610461146124613461446154616461746184619462046214622462346244625462646274628462946304631463246334634463546364637463846394640464146424643464446454646464746484649465046514652465346544655465646574658465946604661466246634664466546664667466846694670467146724673467446754676467746784679468046814682468346844685468646874688468946904691469246934694469546964697469846994700470147024703470447054706470747084709471047114712471347144715471647174718471947204721472247234724472547264727472847294730473147324733473447354736473747384739474047414742474347444745474647474748474947504751475247534754475547564757475847594760476147624763476447654766476747684769477047714772477347744775477647774778477947804781478247834784478547864787478847894790479147924793479447954796479747984799480048014802480348044805480648074808480948104811481248134814481548164817481848194820482148224823482448254826482748284829483048314832483348344835483648374838483948404841484248434844484548464847484848494850485148524853485448554856485748584859486048614862486348644865486648674868486948704871487248734874487548764877487848794880488148824883488448854886488748884889489048914892489348944895489648974898489949004901490249034904490549064907490849094910491149124913491449154916491749184919492049214922492349244925492649274928492949304931493249334934493549364937493849394940494149424943494449454946494749484949495049514952495349544955495649574958495949604961496249634964496549664967496849694970497149724973497449754976497749784979498049814982498349844985498649874988498949904991499249934994499549964997499849995000500150025003500450055006500750085009501050115012501350145015501650175018501950205021502250235024502550265027502850295030503150325033503450355036503750385039504050415042504350445045504650475048504950505051505250535054505550565057505850595060506150625063506450655066506750685069507050715072507350745075507650775078507950805081508250835084508550865087508850895090509150925093509450955096509750985099510051015102510351045105510651075108510951105111511251135114511551165117511851195120512151225123512451255126512751285129513051315132513351345135513651375138513951405141514251435144514551465147514851495150515151525153515451555156515751585159516051615162516351645165516651675168516951705171517251735174517551765177517851795180518151825183518451855186518751885189519051915192519351945195519651975198519952005201520252035204520552065207520852095210521152125213521452155216521752185219522052215222522352245225522652275228522952305231523252335234523552365237523852395240524152425243524452455246524752485249525052515252525352545255525652575258525952605261526252635264526552665267526852695270527152725273527452755276527752785279528052815282528352845285528652875288528952905291529252935294529552965297529852995300530153025303530453055306530753085309531053115312531353145315531653175318531953205321532253235324532553265327532853295330533153325333533453355336533753385339534053415342534353445345534653475348534953505351535253535354535553565357535853595360536153625363536453655366536753685369537053715372537353745375537653775378537953805381538253835384538553865387538853895390539153925393539453955396539753985399540054015402540354045405540654075408540954105411541254135414541554165417541854195420542154225423542454255426542754285429543054315432543354345435543654375438543954405441544254435444544554465447544854495450545154525453545454555456545754585459546054615462546354645465546654675468546954705471547254735474547554765477547854795480548154825483548454855486548754885489549054915492549354945495549654975498549955005501550255035504550555065507550855095510551155125513551455155516551755185519552055215522552355245525552655275528552955305531553255335534553555365537553855395540554155425543554455455546554755485549555055515552555355545555555655575558555955605561556255635564556555665567556855695570557155725573557455755576557755785579558055815582558355845585558655875588558955905591559255935594559555965597559855995600560156025603560456055606560756085609561056115612561356145615561656175618561956205621562256235624562556265627562856295630563156325633563456355636563756385639564056415642564356445645564656475648564956505651565256535654565556565657565856595660566156625663566456655666566756685669567056715672567356745675567656775678567956805681568256835684568556865687568856895690569156925693569456955696569756985699570057015702570357045705570657075708570957105711571257135714571557165717571857195720572157225723572457255726572757285729573057315732573357345735573657375738573957405741574257435744574557465747574857495750575157525753575457555756575757585759576057615762576357645765576657675768576957705771577257735774577557765777577857795780578157825783578457855786578757885789579057915792579357945795579657975798579958005801580258035804580558065807580858095810581158125813581458155816581758185819582058215822582358245825582658275828582958305831583258335834583558365837583858395840584158425843584458455846584758485849585058515852585358545855585658575858585958605861586258635864586558665867586858695870587158725873587458755876587758785879588058815882588358845885588658875888588958905891589258935894589558965897589858995900590159025903590459055906590759085909591059115912591359145915591659175918591959205921592259235924592559265927592859295930593159325933593459355936593759385939594059415942594359445945594659475948594959505951595259535954595559565957595859595960596159625963596459655966596759685969597059715972
  1. import cmath
  2. import warnings
  3. from collections import Counter, abc
  4. from collections.abc import Set
  5. from datetime import datetime, timedelta
  6. from decimal import Decimal
  7. from doctest import DocTestSuite
  8. from fractions import Fraction
  9. from functools import partial, reduce
  10. from io import StringIO
  11. from itertools import (
  12. accumulate,
  13. chain,
  14. combinations,
  15. combinations_with_replacement,
  16. count,
  17. cycle,
  18. groupby,
  19. islice,
  20. permutations,
  21. product,
  22. repeat,
  23. )
  24. from operator import add, mul, itemgetter
  25. from pickle import loads, dumps
  26. from random import Random, random, randrange, seed
  27. from statistics import mean
  28. from string import ascii_letters
  29. from sys import version_info
  30. from time import sleep
  31. from traceback import format_exc
  32. from unittest import skipIf, TestCase
  33. import more_itertools as mi
  34. def load_tests(loader, tests, ignore):
  35. # Add the doctests
  36. tests.addTests(DocTestSuite('more_itertools.more'))
  37. return tests
  38. class ChunkedTests(TestCase):
  39. """Tests for ``chunked()``"""
  40. def test_even(self):
  41. """Test when ``n`` divides evenly into the length of the iterable."""
  42. self.assertEqual(
  43. list(mi.chunked('ABCDEF', 3)), [['A', 'B', 'C'], ['D', 'E', 'F']]
  44. )
  45. def test_odd(self):
  46. """Test when ``n`` does not divide evenly into the length of the
  47. iterable.
  48. """
  49. self.assertEqual(
  50. list(mi.chunked('ABCDE', 3)), [['A', 'B', 'C'], ['D', 'E']]
  51. )
  52. def test_none(self):
  53. """Test when ``n`` has the value ``None``."""
  54. self.assertEqual(
  55. list(mi.chunked('ABCDE', None)), [['A', 'B', 'C', 'D', 'E']]
  56. )
  57. def test_strict_false(self):
  58. """Test when ``n`` does not divide evenly into the length of the
  59. iterable and strict is false.
  60. """
  61. self.assertEqual(
  62. list(mi.chunked('ABCDE', 3, strict=False)),
  63. [['A', 'B', 'C'], ['D', 'E']],
  64. )
  65. def test_strict_being_true(self):
  66. """Test when ``n`` does not divide evenly into the length of the
  67. iterable and strict is True (raising an exception).
  68. """
  69. def f():
  70. return list(mi.chunked('ABCDE', 3, strict=True))
  71. self.assertRaisesRegex(ValueError, "iterable is not divisible by n", f)
  72. self.assertEqual(
  73. list(mi.chunked('ABCDEF', 3, strict=True)),
  74. [['A', 'B', 'C'], ['D', 'E', 'F']],
  75. )
  76. def test_strict_being_true_with_size_none(self):
  77. """Test when ``n`` has value ``None`` and the keyword strict is True
  78. (raising an exception).
  79. """
  80. def f():
  81. return list(mi.chunked('ABCDE', None, strict=True))
  82. self.assertRaisesRegex(
  83. ValueError, "n must not be None when using strict mode.", f
  84. )
  85. class FirstTests(TestCase):
  86. def test_many(self):
  87. # Also try it on a generator expression to make sure it works on
  88. # whatever those return, across Python versions.
  89. self.assertEqual(mi.first(x for x in range(4)), 0)
  90. def test_one(self):
  91. self.assertEqual(mi.first([3]), 3)
  92. def test_empty(self):
  93. with self.assertRaises(ValueError):
  94. mi.first([])
  95. def test_default(self):
  96. self.assertEqual(mi.first([], 'boo'), 'boo')
  97. class IterOnlyRange:
  98. """User-defined iterable class which only support __iter__.
  99. >>> r = IterOnlyRange(5)
  100. >>> r[0] # doctest: +SKIP
  101. AttributeError: IterOnlyRange instance has no attribute '__getitem__'
  102. Note: In Python 3, ``TypeError`` will be raised because ``object`` is
  103. inherited implicitly by default.
  104. >>> r[0] # doctest: +SKIP
  105. TypeError: 'IterOnlyRange' object does not support indexing
  106. """
  107. def __init__(self, n):
  108. """Set the length of the range."""
  109. self.n = n
  110. def __iter__(self):
  111. """Works same as range()."""
  112. return iter(range(self.n))
  113. class LastTests(TestCase):
  114. def test_basic(self):
  115. cases = [
  116. (range(4), 3),
  117. (iter(range(4)), 3),
  118. (range(1), 0),
  119. (iter(range(1)), 0),
  120. (IterOnlyRange(5), 4),
  121. ({n: str(n) for n in range(5)}, 4),
  122. ({0: '0', -1: '-1', 2: '-2'}, 2),
  123. ]
  124. for iterable, expected in cases:
  125. with self.subTest(iterable=iterable):
  126. self.assertEqual(mi.last(iterable), expected)
  127. def test_default(self):
  128. for iterable, default, expected in [
  129. (range(1), None, 0),
  130. ([], None, None),
  131. ({}, None, None),
  132. (iter([]), None, None),
  133. ]:
  134. with self.subTest(args=(iterable, default)):
  135. self.assertEqual(mi.last(iterable, default=default), expected)
  136. def test_empty(self):
  137. for iterable in ([], iter(range(0))):
  138. with self.subTest(iterable=iterable):
  139. with self.assertRaises(ValueError):
  140. mi.last(iterable)
  141. class NthOrLastTests(TestCase):
  142. """Tests for ``nth_or_last()``"""
  143. def test_basic(self):
  144. self.assertEqual(mi.nth_or_last(range(3), 1), 1)
  145. self.assertEqual(mi.nth_or_last(range(3), 3), 2)
  146. def test_default_value(self):
  147. default = 42
  148. self.assertEqual(mi.nth_or_last(range(0), 3, default), default)
  149. def test_empty_iterable_no_default(self):
  150. self.assertRaises(ValueError, lambda: mi.nth_or_last(range(0), 0))
  151. class PeekableMixinTests:
  152. """Common tests for ``peekable()`` and ``seekable()`` behavior"""
  153. cls = None
  154. def test_passthrough(self):
  155. """Iterating a peekable without using ``peek()`` or ``prepend()``
  156. should just give the underlying iterable's elements (a trivial test but
  157. useful to set a baseline in case something goes wrong)"""
  158. expected = [1, 2, 3, 4, 5]
  159. actual = list(self.cls(expected))
  160. self.assertEqual(actual, expected)
  161. def test_peek_default(self):
  162. """Make sure passing a default into ``peek()`` works."""
  163. p = self.cls([])
  164. self.assertEqual(p.peek(7), 7)
  165. def test_truthiness(self):
  166. """Make sure a ``peekable`` tests true iff there are items remaining in
  167. the iterable.
  168. """
  169. p = self.cls([])
  170. self.assertFalse(p)
  171. p = self.cls(range(3))
  172. self.assertTrue(p)
  173. def test_simple_peeking(self):
  174. """Make sure ``next`` and ``peek`` advance and don't advance the
  175. iterator, respectively.
  176. """
  177. p = self.cls(range(10))
  178. self.assertEqual(next(p), 0)
  179. self.assertEqual(p.peek(), 1)
  180. self.assertEqual(p.peek(), 1)
  181. self.assertEqual(next(p), 1)
  182. class PeekableTests(PeekableMixinTests, TestCase):
  183. cls = mi.peekable
  184. def test_indexing(self):
  185. """
  186. Indexing into the peekable shouldn't advance the iterator.
  187. """
  188. p = mi.peekable('abcdefghijkl')
  189. # The 0th index is what ``next()`` will return
  190. self.assertEqual(p[0], 'a')
  191. self.assertEqual(next(p), 'a')
  192. # Indexing further into the peekable shouldn't advance the iterator
  193. self.assertEqual(p[2], 'd')
  194. self.assertEqual(next(p), 'b')
  195. # The 0th index moves up with the iterator; the last index follows
  196. self.assertEqual(p[0], 'c')
  197. self.assertEqual(p[9], 'l')
  198. self.assertEqual(next(p), 'c')
  199. self.assertEqual(p[8], 'l')
  200. # Negative indexing should work too
  201. self.assertEqual(p[-2], 'k')
  202. self.assertEqual(p[-9], 'd')
  203. self.assertRaises(IndexError, lambda: p[-10])
  204. def test_slicing(self):
  205. """Slicing the peekable shouldn't advance the iterator."""
  206. seq = list('abcdefghijkl')
  207. p = mi.peekable(seq)
  208. # Slicing the peekable should just be like slicing a re-iterable
  209. self.assertEqual(p[1:4], seq[1:4])
  210. # Advancing the iterator moves the slices up also
  211. self.assertEqual(next(p), 'a')
  212. self.assertEqual(p[1:4], seq[1:][1:4])
  213. # Implicit starts and stop should work
  214. self.assertEqual(p[:5], seq[1:][:5])
  215. self.assertEqual(p[:], seq[1:][:])
  216. # Indexing past the end should work
  217. self.assertEqual(p[:100], seq[1:][:100])
  218. # Steps should work, including negative
  219. self.assertEqual(p[::2], seq[1:][::2])
  220. self.assertEqual(p[::-1], seq[1:][::-1])
  221. def test_slicing_reset(self):
  222. """Test slicing on a fresh iterable each time"""
  223. iterable = ['0', '1', '2', '3', '4', '5']
  224. indexes = list(range(-4, len(iterable) + 4)) + [None]
  225. steps = [1, 2, 3, 4, -1, -2, -3, 4]
  226. for slice_args in product(indexes, indexes, steps):
  227. it = iter(iterable)
  228. p = mi.peekable(it)
  229. next(p)
  230. index = slice(*slice_args)
  231. actual = p[index]
  232. expected = iterable[1:][index]
  233. self.assertEqual(actual, expected, slice_args)
  234. def test_slicing_error(self):
  235. iterable = '01234567'
  236. p = mi.peekable(iter(iterable))
  237. # Prime the cache
  238. p.peek()
  239. old_cache = list(p._cache)
  240. # Illegal slice
  241. with self.assertRaises(ValueError):
  242. p[1:-1:0]
  243. # Neither the cache nor the iteration should be affected
  244. self.assertEqual(old_cache, list(p._cache))
  245. self.assertEqual(list(p), list(iterable))
  246. # prepend() behavior tests
  247. def test_prepend(self):
  248. """Tests interspersed ``prepend()`` and ``next()`` calls"""
  249. it = mi.peekable(range(2))
  250. actual = []
  251. # Test prepend() before next()
  252. it.prepend(10)
  253. actual += [next(it), next(it)]
  254. # Test prepend() between next()s
  255. it.prepend(11)
  256. actual += [next(it), next(it)]
  257. # Test prepend() after source iterable is consumed
  258. it.prepend(12)
  259. actual += [next(it)]
  260. expected = [10, 0, 11, 1, 12]
  261. self.assertEqual(actual, expected)
  262. def test_multi_prepend(self):
  263. """Tests prepending multiple items and getting them in proper order"""
  264. it = mi.peekable(range(5))
  265. actual = [next(it), next(it)]
  266. it.prepend(10, 11, 12)
  267. it.prepend(20, 21)
  268. actual += list(it)
  269. expected = [0, 1, 20, 21, 10, 11, 12, 2, 3, 4]
  270. self.assertEqual(actual, expected)
  271. def test_empty(self):
  272. """Tests prepending in front of an empty iterable"""
  273. it = mi.peekable([])
  274. it.prepend(10)
  275. actual = list(it)
  276. expected = [10]
  277. self.assertEqual(actual, expected)
  278. def test_prepend_truthiness(self):
  279. """Tests that ``__bool__()`` or ``__nonzero__()`` works properly
  280. with ``prepend()``"""
  281. it = mi.peekable(range(5))
  282. self.assertTrue(it)
  283. actual = list(it)
  284. self.assertFalse(it)
  285. it.prepend(10)
  286. self.assertTrue(it)
  287. actual += [next(it)]
  288. self.assertFalse(it)
  289. expected = [0, 1, 2, 3, 4, 10]
  290. self.assertEqual(actual, expected)
  291. def test_multi_prepend_peek(self):
  292. """Tests prepending multiple elements and getting them in reverse order
  293. while peeking"""
  294. it = mi.peekable(range(5))
  295. actual = [next(it), next(it)]
  296. self.assertEqual(it.peek(), 2)
  297. it.prepend(10, 11, 12)
  298. self.assertEqual(it.peek(), 10)
  299. it.prepend(20, 21)
  300. self.assertEqual(it.peek(), 20)
  301. actual += list(it)
  302. self.assertFalse(it)
  303. expected = [0, 1, 20, 21, 10, 11, 12, 2, 3, 4]
  304. self.assertEqual(actual, expected)
  305. def test_prepend_after_stop(self):
  306. """Test resuming iteration after a previous exhaustion"""
  307. it = mi.peekable(range(3))
  308. self.assertEqual(list(it), [0, 1, 2])
  309. self.assertRaises(StopIteration, lambda: next(it))
  310. it.prepend(10)
  311. self.assertEqual(next(it), 10)
  312. self.assertRaises(StopIteration, lambda: next(it))
  313. def test_prepend_slicing(self):
  314. """Tests interaction between prepending and slicing"""
  315. seq = list(range(20))
  316. p = mi.peekable(seq)
  317. p.prepend(30, 40, 50)
  318. pseq = [30, 40, 50] + seq # pseq for prepended_seq
  319. # adapt the specific tests from test_slicing
  320. self.assertEqual(p[0], 30)
  321. self.assertEqual(p[1:8], pseq[1:8])
  322. self.assertEqual(p[1:], pseq[1:])
  323. self.assertEqual(p[:5], pseq[:5])
  324. self.assertEqual(p[:], pseq[:])
  325. self.assertEqual(p[:100], pseq[:100])
  326. self.assertEqual(p[::2], pseq[::2])
  327. self.assertEqual(p[::-1], pseq[::-1])
  328. def test_prepend_indexing(self):
  329. """Tests interaction between prepending and indexing"""
  330. seq = list(range(20))
  331. p = mi.peekable(seq)
  332. p.prepend(30, 40, 50)
  333. self.assertEqual(p[0], 30)
  334. self.assertEqual(next(p), 30)
  335. self.assertEqual(p[2], 0)
  336. self.assertEqual(next(p), 40)
  337. self.assertEqual(p[0], 50)
  338. self.assertEqual(p[9], 8)
  339. self.assertEqual(next(p), 50)
  340. self.assertEqual(p[8], 8)
  341. self.assertEqual(p[-2], 18)
  342. self.assertEqual(p[-9], 11)
  343. self.assertRaises(IndexError, lambda: p[-21])
  344. def test_prepend_iterable(self):
  345. """Tests prepending from an iterable"""
  346. it = mi.peekable(range(5))
  347. # Don't directly use the range() object to avoid any range-specific
  348. # optimizations
  349. it.prepend(*(x for x in range(5)))
  350. actual = list(it)
  351. expected = list(chain(range(5), range(5)))
  352. self.assertEqual(actual, expected)
  353. def test_prepend_many(self):
  354. """Tests that prepending a huge number of elements works"""
  355. it = mi.peekable(range(5))
  356. # Don't directly use the range() object to avoid any range-specific
  357. # optimizations
  358. it.prepend(*(x for x in range(20000)))
  359. actual = list(it)
  360. expected = list(chain(range(20000), range(5)))
  361. self.assertEqual(actual, expected)
  362. def test_prepend_reversed(self):
  363. """Tests prepending from a reversed iterable"""
  364. it = mi.peekable(range(3))
  365. it.prepend(*reversed((10, 11, 12)))
  366. actual = list(it)
  367. expected = [12, 11, 10, 0, 1, 2]
  368. self.assertEqual(actual, expected)
  369. class ConsumerTests(TestCase):
  370. """Tests for ``consumer()``"""
  371. def test_consumer(self):
  372. @mi.consumer
  373. def eater():
  374. while True:
  375. x = yield # noqa
  376. e = eater()
  377. e.send('hi') # without @consumer, would raise TypeError
  378. class DistinctPermutationsTests(TestCase):
  379. def test_basic(self):
  380. iterable = ['z', 'a', 'a', 'q', 'q', 'q', 'y']
  381. actual = list(mi.distinct_permutations(iterable))
  382. expected = set(permutations(iterable))
  383. self.assertCountEqual(actual, expected)
  384. def test_r(self):
  385. for iterable, r in (
  386. ('mississippi', 0),
  387. ('mississippi', 1),
  388. ('mississippi', 6),
  389. ('mississippi', 7),
  390. ('mississippi', 12),
  391. ([0, 1, 1, 0], 0),
  392. ([0, 1, 1, 0], 1),
  393. ([0, 1, 1, 0], 2),
  394. ([0, 1, 1, 0], 3),
  395. ([0, 1, 1, 0], 4),
  396. (['a'], 0),
  397. (['a'], 1),
  398. (['a'], 5),
  399. ([], 0),
  400. ([], 1),
  401. ([], 4),
  402. ):
  403. with self.subTest(iterable=iterable, r=r):
  404. expected = set(permutations(iterable, r))
  405. actual = list(mi.distinct_permutations(iter(iterable), r))
  406. self.assertCountEqual(actual, expected)
  407. def test_unsortable(self):
  408. iterable = ['1', 2, 2, 3, 3, 3]
  409. actual = list(mi.distinct_permutations(iterable))
  410. expected = set(permutations(iterable))
  411. self.assertCountEqual(actual, expected)
  412. def test_unsortable_r(self):
  413. iterable = ['1', 2, 2, 3, 3, 3]
  414. for r in range(len(iterable) + 1):
  415. with self.subTest(iterable=iterable, r=r):
  416. actual = list(mi.distinct_permutations(iterable, r=r))
  417. expected = set(permutations(iterable, r=r))
  418. self.assertCountEqual(actual, expected)
  419. def test_unsorted_equivalent(self):
  420. iterable = [1, True, '3']
  421. actual = list(mi.distinct_permutations(iterable))
  422. expected = set(permutations(iterable))
  423. self.assertCountEqual(actual, expected)
  424. def test_unhashable(self):
  425. iterable = ([1], [1], 2)
  426. actual = list(mi.distinct_permutations(iterable))
  427. expected = list(mi.unique_everseen(permutations(iterable)))
  428. self.assertCountEqual(actual, expected)
  429. class IlenTests(TestCase):
  430. def test_ilen(self):
  431. """Sanity-checks for ``ilen()``."""
  432. # Non-empty
  433. self.assertEqual(
  434. mi.ilen(filter(lambda x: x % 10 == 0, range(101))), 11
  435. )
  436. # Empty
  437. self.assertEqual(mi.ilen(x for x in range(0)), 0)
  438. # Iterable with __len__
  439. self.assertEqual(mi.ilen(list(range(6))), 6)
  440. class MinMaxTests(TestCase):
  441. def test_basic(self):
  442. for iterable, expected in (
  443. # easy case
  444. ([0, 1, 2, 3], (0, 3)),
  445. # min and max are not in the extremes + we have `int`s and `float`s
  446. ([3, 5.5, -1, 2], (-1, 5.5)),
  447. # unordered collection
  448. ({3, 5.5, -1, 2}, (-1, 5.5)),
  449. # with repetitions
  450. ([3, 5.5, float('-Inf'), 5.5], (float('-Inf'), 5.5)),
  451. # other collections
  452. ('banana', ('a', 'n')),
  453. ({0: 1, 2: 100, 1: 10}, (0, 2)),
  454. (range(3, 14), (3, 13)),
  455. ):
  456. with self.subTest(iterable=iterable, expected=expected):
  457. # check for expected results
  458. self.assertTupleEqual(mi.minmax(iterable), expected)
  459. # check for equality with built-in `min` and `max`
  460. self.assertTupleEqual(
  461. mi.minmax(iterable), (min(iterable), max(iterable))
  462. )
  463. def test_unpacked(self):
  464. self.assertTupleEqual(mi.minmax(2, 3, 1), (1, 3))
  465. self.assertTupleEqual(mi.minmax(12, 3, 4, key=str), (12, 4))
  466. def test_iterables(self):
  467. self.assertTupleEqual(mi.minmax(x for x in [0, 1, 2, 3]), (0, 3))
  468. self.assertTupleEqual(
  469. mi.minmax(map(str, [3, 5.5, 'a', 2])), ('2', 'a')
  470. )
  471. self.assertTupleEqual(
  472. mi.minmax(filter(None, [0, 3, '', None, 10])), (3, 10)
  473. )
  474. def test_key(self):
  475. self.assertTupleEqual(
  476. mi.minmax({(), (1, 4, 2), 'abcde', range(4)}, key=len),
  477. ((), 'abcde'),
  478. )
  479. self.assertTupleEqual(
  480. mi.minmax((x for x in [10, 3, 25]), key=str), (10, 3)
  481. )
  482. def test_default(self):
  483. with self.assertRaises(ValueError):
  484. mi.minmax([])
  485. self.assertIs(mi.minmax([], default=None), None)
  486. self.assertListEqual(mi.minmax([], default=[1, 'a']), [1, 'a'])
  487. class WithIterTests(TestCase):
  488. def test_with_iter(self):
  489. s = StringIO('One fish\nTwo fish')
  490. initial_words = [line.split()[0] for line in mi.with_iter(s)]
  491. # Iterable's items should be faithfully represented
  492. self.assertEqual(initial_words, ['One', 'Two'])
  493. # The file object should be closed
  494. self.assertTrue(s.closed)
  495. class OneTests(TestCase):
  496. def test_basic(self):
  497. it = iter(['item'])
  498. self.assertEqual(mi.one(it), 'item')
  499. def test_too_short(self):
  500. it = iter([])
  501. for too_short, exc_type in [
  502. (None, ValueError),
  503. (IndexError, IndexError),
  504. ]:
  505. with self.subTest(too_short=too_short):
  506. try:
  507. mi.one(it, too_short=too_short)
  508. except exc_type:
  509. formatted_exc = format_exc()
  510. self.assertIn('StopIteration', formatted_exc)
  511. self.assertIn(
  512. 'The above exception was the direct cause',
  513. formatted_exc,
  514. )
  515. else:
  516. self.fail()
  517. def test_too_long(self):
  518. it = count()
  519. self.assertRaises(ValueError, lambda: mi.one(it)) # burn 0 and 1
  520. self.assertEqual(next(it), 2)
  521. self.assertRaises(
  522. OverflowError, lambda: mi.one(it, too_long=OverflowError)
  523. )
  524. def test_too_long_default_message(self):
  525. it = count()
  526. self.assertRaisesRegex(
  527. ValueError,
  528. "Expected exactly one item in "
  529. "iterable, but got 0, 1, and "
  530. "perhaps more.",
  531. lambda: mi.one(it),
  532. )
  533. class IntersperseTest(TestCase):
  534. """Tests for intersperse()"""
  535. def test_even(self):
  536. iterable = (x for x in '01')
  537. self.assertEqual(
  538. list(mi.intersperse(None, iterable)), ['0', None, '1']
  539. )
  540. def test_odd(self):
  541. iterable = (x for x in '012')
  542. self.assertEqual(
  543. list(mi.intersperse(None, iterable)), ['0', None, '1', None, '2']
  544. )
  545. def test_nested(self):
  546. element = ('a', 'b')
  547. iterable = (x for x in '012')
  548. actual = list(mi.intersperse(element, iterable))
  549. expected = ['0', ('a', 'b'), '1', ('a', 'b'), '2']
  550. self.assertEqual(actual, expected)
  551. def test_not_iterable(self):
  552. self.assertRaises(TypeError, lambda: mi.intersperse('x', 1))
  553. def test_n(self):
  554. for n, element, expected in [
  555. (1, '_', ['0', '_', '1', '_', '2', '_', '3', '_', '4', '_', '5']),
  556. (2, '_', ['0', '1', '_', '2', '3', '_', '4', '5']),
  557. (3, '_', ['0', '1', '2', '_', '3', '4', '5']),
  558. (4, '_', ['0', '1', '2', '3', '_', '4', '5']),
  559. (5, '_', ['0', '1', '2', '3', '4', '_', '5']),
  560. (6, '_', ['0', '1', '2', '3', '4', '5']),
  561. (7, '_', ['0', '1', '2', '3', '4', '5']),
  562. (3, ['a', 'b'], ['0', '1', '2', ['a', 'b'], '3', '4', '5']),
  563. ]:
  564. iterable = (x for x in '012345')
  565. actual = list(mi.intersperse(element, iterable, n=n))
  566. self.assertEqual(actual, expected)
  567. def test_n_zero(self):
  568. self.assertRaises(
  569. ValueError, lambda: list(mi.intersperse('x', '012', n=0))
  570. )
  571. class UniqueToEachTests(TestCase):
  572. """Tests for ``unique_to_each()``"""
  573. def test_all_unique(self):
  574. """When all the input iterables are unique the output should match
  575. the input."""
  576. iterables = [[1, 2], [3, 4, 5], [6, 7, 8]]
  577. self.assertEqual(mi.unique_to_each(*iterables), iterables)
  578. def test_duplicates(self):
  579. """When there are duplicates in any of the input iterables that aren't
  580. in the rest, those duplicates should be emitted."""
  581. iterables = ["mississippi", "missouri"]
  582. self.assertEqual(
  583. mi.unique_to_each(*iterables), [['p', 'p'], ['o', 'u', 'r']]
  584. )
  585. def test_mixed(self):
  586. """When the input iterables contain different types the function should
  587. still behave properly"""
  588. iterables = ['x', (i for i in range(3)), [1, 2, 3], tuple()]
  589. self.assertEqual(mi.unique_to_each(*iterables), [['x'], [0], [3], []])
  590. class WindowedTests(TestCase):
  591. def test_basic(self):
  592. iterable = [1, 2, 3, 4, 5]
  593. for n, expected in (
  594. (6, [(1, 2, 3, 4, 5, None)]),
  595. (5, [(1, 2, 3, 4, 5)]),
  596. (4, [(1, 2, 3, 4), (2, 3, 4, 5)]),
  597. (3, [(1, 2, 3), (2, 3, 4), (3, 4, 5)]),
  598. (2, [(1, 2), (2, 3), (3, 4), (4, 5)]),
  599. (1, [(1,), (2,), (3,), (4,), (5,)]),
  600. (0, [()]),
  601. ):
  602. with self.subTest(n=n):
  603. actual = list(mi.windowed(iterable, n))
  604. self.assertEqual(actual, expected)
  605. def test_fillvalue(self):
  606. actual = list(mi.windowed([1, 2, 3, 4, 5], 6, fillvalue='!'))
  607. expected = [(1, 2, 3, 4, 5, '!')]
  608. self.assertEqual(actual, expected)
  609. def test_step(self):
  610. iterable = [1, 2, 3, 4, 5, 6, 7]
  611. for n, step, expected in [
  612. (3, 2, [(1, 2, 3), (3, 4, 5), (5, 6, 7)]), # n > step
  613. (3, 3, [(1, 2, 3), (4, 5, 6), (7, None, None)]), # n == step
  614. (3, 4, [(1, 2, 3), (5, 6, 7)]), # lines up nicely
  615. (3, 5, [(1, 2, 3), (6, 7, None)]), # off by one
  616. (3, 6, [(1, 2, 3), (7, None, None)]), # off by two
  617. (3, 7, [(1, 2, 3)]), # step past the end
  618. (7, 8, [(1, 2, 3, 4, 5, 6, 7)]), # step > len(iterable)
  619. ]:
  620. with self.subTest(n=n, step=step):
  621. actual = list(mi.windowed(iterable, n, step=step))
  622. self.assertEqual(actual, expected)
  623. def test_invalid_step(self):
  624. # Step must be greater than or equal to 1
  625. with self.assertRaises(ValueError):
  626. list(mi.windowed([1, 2, 3, 4, 5], 3, step=0))
  627. def test_fillvalue_step(self):
  628. actual = list(mi.windowed([1, 2, 3, 4, 5], 3, fillvalue='!', step=3))
  629. expected = [(1, 2, 3), (4, 5, '!')]
  630. self.assertEqual(actual, expected)
  631. def test_negative(self):
  632. with self.assertRaises(ValueError):
  633. list(mi.windowed([1, 2, 3, 4, 5], -1))
  634. def test_empty_seq(self):
  635. actual = list(mi.windowed([], 3))
  636. expected = []
  637. self.assertEqual(actual, expected)
  638. class SubstringsTests(TestCase):
  639. def test_basic(self):
  640. iterable = (x for x in range(4))
  641. actual = list(mi.substrings(iterable))
  642. expected = [
  643. (0,),
  644. (1,),
  645. (2,),
  646. (3,),
  647. (0, 1),
  648. (1, 2),
  649. (2, 3),
  650. (0, 1, 2),
  651. (1, 2, 3),
  652. (0, 1, 2, 3),
  653. ]
  654. self.assertEqual(actual, expected)
  655. def test_strings(self):
  656. iterable = 'abc'
  657. actual = list(mi.substrings(iterable))
  658. expected = [
  659. ('a',),
  660. ('b',),
  661. ('c',),
  662. ('a', 'b'),
  663. ('b', 'c'),
  664. ('a', 'b', 'c'),
  665. ]
  666. self.assertEqual(actual, expected)
  667. def test_empty(self):
  668. iterable = iter([])
  669. actual = list(mi.substrings(iterable))
  670. expected = []
  671. self.assertEqual(actual, expected)
  672. def test_order(self):
  673. iterable = [2, 0, 1]
  674. actual = list(mi.substrings(iterable))
  675. expected = [(2,), (0,), (1,), (2, 0), (0, 1), (2, 0, 1)]
  676. self.assertEqual(actual, expected)
  677. class SubstringsIndexesTests(TestCase):
  678. def test_basic(self):
  679. sequence = [x for x in range(4)]
  680. actual = list(mi.substrings_indexes(sequence))
  681. expected = [
  682. ([0], 0, 1),
  683. ([1], 1, 2),
  684. ([2], 2, 3),
  685. ([3], 3, 4),
  686. ([0, 1], 0, 2),
  687. ([1, 2], 1, 3),
  688. ([2, 3], 2, 4),
  689. ([0, 1, 2], 0, 3),
  690. ([1, 2, 3], 1, 4),
  691. ([0, 1, 2, 3], 0, 4),
  692. ]
  693. self.assertEqual(actual, expected)
  694. def test_strings(self):
  695. sequence = 'abc'
  696. actual = list(mi.substrings_indexes(sequence))
  697. expected = [
  698. ('a', 0, 1),
  699. ('b', 1, 2),
  700. ('c', 2, 3),
  701. ('ab', 0, 2),
  702. ('bc', 1, 3),
  703. ('abc', 0, 3),
  704. ]
  705. self.assertEqual(actual, expected)
  706. def test_empty(self):
  707. sequence = []
  708. actual = list(mi.substrings_indexes(sequence))
  709. expected = []
  710. self.assertEqual(actual, expected)
  711. def test_order(self):
  712. sequence = [2, 0, 1]
  713. actual = list(mi.substrings_indexes(sequence))
  714. expected = [
  715. ([2], 0, 1),
  716. ([0], 1, 2),
  717. ([1], 2, 3),
  718. ([2, 0], 0, 2),
  719. ([0, 1], 1, 3),
  720. ([2, 0, 1], 0, 3),
  721. ]
  722. self.assertEqual(actual, expected)
  723. def test_reverse(self):
  724. sequence = [2, 0, 1]
  725. actual = list(mi.substrings_indexes(sequence, reverse=True))
  726. expected = [
  727. ([2, 0, 1], 0, 3),
  728. ([2, 0], 0, 2),
  729. ([0, 1], 1, 3),
  730. ([2], 0, 1),
  731. ([0], 1, 2),
  732. ([1], 2, 3),
  733. ]
  734. self.assertEqual(actual, expected)
  735. class BucketTests(TestCase):
  736. def test_basic(self):
  737. iterable = [10, 20, 30, 11, 21, 31, 12, 22, 23, 33]
  738. D = mi.bucket(iterable, key=lambda x: 10 * (x // 10))
  739. # In-order access
  740. self.assertEqual(list(D[10]), [10, 11, 12])
  741. # Out of order access
  742. self.assertEqual(list(D[30]), [30, 31, 33])
  743. self.assertEqual(list(D[20]), [20, 21, 22, 23])
  744. self.assertEqual(list(D[40]), []) # Nothing in here!
  745. def test_in(self):
  746. iterable = [10, 20, 30, 11, 21, 31, 12, 22, 23, 33]
  747. D = mi.bucket(iterable, key=lambda x: 10 * (x // 10))
  748. self.assertIn(10, D)
  749. self.assertNotIn(40, D)
  750. self.assertIn(20, D)
  751. self.assertNotIn(21, D)
  752. # Checking in-ness shouldn't advance the iterator
  753. self.assertEqual(next(D[10]), 10)
  754. def test_validator(self):
  755. iterable = count(0)
  756. key = lambda x: int(str(x)[0]) # First digit of each number
  757. validator = lambda x: 0 < x < 10 # No leading zeros
  758. D = mi.bucket(iterable, key, validator=validator)
  759. self.assertEqual(mi.take(3, D[1]), [1, 10, 11])
  760. self.assertNotIn(0, D) # Non-valid entries don't return True
  761. self.assertNotIn(0, D._cache) # Don't store non-valid entries
  762. self.assertEqual(list(D[0]), [])
  763. def test_list(self):
  764. iterable = [10, 20, 30, 11, 21, 31, 12, 22, 23, 33]
  765. D = mi.bucket(iterable, key=lambda x: 10 * (x // 10))
  766. self.assertEqual(list(D[10]), [10, 11, 12])
  767. self.assertEqual(list(D[20]), [20, 21, 22, 23])
  768. self.assertEqual(list(D[30]), [30, 31, 33])
  769. self.assertEqual(set(D), {10, 20, 30})
  770. def test_list_validator(self):
  771. iterable = [10, 20, 30, 11, 21, 31, 12, 22, 23, 33]
  772. key = lambda x: 10 * (x // 10)
  773. validator = lambda x: x != 20
  774. D = mi.bucket(iterable, key, validator=validator)
  775. self.assertEqual(set(D), {10, 30})
  776. self.assertEqual(list(D[10]), [10, 11, 12])
  777. self.assertEqual(list(D[20]), [])
  778. self.assertEqual(list(D[30]), [30, 31, 33])
  779. class SpyTests(TestCase):
  780. """Tests for ``spy()``"""
  781. def test_basic(self):
  782. original_iterable = iter('abcdefg')
  783. head, new_iterable = mi.spy(original_iterable)
  784. self.assertEqual(head, ['a'])
  785. self.assertEqual(
  786. list(new_iterable), ['a', 'b', 'c', 'd', 'e', 'f', 'g']
  787. )
  788. def test_unpacking(self):
  789. original_iterable = iter('abcdefg')
  790. (first, second, third), new_iterable = mi.spy(original_iterable, 3)
  791. self.assertEqual(first, 'a')
  792. self.assertEqual(second, 'b')
  793. self.assertEqual(third, 'c')
  794. self.assertEqual(
  795. list(new_iterable), ['a', 'b', 'c', 'd', 'e', 'f', 'g']
  796. )
  797. def test_too_many(self):
  798. original_iterable = iter('abc')
  799. head, new_iterable = mi.spy(original_iterable, 4)
  800. self.assertEqual(head, ['a', 'b', 'c'])
  801. self.assertEqual(list(new_iterable), ['a', 'b', 'c'])
  802. def test_zero(self):
  803. original_iterable = iter('abc')
  804. head, new_iterable = mi.spy(original_iterable, 0)
  805. self.assertEqual(head, [])
  806. self.assertEqual(list(new_iterable), ['a', 'b', 'c'])
  807. def test_immutable(self):
  808. original_iterable = iter('abcdefg')
  809. head, new_iterable = mi.spy(original_iterable, 3)
  810. head[0] = 'A'
  811. self.assertEqual(head, ['A', 'b', 'c'])
  812. self.assertEqual(
  813. list(new_iterable), ['a', 'b', 'c', 'd', 'e', 'f', 'g']
  814. )
  815. class InterleaveTests(TestCase):
  816. def test_even(self):
  817. actual = list(mi.interleave([1, 4, 7], [2, 5, 8], [3, 6, 9]))
  818. expected = [1, 2, 3, 4, 5, 6, 7, 8, 9]
  819. self.assertEqual(actual, expected)
  820. def test_short(self):
  821. actual = list(mi.interleave([1, 4], [2, 5, 7], [3, 6, 8]))
  822. expected = [1, 2, 3, 4, 5, 6]
  823. self.assertEqual(actual, expected)
  824. def test_mixed_types(self):
  825. it_list = ['a', 'b', 'c', 'd']
  826. it_str = '12345'
  827. it_inf = count()
  828. actual = list(mi.interleave(it_list, it_str, it_inf))
  829. expected = ['a', '1', 0, 'b', '2', 1, 'c', '3', 2, 'd', '4', 3]
  830. self.assertEqual(actual, expected)
  831. class InterleaveLongestTests(TestCase):
  832. def test_even(self):
  833. actual = list(mi.interleave_longest([1, 4, 7], [2, 5, 8], [3, 6, 9]))
  834. expected = [1, 2, 3, 4, 5, 6, 7, 8, 9]
  835. self.assertEqual(actual, expected)
  836. def test_short(self):
  837. actual = list(mi.interleave_longest([1, 4], [2, 5, 7], [3, 6, 8]))
  838. expected = [1, 2, 3, 4, 5, 6, 7, 8]
  839. self.assertEqual(actual, expected)
  840. def test_mixed_types(self):
  841. it_list = ['a', 'b', 'c', 'd']
  842. it_str = '12345'
  843. it_gen = (x for x in range(3))
  844. actual = list(mi.interleave_longest(it_list, it_str, it_gen))
  845. expected = ['a', '1', 0, 'b', '2', 1, 'c', '3', 2, 'd', '4', '5']
  846. self.assertEqual(actual, expected)
  847. class InterleaveEvenlyTests(TestCase):
  848. def test_equal_lengths(self):
  849. # when lengths are equal, the relative order shouldn't change
  850. a = [1, 2, 3]
  851. b = [5, 6, 7]
  852. actual = list(mi.interleave_evenly([a, b]))
  853. expected = [1, 5, 2, 6, 3, 7]
  854. self.assertEqual(actual, expected)
  855. def test_proportional(self):
  856. # easy case where the iterables have proportional length
  857. a = [1, 2, 3, 4]
  858. b = [5, 6]
  859. actual = list(mi.interleave_evenly([a, b]))
  860. expected = [1, 2, 5, 3, 4, 6]
  861. self.assertEqual(actual, expected)
  862. # swapping a and b should yield the same result
  863. actual_swapped = list(mi.interleave_evenly([b, a]))
  864. self.assertEqual(actual_swapped, expected)
  865. def test_not_proportional(self):
  866. a = [1, 2, 3, 4, 5, 6, 7]
  867. b = [8, 9, 10]
  868. expected = [1, 2, 8, 3, 4, 9, 5, 6, 10, 7]
  869. actual = list(mi.interleave_evenly([a, b]))
  870. self.assertEqual(actual, expected)
  871. def test_degenerate_one(self):
  872. a = [0, 1, 2, 3, 4]
  873. b = [5]
  874. expected = [0, 1, 2, 5, 3, 4]
  875. actual = list(mi.interleave_evenly([a, b]))
  876. self.assertEqual(actual, expected)
  877. def test_degenerate_empty(self):
  878. a = [1, 2, 3]
  879. b = []
  880. expected = [1, 2, 3]
  881. actual = list(mi.interleave_evenly([a, b]))
  882. self.assertEqual(actual, expected)
  883. def test_three_iters(self):
  884. a = ["a1", "a2", "a3", "a4", "a5"]
  885. b = ["b1", "b2", "b3"]
  886. c = ["c1"]
  887. actual = list(mi.interleave_evenly([a, b, c]))
  888. expected = ["a1", "b1", "a2", "c1", "a3", "b2", "a4", "b3", "a5"]
  889. self.assertEqual(actual, expected)
  890. def test_many_iters(self):
  891. # smoke test with many iterables: create iterables with a random
  892. # number of elements starting with a character ("a0", "a1", ...)
  893. rng = Random(0)
  894. iterables = []
  895. for ch in ascii_letters:
  896. length = rng.randint(0, 100)
  897. iterable = [f"{ch}{i}" for i in range(length)]
  898. iterables.append(iterable)
  899. interleaved = list(mi.interleave_evenly(iterables))
  900. # for each iterable, check that the result contains all its items
  901. for iterable, ch_expect in zip(iterables, ascii_letters):
  902. interleaved_actual = [
  903. e for e in interleaved if e.startswith(ch_expect)
  904. ]
  905. assert len(set(interleaved_actual)) == len(iterable)
  906. def test_manual_lengths(self):
  907. a = combinations(range(4), 2)
  908. len_a = 4 * (4 - 1) // 2 # == 6
  909. b = combinations(range(4), 3)
  910. len_b = 4
  911. expected = [
  912. (0, 1),
  913. (0, 1, 2),
  914. (0, 2),
  915. (0, 3),
  916. (0, 1, 3),
  917. (1, 2),
  918. (0, 2, 3),
  919. (1, 3),
  920. (2, 3),
  921. (1, 2, 3),
  922. ]
  923. actual = list(mi.interleave_evenly([a, b], lengths=[len_a, len_b]))
  924. self.assertEqual(expected, actual)
  925. def test_no_length_raises(self):
  926. # combinations doesn't have __len__, should trigger ValueError
  927. iterables = [range(5), combinations(range(5), 2)]
  928. with self.assertRaises(ValueError):
  929. list(mi.interleave_evenly(iterables))
  930. def test_argument_mismatch_raises(self):
  931. # pass mismatching number of iterables and lengths
  932. iterables = [range(3)]
  933. lengths = [3, 4]
  934. with self.assertRaises(ValueError):
  935. list(mi.interleave_evenly(iterables, lengths=lengths))
  936. class TestCollapse(TestCase):
  937. """Tests for ``collapse()``"""
  938. def test_collapse(self):
  939. l = [[1], 2, [[3], 4], [[[5]]]]
  940. self.assertEqual(list(mi.collapse(l)), [1, 2, 3, 4, 5])
  941. def test_collapse_to_string(self):
  942. l = [["s1"], "s2", [["s3"], "s4"], [[["s5"]]]]
  943. self.assertEqual(list(mi.collapse(l)), ["s1", "s2", "s3", "s4", "s5"])
  944. def test_collapse_to_bytes(self):
  945. l = [[b"s1"], b"s2", [[b"s3"], b"s4"], [[[b"s5"]]]]
  946. self.assertEqual(
  947. list(mi.collapse(l)), [b"s1", b"s2", b"s3", b"s4", b"s5"]
  948. )
  949. def test_collapse_flatten(self):
  950. l = [[1], [2], [[3], 4], [[[5]]]]
  951. self.assertEqual(list(mi.collapse(l, levels=1)), list(mi.flatten(l)))
  952. def test_collapse_to_level(self):
  953. l = [[1], 2, [[3], 4], [[[5]]]]
  954. self.assertEqual(list(mi.collapse(l, levels=2)), [1, 2, 3, 4, [5]])
  955. self.assertEqual(
  956. list(mi.collapse(mi.collapse(l, levels=1), levels=1)),
  957. list(mi.collapse(l, levels=2)),
  958. )
  959. def test_collapse_to_list(self):
  960. l = (1, [2], (3, [4, (5,)], 'ab'))
  961. actual = list(mi.collapse(l, base_type=list))
  962. expected = [1, [2], 3, [4, (5,)], 'ab']
  963. self.assertEqual(actual, expected)
  964. class SideEffectTests(TestCase):
  965. """Tests for ``side_effect()``"""
  966. def test_individual(self):
  967. # The function increments the counter for each call
  968. counter = [0]
  969. def func(arg):
  970. counter[0] += 1
  971. result = list(mi.side_effect(func, range(10)))
  972. self.assertEqual(result, list(range(10)))
  973. self.assertEqual(counter[0], 10)
  974. def test_chunked(self):
  975. # The function increments the counter for each call
  976. counter = [0]
  977. def func(arg):
  978. counter[0] += 1
  979. result = list(mi.side_effect(func, range(10), 2))
  980. self.assertEqual(result, list(range(10)))
  981. self.assertEqual(counter[0], 5)
  982. def test_before_after(self):
  983. f = StringIO()
  984. collector = []
  985. def func(item):
  986. print(item, file=f)
  987. collector.append(f.getvalue())
  988. def it():
  989. yield 'a'
  990. yield 'b'
  991. raise RuntimeError('kaboom')
  992. before = lambda: print('HEADER', file=f)
  993. after = f.close
  994. try:
  995. mi.consume(mi.side_effect(func, it(), before=before, after=after))
  996. except RuntimeError:
  997. pass
  998. # The iterable should have been written to the file
  999. self.assertEqual(collector, ['HEADER\na\n', 'HEADER\na\nb\n'])
  1000. # The file should be closed even though something bad happened
  1001. self.assertTrue(f.closed)
  1002. def test_before_fails(self):
  1003. f = StringIO()
  1004. func = lambda x: print(x, file=f)
  1005. def before():
  1006. raise RuntimeError('ouch')
  1007. try:
  1008. mi.consume(
  1009. mi.side_effect(func, 'abc', before=before, after=f.close)
  1010. )
  1011. except RuntimeError:
  1012. pass
  1013. # The file should be closed even though something bad happened in the
  1014. # before function
  1015. self.assertTrue(f.closed)
  1016. class SlicedTests(TestCase):
  1017. """Tests for ``sliced()``"""
  1018. def test_even(self):
  1019. """Test when the length of the sequence is divisible by *n*"""
  1020. seq = 'ABCDEFGHI'
  1021. self.assertEqual(list(mi.sliced(seq, 3)), ['ABC', 'DEF', 'GHI'])
  1022. def test_odd(self):
  1023. """Test when the length of the sequence is not divisible by *n*"""
  1024. seq = 'ABCDEFGHI'
  1025. self.assertEqual(list(mi.sliced(seq, 4)), ['ABCD', 'EFGH', 'I'])
  1026. def test_not_sliceable(self):
  1027. seq = (x for x in 'ABCDEFGHI')
  1028. with self.assertRaises(TypeError):
  1029. list(mi.sliced(seq, 3))
  1030. def test_odd_and_strict(self):
  1031. seq = [x for x in 'ABCDEFGHI']
  1032. with self.assertRaises(ValueError):
  1033. list(mi.sliced(seq, 4, strict=True))
  1034. def test_numpy_like_array(self):
  1035. # Numpy arrays don't behave like Python lists - calling bool()
  1036. # on them doesn't return False for empty lists and True for non-empty
  1037. # ones. Emulate that behavior.
  1038. class FalseList(list):
  1039. def __getitem__(self, key):
  1040. ret = super().__getitem__(key)
  1041. if isinstance(key, slice):
  1042. return FalseList(ret)
  1043. return ret
  1044. def __bool__(self):
  1045. return False
  1046. seq = FalseList(range(9))
  1047. actual = list(mi.sliced(seq, 3))
  1048. expected = [[0, 1, 2], [3, 4, 5], [6, 7, 8]]
  1049. self.assertEqual(actual, expected)
  1050. class SplitAtTests(TestCase):
  1051. def test_basic(self):
  1052. for iterable, separator in [
  1053. ('a,bb,ccc,dddd', ','),
  1054. (',a,bb,ccc,dddd', ','),
  1055. ('a,bb,ccc,dddd,', ','),
  1056. ('a,bb,ccc,,dddd', ','),
  1057. ('', ','),
  1058. (',', ','),
  1059. ('a,bb,ccc,dddd', ';'),
  1060. ]:
  1061. with self.subTest(iterable=iterable, separator=separator):
  1062. it = iter(iterable)
  1063. pred = lambda x: x == separator
  1064. actual = [''.join(x) for x in mi.split_at(it, pred)]
  1065. expected = iterable.split(separator)
  1066. self.assertEqual(actual, expected)
  1067. def test_maxsplit(self):
  1068. iterable = 'a,bb,ccc,dddd'
  1069. separator = ','
  1070. pred = lambda x: x == separator
  1071. for maxsplit in range(-1, 4):
  1072. with self.subTest(maxsplit=maxsplit):
  1073. it = iter(iterable)
  1074. result = mi.split_at(it, pred, maxsplit=maxsplit)
  1075. actual = [''.join(x) for x in result]
  1076. expected = iterable.split(separator, maxsplit)
  1077. self.assertEqual(actual, expected)
  1078. def test_keep_separator(self):
  1079. separator = ','
  1080. pred = lambda x: x == separator
  1081. for iterable, expected in [
  1082. ('a,bb,ccc', ['a', ',', 'bb', ',', 'ccc']),
  1083. (',a,bb,ccc', ['', ',', 'a', ',', 'bb', ',', 'ccc']),
  1084. ('a,bb,ccc,', ['a', ',', 'bb', ',', 'ccc', ',', '']),
  1085. ]:
  1086. with self.subTest(iterable=iterable):
  1087. it = iter(iterable)
  1088. result = mi.split_at(it, pred, keep_separator=True)
  1089. actual = [''.join(x) for x in result]
  1090. self.assertEqual(actual, expected)
  1091. def test_combination(self):
  1092. iterable = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
  1093. pred = lambda x: x % 3 == 0
  1094. actual = list(
  1095. mi.split_at(iterable, pred, maxsplit=2, keep_separator=True)
  1096. )
  1097. expected = [[1, 2], [3], [4, 5], [6], [7, 8, 9, 10]]
  1098. self.assertEqual(actual, expected)
  1099. class SplitBeforeTest(TestCase):
  1100. """Tests for ``split_before()``"""
  1101. def test_starts_with_sep(self):
  1102. actual = list(mi.split_before('xooxoo', lambda c: c == 'x'))
  1103. expected = [['x', 'o', 'o'], ['x', 'o', 'o']]
  1104. self.assertEqual(actual, expected)
  1105. def test_ends_with_sep(self):
  1106. actual = list(mi.split_before('ooxoox', lambda c: c == 'x'))
  1107. expected = [['o', 'o'], ['x', 'o', 'o'], ['x']]
  1108. self.assertEqual(actual, expected)
  1109. def test_no_sep(self):
  1110. actual = list(mi.split_before('ooo', lambda c: c == 'x'))
  1111. expected = [['o', 'o', 'o']]
  1112. self.assertEqual(actual, expected)
  1113. def test_empty_collection(self):
  1114. actual = list(mi.split_before([], lambda c: bool(c)))
  1115. expected = []
  1116. self.assertEqual(actual, expected)
  1117. def test_max_split(self):
  1118. for args, expected in [
  1119. (
  1120. ('a,b,c,d', lambda c: c == ',', -1),
  1121. [['a'], [',', 'b'], [',', 'c'], [',', 'd']],
  1122. ),
  1123. (
  1124. ('a,b,c,d', lambda c: c == ',', 0),
  1125. [['a', ',', 'b', ',', 'c', ',', 'd']],
  1126. ),
  1127. (
  1128. ('a,b,c,d', lambda c: c == ',', 1),
  1129. [['a'], [',', 'b', ',', 'c', ',', 'd']],
  1130. ),
  1131. (
  1132. ('a,b,c,d', lambda c: c == ',', 2),
  1133. [['a'], [',', 'b'], [',', 'c', ',', 'd']],
  1134. ),
  1135. (
  1136. ('a,b,c,d', lambda c: c == ',', 10),
  1137. [['a'], [',', 'b'], [',', 'c'], [',', 'd']],
  1138. ),
  1139. (
  1140. ('a,b,c,d', lambda c: c == '@', 2),
  1141. [['a', ',', 'b', ',', 'c', ',', 'd']],
  1142. ),
  1143. (
  1144. ('a,b,c,d', lambda c: c != ',', 2),
  1145. [['a', ','], ['b', ','], ['c', ',', 'd']],
  1146. ),
  1147. ]:
  1148. actual = list(mi.split_before(*args))
  1149. self.assertEqual(actual, expected)
  1150. class SplitAfterTest(TestCase):
  1151. """Tests for ``split_after()``"""
  1152. def test_starts_with_sep(self):
  1153. actual = list(mi.split_after('xooxoo', lambda c: c == 'x'))
  1154. expected = [['x'], ['o', 'o', 'x'], ['o', 'o']]
  1155. self.assertEqual(actual, expected)
  1156. def test_ends_with_sep(self):
  1157. actual = list(mi.split_after('ooxoox', lambda c: c == 'x'))
  1158. expected = [['o', 'o', 'x'], ['o', 'o', 'x']]
  1159. self.assertEqual(actual, expected)
  1160. def test_no_sep(self):
  1161. actual = list(mi.split_after('ooo', lambda c: c == 'x'))
  1162. expected = [['o', 'o', 'o']]
  1163. self.assertEqual(actual, expected)
  1164. def test_max_split(self):
  1165. for args, expected in [
  1166. (
  1167. ('a,b,c,d', lambda c: c == ',', -1),
  1168. [['a', ','], ['b', ','], ['c', ','], ['d']],
  1169. ),
  1170. (
  1171. ('a,b,c,d', lambda c: c == ',', 0),
  1172. [['a', ',', 'b', ',', 'c', ',', 'd']],
  1173. ),
  1174. (
  1175. ('a,b,c,d', lambda c: c == ',', 1),
  1176. [['a', ','], ['b', ',', 'c', ',', 'd']],
  1177. ),
  1178. (
  1179. ('a,b,c,d', lambda c: c == ',', 2),
  1180. [['a', ','], ['b', ','], ['c', ',', 'd']],
  1181. ),
  1182. (
  1183. ('a,b,c,d', lambda c: c == ',', 10),
  1184. [['a', ','], ['b', ','], ['c', ','], ['d']],
  1185. ),
  1186. (
  1187. ('a,b,c,d', lambda c: c == '@', 2),
  1188. [['a', ',', 'b', ',', 'c', ',', 'd']],
  1189. ),
  1190. (
  1191. ('a,b,c,d', lambda c: c != ',', 2),
  1192. [['a'], [',', 'b'], [',', 'c', ',', 'd']],
  1193. ),
  1194. (
  1195. ([1], lambda x: x == 1, 1),
  1196. [[1]],
  1197. ),
  1198. ]:
  1199. actual = list(mi.split_after(*args))
  1200. self.assertEqual(actual, expected)
  1201. class SplitWhenTests(TestCase):
  1202. """Tests for ``split_when()``"""
  1203. @staticmethod
  1204. def _split_when_before(iterable, pred):
  1205. return mi.split_when(iterable, lambda _, c: pred(c))
  1206. @staticmethod
  1207. def _split_when_after(iterable, pred):
  1208. return mi.split_when(iterable, lambda c, _: pred(c))
  1209. # split_before emulation
  1210. def test_before_emulation_starts_with_sep(self):
  1211. actual = list(self._split_when_before('xooxoo', lambda c: c == 'x'))
  1212. expected = [['x', 'o', 'o'], ['x', 'o', 'o']]
  1213. self.assertEqual(actual, expected)
  1214. def test_before_emulation_ends_with_sep(self):
  1215. actual = list(self._split_when_before('ooxoox', lambda c: c == 'x'))
  1216. expected = [['o', 'o'], ['x', 'o', 'o'], ['x']]
  1217. self.assertEqual(actual, expected)
  1218. def test_before_emulation_no_sep(self):
  1219. actual = list(self._split_when_before('ooo', lambda c: c == 'x'))
  1220. expected = [['o', 'o', 'o']]
  1221. self.assertEqual(actual, expected)
  1222. # split_after emulation
  1223. def test_after_emulation_starts_with_sep(self):
  1224. actual = list(self._split_when_after('xooxoo', lambda c: c == 'x'))
  1225. expected = [['x'], ['o', 'o', 'x'], ['o', 'o']]
  1226. self.assertEqual(actual, expected)
  1227. def test_after_emulation_ends_with_sep(self):
  1228. actual = list(self._split_when_after('ooxoox', lambda c: c == 'x'))
  1229. expected = [['o', 'o', 'x'], ['o', 'o', 'x']]
  1230. self.assertEqual(actual, expected)
  1231. def test_after_emulation_no_sep(self):
  1232. actual = list(self._split_when_after('ooo', lambda c: c == 'x'))
  1233. expected = [['o', 'o', 'o']]
  1234. self.assertEqual(actual, expected)
  1235. # edge cases
  1236. def test_empty_iterable(self):
  1237. actual = list(mi.split_when('', lambda a, b: a != b))
  1238. expected = []
  1239. self.assertEqual(actual, expected)
  1240. def test_one_element(self):
  1241. actual = list(mi.split_when('o', lambda a, b: a == b))
  1242. expected = [['o']]
  1243. self.assertEqual(actual, expected)
  1244. def test_one_element_is_second_item(self):
  1245. actual = list(self._split_when_before('x', lambda c: c == 'x'))
  1246. expected = [['x']]
  1247. self.assertEqual(actual, expected)
  1248. def test_one_element_is_first_item(self):
  1249. actual = list(self._split_when_after('x', lambda c: c == 'x'))
  1250. expected = [['x']]
  1251. self.assertEqual(actual, expected)
  1252. def test_max_split(self):
  1253. for args, expected in [
  1254. (
  1255. ('a,b,c,d', lambda a, _: a == ',', -1),
  1256. [['a', ','], ['b', ','], ['c', ','], ['d']],
  1257. ),
  1258. (
  1259. ('a,b,c,d', lambda a, _: a == ',', 0),
  1260. [['a', ',', 'b', ',', 'c', ',', 'd']],
  1261. ),
  1262. (
  1263. ('a,b,c,d', lambda _, b: b == ',', 1),
  1264. [['a'], [',', 'b', ',', 'c', ',', 'd']],
  1265. ),
  1266. (
  1267. ('a,b,c,d', lambda a, _: a == ',', 2),
  1268. [['a', ','], ['b', ','], ['c', ',', 'd']],
  1269. ),
  1270. (
  1271. ('0124376', lambda a, b: a > b, -1),
  1272. [['0', '1', '2', '4'], ['3', '7'], ['6']],
  1273. ),
  1274. (
  1275. ('0124376', lambda a, b: a > b, 0),
  1276. [['0', '1', '2', '4', '3', '7', '6']],
  1277. ),
  1278. (
  1279. ('0124376', lambda a, b: a > b, 1),
  1280. [['0', '1', '2', '4'], ['3', '7', '6']],
  1281. ),
  1282. (
  1283. ('0124376', lambda a, b: a > b, 2),
  1284. [['0', '1', '2', '4'], ['3', '7'], ['6']],
  1285. ),
  1286. ]:
  1287. actual = list(mi.split_when(*args))
  1288. self.assertEqual(actual, expected, str(args))
  1289. class SplitIntoTests(TestCase):
  1290. """Tests for ``split_into()``"""
  1291. def test_iterable_just_right(self):
  1292. """Size of ``iterable`` equals the sum of ``sizes``."""
  1293. iterable = [1, 2, 3, 4, 5, 6, 7, 8, 9]
  1294. sizes = [2, 3, 4]
  1295. expected = [[1, 2], [3, 4, 5], [6, 7, 8, 9]]
  1296. actual = list(mi.split_into(iterable, sizes))
  1297. self.assertEqual(actual, expected)
  1298. def test_iterable_too_small(self):
  1299. """Size of ``iterable`` is smaller than sum of ``sizes``. Last return
  1300. list is shorter as a result."""
  1301. iterable = [1, 2, 3, 4, 5, 6, 7]
  1302. sizes = [2, 3, 4]
  1303. expected = [[1, 2], [3, 4, 5], [6, 7]]
  1304. actual = list(mi.split_into(iterable, sizes))
  1305. self.assertEqual(actual, expected)
  1306. def test_iterable_too_small_extra(self):
  1307. """Size of ``iterable`` is smaller than sum of ``sizes``. Second last
  1308. return list is shorter and last return list is empty as a result."""
  1309. iterable = [1, 2, 3, 4, 5, 6, 7]
  1310. sizes = [2, 3, 4, 5]
  1311. expected = [[1, 2], [3, 4, 5], [6, 7], []]
  1312. actual = list(mi.split_into(iterable, sizes))
  1313. self.assertEqual(actual, expected)
  1314. def test_iterable_too_large(self):
  1315. """Size of ``iterable`` is larger than sum of ``sizes``. Not all
  1316. items of iterable are returned."""
  1317. iterable = [1, 2, 3, 4, 5, 6, 7, 8, 9]
  1318. sizes = [2, 3, 2]
  1319. expected = [[1, 2], [3, 4, 5], [6, 7]]
  1320. actual = list(mi.split_into(iterable, sizes))
  1321. self.assertEqual(actual, expected)
  1322. def test_using_none_with_leftover(self):
  1323. """Last item of ``sizes`` is None when items still remain in
  1324. ``iterable``. Last list returned stretches to fit all remaining items
  1325. of ``iterable``."""
  1326. iterable = [1, 2, 3, 4, 5, 6, 7, 8, 9]
  1327. sizes = [2, 3, None]
  1328. expected = [[1, 2], [3, 4, 5], [6, 7, 8, 9]]
  1329. actual = list(mi.split_into(iterable, sizes))
  1330. self.assertEqual(actual, expected)
  1331. def test_using_none_without_leftover(self):
  1332. """Last item of ``sizes`` is None when no items remain in
  1333. ``iterable``. Last list returned is empty."""
  1334. iterable = [1, 2, 3, 4, 5, 6, 7, 8, 9]
  1335. sizes = [2, 3, 4, None]
  1336. expected = [[1, 2], [3, 4, 5], [6, 7, 8, 9], []]
  1337. actual = list(mi.split_into(iterable, sizes))
  1338. self.assertEqual(actual, expected)
  1339. def test_using_none_mid_sizes(self):
  1340. """None is present in ``sizes`` but is not the last item. Last list
  1341. returned stretches to fit all remaining items of ``iterable`` but
  1342. all items in ``sizes`` after None are ignored."""
  1343. iterable = [1, 2, 3, 4, 5, 6, 7, 8, 9]
  1344. sizes = [2, 3, None, 4]
  1345. expected = [[1, 2], [3, 4, 5], [6, 7, 8, 9]]
  1346. actual = list(mi.split_into(iterable, sizes))
  1347. self.assertEqual(actual, expected)
  1348. def test_iterable_empty(self):
  1349. """``iterable`` argument is empty but ``sizes`` is not. An empty
  1350. list is returned for each item in ``sizes``."""
  1351. iterable = []
  1352. sizes = [2, 4, 2]
  1353. expected = [[], [], []]
  1354. actual = list(mi.split_into(iterable, sizes))
  1355. self.assertEqual(actual, expected)
  1356. def test_iterable_empty_using_none(self):
  1357. """``iterable`` argument is empty but ``sizes`` is not. An empty
  1358. list is returned for each item in ``sizes`` that is not after a
  1359. None item."""
  1360. iterable = []
  1361. sizes = [2, 4, None, 2]
  1362. expected = [[], [], []]
  1363. actual = list(mi.split_into(iterable, sizes))
  1364. self.assertEqual(actual, expected)
  1365. def test_sizes_empty(self):
  1366. """``sizes`` argument is empty but ``iterable`` is not. An empty
  1367. generator is returned."""
  1368. iterable = [1, 2, 3, 4, 5, 6, 7, 8, 9]
  1369. sizes = []
  1370. expected = []
  1371. actual = list(mi.split_into(iterable, sizes))
  1372. self.assertEqual(actual, expected)
  1373. def test_both_empty(self):
  1374. """Both ``sizes`` and ``iterable`` arguments are empty. An empty
  1375. generator is returned."""
  1376. iterable = []
  1377. sizes = []
  1378. expected = []
  1379. actual = list(mi.split_into(iterable, sizes))
  1380. self.assertEqual(actual, expected)
  1381. def test_bool_in_sizes(self):
  1382. """A bool object is present in ``sizes`` is treated as a 1 or 0 for
  1383. ``True`` or ``False`` due to bool being an instance of int."""
  1384. iterable = [1, 2, 3, 4, 5, 6, 7, 8, 9]
  1385. sizes = [3, True, 2, False]
  1386. expected = [[1, 2, 3], [4], [5, 6], []]
  1387. actual = list(mi.split_into(iterable, sizes))
  1388. self.assertEqual(actual, expected)
  1389. def test_invalid_in_sizes(self):
  1390. """A ValueError is raised if an object in ``sizes`` is neither ``None``
  1391. or an integer."""
  1392. iterable = [1, 2, 3, 4, 5, 6, 7, 8, 9]
  1393. sizes = [1, [], 3]
  1394. with self.assertRaises(ValueError):
  1395. list(mi.split_into(iterable, sizes))
  1396. def test_invalid_in_sizes_after_none(self):
  1397. """A item in ``sizes`` that is invalid will not raise a TypeError if it
  1398. comes after a ``None`` item."""
  1399. iterable = [1, 2, 3, 4, 5, 6, 7, 8, 9]
  1400. sizes = [3, 4, None, []]
  1401. expected = [[1, 2, 3], [4, 5, 6, 7], [8, 9]]
  1402. actual = list(mi.split_into(iterable, sizes))
  1403. self.assertEqual(actual, expected)
  1404. def test_generator_iterable_integrity(self):
  1405. """Check that if ``iterable`` is an iterator, it is consumed only by as
  1406. many items as the sum of ``sizes``."""
  1407. iterable = (i for i in range(10))
  1408. sizes = [2, 3]
  1409. expected = [[0, 1], [2, 3, 4]]
  1410. actual = list(mi.split_into(iterable, sizes))
  1411. self.assertEqual(actual, expected)
  1412. iterable_expected = [5, 6, 7, 8, 9]
  1413. iterable_actual = list(iterable)
  1414. self.assertEqual(iterable_actual, iterable_expected)
  1415. def test_generator_sizes_integrity(self):
  1416. """Check that if ``sizes`` is an iterator, it is consumed only until a
  1417. ``None`` item is reached"""
  1418. iterable = [1, 2, 3, 4, 5, 6, 7, 8, 9]
  1419. sizes = (i for i in [1, 2, None, 3, 4])
  1420. expected = [[1], [2, 3], [4, 5, 6, 7, 8, 9]]
  1421. actual = list(mi.split_into(iterable, sizes))
  1422. self.assertEqual(actual, expected)
  1423. sizes_expected = [3, 4]
  1424. sizes_actual = list(sizes)
  1425. self.assertEqual(sizes_actual, sizes_expected)
  1426. class PaddedTest(TestCase):
  1427. """Tests for ``padded()``"""
  1428. def test_no_n(self):
  1429. seq = [1, 2, 3]
  1430. # No fillvalue
  1431. self.assertEqual(mi.take(5, mi.padded(seq)), [1, 2, 3, None, None])
  1432. # With fillvalue
  1433. self.assertEqual(
  1434. mi.take(5, mi.padded(seq, fillvalue='')), [1, 2, 3, '', '']
  1435. )
  1436. def test_invalid_n(self):
  1437. self.assertRaises(ValueError, lambda: list(mi.padded([1, 2, 3], n=-1)))
  1438. self.assertRaises(ValueError, lambda: list(mi.padded([1, 2, 3], n=0)))
  1439. def test_valid_n(self):
  1440. seq = [1, 2, 3, 4, 5]
  1441. # No need for padding: len(seq) <= n
  1442. self.assertEqual(list(mi.padded(seq, n=4)), [1, 2, 3, 4, 5])
  1443. self.assertEqual(list(mi.padded(seq, n=5)), [1, 2, 3, 4, 5])
  1444. # No fillvalue
  1445. self.assertEqual(
  1446. list(mi.padded(seq, n=7)), [1, 2, 3, 4, 5, None, None]
  1447. )
  1448. # With fillvalue
  1449. self.assertEqual(
  1450. list(mi.padded(seq, fillvalue='', n=7)), [1, 2, 3, 4, 5, '', '']
  1451. )
  1452. def test_next_multiple(self):
  1453. seq = [1, 2, 3, 4, 5, 6]
  1454. # No need for padding: len(seq) % n == 0
  1455. self.assertEqual(
  1456. list(mi.padded(seq, n=3, next_multiple=True)), [1, 2, 3, 4, 5, 6]
  1457. )
  1458. # Padding needed: len(seq) < n
  1459. self.assertEqual(
  1460. list(mi.padded(seq, n=8, next_multiple=True)),
  1461. [1, 2, 3, 4, 5, 6, None, None],
  1462. )
  1463. # No padding needed: len(seq) == n
  1464. self.assertEqual(
  1465. list(mi.padded(seq, n=6, next_multiple=True)), [1, 2, 3, 4, 5, 6]
  1466. )
  1467. # Padding needed: len(seq) > n
  1468. self.assertEqual(
  1469. list(mi.padded(seq, n=4, next_multiple=True)),
  1470. [1, 2, 3, 4, 5, 6, None, None],
  1471. )
  1472. # With fillvalue
  1473. self.assertEqual(
  1474. list(mi.padded(seq, fillvalue='', n=4, next_multiple=True)),
  1475. [1, 2, 3, 4, 5, 6, '', ''],
  1476. )
  1477. class RepeatEachTests(TestCase):
  1478. """Tests for repeat_each()"""
  1479. def test_default(self):
  1480. actual = list(mi.repeat_each('ABC'))
  1481. expected = ['A', 'A', 'B', 'B', 'C', 'C']
  1482. self.assertEqual(actual, expected)
  1483. def test_basic(self):
  1484. actual = list(mi.repeat_each('ABC', 3))
  1485. expected = ['A', 'A', 'A', 'B', 'B', 'B', 'C', 'C', 'C']
  1486. self.assertEqual(actual, expected)
  1487. def test_empty(self):
  1488. actual = list(mi.repeat_each(''))
  1489. expected = []
  1490. self.assertEqual(actual, expected)
  1491. def test_no_repeat(self):
  1492. actual = list(mi.repeat_each('ABC', 0))
  1493. expected = []
  1494. self.assertEqual(actual, expected)
  1495. def test_negative_repeat(self):
  1496. actual = list(mi.repeat_each('ABC', -1))
  1497. expected = []
  1498. self.assertEqual(actual, expected)
  1499. def test_infinite_input(self):
  1500. repeater = mi.repeat_each(cycle('AB'))
  1501. actual = mi.take(6, repeater)
  1502. expected = ['A', 'A', 'B', 'B', 'A', 'A']
  1503. self.assertEqual(actual, expected)
  1504. class RepeatLastTests(TestCase):
  1505. def test_empty_iterable(self):
  1506. slice_length = 3
  1507. iterable = iter([])
  1508. actual = mi.take(slice_length, mi.repeat_last(iterable))
  1509. expected = [None] * slice_length
  1510. self.assertEqual(actual, expected)
  1511. def test_default_value(self):
  1512. slice_length = 3
  1513. iterable = iter([])
  1514. default = '3'
  1515. actual = mi.take(slice_length, mi.repeat_last(iterable, default))
  1516. expected = ['3'] * slice_length
  1517. self.assertEqual(actual, expected)
  1518. def test_basic(self):
  1519. slice_length = 10
  1520. iterable = (str(x) for x in range(5))
  1521. actual = mi.take(slice_length, mi.repeat_last(iterable))
  1522. expected = ['0', '1', '2', '3', '4', '4', '4', '4', '4', '4']
  1523. self.assertEqual(actual, expected)
  1524. class DistributeTest(TestCase):
  1525. """Tests for distribute()"""
  1526. def test_invalid_n(self):
  1527. self.assertRaises(ValueError, lambda: mi.distribute(-1, [1, 2, 3]))
  1528. self.assertRaises(ValueError, lambda: mi.distribute(0, [1, 2, 3]))
  1529. def test_basic(self):
  1530. iterable = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
  1531. for n, expected in [
  1532. (1, [iterable]),
  1533. (2, [[1, 3, 5, 7, 9], [2, 4, 6, 8, 10]]),
  1534. (3, [[1, 4, 7, 10], [2, 5, 8], [3, 6, 9]]),
  1535. (10, [[n] for n in range(1, 10 + 1)]),
  1536. ]:
  1537. self.assertEqual(
  1538. [list(x) for x in mi.distribute(n, iterable)], expected
  1539. )
  1540. def test_large_n(self):
  1541. iterable = [1, 2, 3, 4]
  1542. self.assertEqual(
  1543. [list(x) for x in mi.distribute(6, iterable)],
  1544. [[1], [2], [3], [4], [], []],
  1545. )
  1546. class StaggerTest(TestCase):
  1547. """Tests for ``stagger()``"""
  1548. def test_default(self):
  1549. iterable = [0, 1, 2, 3]
  1550. actual = list(mi.stagger(iterable))
  1551. expected = [(None, 0, 1), (0, 1, 2), (1, 2, 3)]
  1552. self.assertEqual(actual, expected)
  1553. def test_offsets(self):
  1554. iterable = [0, 1, 2, 3]
  1555. for offsets, expected in [
  1556. ((-2, 0, 2), [('', 0, 2), ('', 1, 3)]),
  1557. ((-2, -1), [('', ''), ('', 0), (0, 1), (1, 2), (2, 3)]),
  1558. ((1, 2), [(1, 2), (2, 3)]),
  1559. ]:
  1560. all_groups = mi.stagger(iterable, offsets=offsets, fillvalue='')
  1561. self.assertEqual(list(all_groups), expected)
  1562. def test_longest(self):
  1563. iterable = [0, 1, 2, 3]
  1564. for offsets, expected in [
  1565. (
  1566. (-1, 0, 1),
  1567. [('', 0, 1), (0, 1, 2), (1, 2, 3), (2, 3, ''), (3, '', '')],
  1568. ),
  1569. ((-2, -1), [('', ''), ('', 0), (0, 1), (1, 2), (2, 3), (3, '')]),
  1570. ((1, 2), [(1, 2), (2, 3), (3, '')]),
  1571. ]:
  1572. all_groups = mi.stagger(
  1573. iterable, offsets=offsets, fillvalue='', longest=True
  1574. )
  1575. self.assertEqual(list(all_groups), expected)
  1576. class ZipEqualTest(TestCase):
  1577. @skipIf(version_info[:2] < (3, 10), 'zip_equal deprecated for 3.10+')
  1578. def test_deprecation(self):
  1579. with warnings.catch_warnings(record=True) as caught:
  1580. warnings.simplefilter('always')
  1581. self.assertEqual(
  1582. list(mi.zip_equal([1, 2], [3, 4])), [(1, 3), (2, 4)]
  1583. )
  1584. (warning,) = caught
  1585. assert warning.category == DeprecationWarning
  1586. def test_equal(self):
  1587. lists = [0, 1, 2], [2, 3, 4]
  1588. for iterables in [lists, map(iter, lists)]:
  1589. actual = list(mi.zip_equal(*iterables))
  1590. expected = [(0, 2), (1, 3), (2, 4)]
  1591. self.assertEqual(actual, expected)
  1592. def test_unequal_lists(self):
  1593. two_items = [0, 1]
  1594. three_items = [2, 3, 4]
  1595. four_items = [5, 6, 7, 8]
  1596. # the mismatch is at index 1
  1597. try:
  1598. list(mi.zip_equal(two_items, three_items, four_items))
  1599. except mi.UnequalIterablesError as e:
  1600. self.assertEqual(
  1601. e.args[0],
  1602. (
  1603. 'Iterables have different lengths: '
  1604. 'index 0 has length 2; index 1 has length 3'
  1605. ),
  1606. )
  1607. # the mismatch is at index 2
  1608. try:
  1609. list(mi.zip_equal(two_items, two_items, four_items, four_items))
  1610. except mi.UnequalIterablesError as e:
  1611. self.assertEqual(
  1612. e.args[0],
  1613. (
  1614. 'Iterables have different lengths: '
  1615. 'index 0 has length 2; index 2 has length 4'
  1616. ),
  1617. )
  1618. # One without length: delegate to _zip_equal_generator
  1619. try:
  1620. list(mi.zip_equal(two_items, iter(two_items), three_items))
  1621. except mi.UnequalIterablesError as e:
  1622. self.assertEqual(e.args[0], 'Iterables have different lengths')
  1623. class ZipOffsetTest(TestCase):
  1624. """Tests for ``zip_offset()``"""
  1625. def test_shortest(self):
  1626. a_1 = [0, 1, 2, 3]
  1627. a_2 = [0, 1, 2, 3, 4, 5]
  1628. a_3 = [0, 1, 2, 3, 4, 5, 6, 7]
  1629. actual = list(
  1630. mi.zip_offset(a_1, a_2, a_3, offsets=(-1, 0, 1), fillvalue='')
  1631. )
  1632. expected = [('', 0, 1), (0, 1, 2), (1, 2, 3), (2, 3, 4), (3, 4, 5)]
  1633. self.assertEqual(actual, expected)
  1634. def test_longest(self):
  1635. a_1 = [0, 1, 2, 3]
  1636. a_2 = [0, 1, 2, 3, 4, 5]
  1637. a_3 = [0, 1, 2, 3, 4, 5, 6, 7]
  1638. actual = list(
  1639. mi.zip_offset(a_1, a_2, a_3, offsets=(-1, 0, 1), longest=True)
  1640. )
  1641. expected = [
  1642. (None, 0, 1),
  1643. (0, 1, 2),
  1644. (1, 2, 3),
  1645. (2, 3, 4),
  1646. (3, 4, 5),
  1647. (None, 5, 6),
  1648. (None, None, 7),
  1649. ]
  1650. self.assertEqual(actual, expected)
  1651. def test_mismatch(self):
  1652. iterables = [0, 1, 2], [2, 3, 4]
  1653. offsets = (-1, 0, 1)
  1654. self.assertRaises(
  1655. ValueError,
  1656. lambda: list(mi.zip_offset(*iterables, offsets=offsets)),
  1657. )
  1658. class UnzipTests(TestCase):
  1659. """Tests for unzip()"""
  1660. def test_empty_iterable(self):
  1661. self.assertEqual(list(mi.unzip([])), [])
  1662. # in reality zip([], [], []) is equivalent to iter([])
  1663. # but it doesn't hurt to test both
  1664. self.assertEqual(list(mi.unzip(zip([], [], []))), [])
  1665. def test_length_one_iterable(self):
  1666. xs, ys, zs = mi.unzip(zip([1], [2], [3]))
  1667. self.assertEqual(list(xs), [1])
  1668. self.assertEqual(list(ys), [2])
  1669. self.assertEqual(list(zs), [3])
  1670. def test_normal_case(self):
  1671. xs, ys, zs = range(10), range(1, 11), range(2, 12)
  1672. zipped = zip(xs, ys, zs)
  1673. xs, ys, zs = mi.unzip(zipped)
  1674. self.assertEqual(list(xs), list(range(10)))
  1675. self.assertEqual(list(ys), list(range(1, 11)))
  1676. self.assertEqual(list(zs), list(range(2, 12)))
  1677. def test_improperly_zipped(self):
  1678. zipped = iter([(1, 2, 3), (4, 5), (6,)])
  1679. xs, ys, zs = mi.unzip(zipped)
  1680. self.assertEqual(list(xs), [1, 4, 6])
  1681. self.assertEqual(list(ys), [2, 5])
  1682. self.assertEqual(list(zs), [3])
  1683. def test_increasingly_zipped(self):
  1684. zipped = iter([(1, 2), (3, 4, 5), (6, 7, 8, 9)])
  1685. unzipped = mi.unzip(zipped)
  1686. # from the docstring:
  1687. # len(first tuple) is the number of iterables zipped
  1688. self.assertEqual(len(unzipped), 2)
  1689. xs, ys = unzipped
  1690. self.assertEqual(list(xs), [1, 3, 6])
  1691. self.assertEqual(list(ys), [2, 4, 7])
  1692. class SortTogetherTest(TestCase):
  1693. """Tests for sort_together()"""
  1694. def test_key_list(self):
  1695. """tests `key_list` including default, iterables include duplicates"""
  1696. iterables = [
  1697. ['GA', 'GA', 'GA', 'CT', 'CT', 'CT'],
  1698. ['May', 'Aug.', 'May', 'June', 'July', 'July'],
  1699. [97, 20, 100, 70, 100, 20],
  1700. ]
  1701. self.assertEqual(
  1702. mi.sort_together(iterables),
  1703. [
  1704. ('CT', 'CT', 'CT', 'GA', 'GA', 'GA'),
  1705. ('June', 'July', 'July', 'May', 'Aug.', 'May'),
  1706. (70, 100, 20, 97, 20, 100),
  1707. ],
  1708. )
  1709. self.assertEqual(
  1710. mi.sort_together(iterables, key_list=(0, 1)),
  1711. [
  1712. ('CT', 'CT', 'CT', 'GA', 'GA', 'GA'),
  1713. ('July', 'July', 'June', 'Aug.', 'May', 'May'),
  1714. (100, 20, 70, 20, 97, 100),
  1715. ],
  1716. )
  1717. self.assertEqual(
  1718. mi.sort_together(iterables, key_list=(0, 1, 2)),
  1719. [
  1720. ('CT', 'CT', 'CT', 'GA', 'GA', 'GA'),
  1721. ('July', 'July', 'June', 'Aug.', 'May', 'May'),
  1722. (20, 100, 70, 20, 97, 100),
  1723. ],
  1724. )
  1725. self.assertEqual(
  1726. mi.sort_together(iterables, key_list=(2,)),
  1727. [
  1728. ('GA', 'CT', 'CT', 'GA', 'GA', 'CT'),
  1729. ('Aug.', 'July', 'June', 'May', 'May', 'July'),
  1730. (20, 20, 70, 97, 100, 100),
  1731. ],
  1732. )
  1733. def test_invalid_key_list(self):
  1734. """tests `key_list` for indexes not available in `iterables`"""
  1735. iterables = [
  1736. ['GA', 'GA', 'GA', 'CT', 'CT', 'CT'],
  1737. ['May', 'Aug.', 'May', 'June', 'July', 'July'],
  1738. [97, 20, 100, 70, 100, 20],
  1739. ]
  1740. self.assertRaises(
  1741. IndexError, lambda: mi.sort_together(iterables, key_list=(5,))
  1742. )
  1743. def test_key_function(self):
  1744. """tests `key` function, including interaction with `key_list`"""
  1745. iterables = [
  1746. ['GA', 'GA', 'GA', 'CT', 'CT', 'CT'],
  1747. ['May', 'Aug.', 'May', 'June', 'July', 'July'],
  1748. [97, 20, 100, 70, 100, 20],
  1749. ]
  1750. self.assertEqual(
  1751. mi.sort_together(iterables, key=lambda x: x),
  1752. [
  1753. ('CT', 'CT', 'CT', 'GA', 'GA', 'GA'),
  1754. ('June', 'July', 'July', 'May', 'Aug.', 'May'),
  1755. (70, 100, 20, 97, 20, 100),
  1756. ],
  1757. )
  1758. self.assertEqual(
  1759. mi.sort_together(iterables, key=lambda x: x[::-1]),
  1760. [
  1761. ('GA', 'GA', 'GA', 'CT', 'CT', 'CT'),
  1762. ('May', 'Aug.', 'May', 'June', 'July', 'July'),
  1763. (97, 20, 100, 70, 100, 20),
  1764. ],
  1765. )
  1766. self.assertEqual(
  1767. mi.sort_together(
  1768. iterables,
  1769. key_list=(0, 2),
  1770. key=lambda state, number: (
  1771. number if state == 'CT' else 2 * number
  1772. ),
  1773. ),
  1774. [
  1775. ('CT', 'GA', 'CT', 'CT', 'GA', 'GA'),
  1776. ('July', 'Aug.', 'June', 'July', 'May', 'May'),
  1777. (20, 20, 70, 100, 97, 100),
  1778. ],
  1779. )
  1780. def test_reverse(self):
  1781. """tests `reverse` to ensure a reverse sort for `key_list` iterables"""
  1782. iterables = [
  1783. ['GA', 'GA', 'GA', 'CT', 'CT', 'CT'],
  1784. ['May', 'Aug.', 'May', 'June', 'July', 'July'],
  1785. [97, 20, 100, 70, 100, 20],
  1786. ]
  1787. self.assertEqual(
  1788. mi.sort_together(iterables, key_list=(0, 1, 2), reverse=True),
  1789. [
  1790. ('GA', 'GA', 'GA', 'CT', 'CT', 'CT'),
  1791. ('May', 'May', 'Aug.', 'June', 'July', 'July'),
  1792. (100, 97, 20, 70, 100, 20),
  1793. ],
  1794. )
  1795. def test_uneven_iterables(self):
  1796. """tests trimming of iterables to the shortest length before sorting"""
  1797. iterables = [
  1798. ['GA', 'GA', 'GA', 'CT', 'CT', 'CT', 'MA'],
  1799. ['May', 'Aug.', 'May', 'June', 'July', 'July'],
  1800. [97, 20, 100, 70, 100, 20, 0],
  1801. ]
  1802. self.assertEqual(
  1803. mi.sort_together(iterables),
  1804. [
  1805. ('CT', 'CT', 'CT', 'GA', 'GA', 'GA'),
  1806. ('June', 'July', 'July', 'May', 'Aug.', 'May'),
  1807. (70, 100, 20, 97, 20, 100),
  1808. ],
  1809. )
  1810. def test_strict(self):
  1811. # Test for list of lists or tuples
  1812. self.assertRaises(
  1813. mi.UnequalIterablesError,
  1814. lambda: mi.sort_together(
  1815. [(4, 3, 2, 1), ('a', 'b', 'c')], strict=True
  1816. ),
  1817. )
  1818. # Test for list of iterables
  1819. self.assertRaises(
  1820. mi.UnequalIterablesError,
  1821. lambda: mi.sort_together([range(4), range(5)], strict=True),
  1822. )
  1823. # Test for iterable of iterables
  1824. self.assertRaises(
  1825. mi.UnequalIterablesError,
  1826. lambda: mi.sort_together(
  1827. (range(i) for i in range(4)), strict=True
  1828. ),
  1829. )
  1830. class DivideTest(TestCase):
  1831. """Tests for divide()"""
  1832. def test_invalid_n(self):
  1833. self.assertRaises(ValueError, lambda: mi.divide(-1, [1, 2, 3]))
  1834. self.assertRaises(ValueError, lambda: mi.divide(0, [1, 2, 3]))
  1835. def test_basic(self):
  1836. iterable = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
  1837. for n, expected in [
  1838. (1, [iterable]),
  1839. (2, [[1, 2, 3, 4, 5], [6, 7, 8, 9, 10]]),
  1840. (3, [[1, 2, 3, 4], [5, 6, 7], [8, 9, 10]]),
  1841. (10, [[n] for n in range(1, 10 + 1)]),
  1842. ]:
  1843. self.assertEqual(
  1844. [list(x) for x in mi.divide(n, iterable)], expected
  1845. )
  1846. def test_large_n(self):
  1847. self.assertEqual(
  1848. [list(x) for x in mi.divide(6, iter(range(1, 4 + 1)))],
  1849. [[1], [2], [3], [4], [], []],
  1850. )
  1851. class TestAlwaysIterable(TestCase):
  1852. """Tests for always_iterable()"""
  1853. def test_single(self):
  1854. self.assertEqual(list(mi.always_iterable(1)), [1])
  1855. def test_strings(self):
  1856. for obj in ['foo', b'bar', 'baz']:
  1857. actual = list(mi.always_iterable(obj))
  1858. expected = [obj]
  1859. self.assertEqual(actual, expected)
  1860. def test_base_type(self):
  1861. dict_obj = {'a': 1, 'b': 2}
  1862. str_obj = '123'
  1863. # Default: dicts are iterable like they normally are
  1864. default_actual = list(mi.always_iterable(dict_obj))
  1865. default_expected = list(dict_obj)
  1866. self.assertEqual(default_actual, default_expected)
  1867. # Unitary types set: dicts are not iterable
  1868. custom_actual = list(mi.always_iterable(dict_obj, base_type=dict))
  1869. custom_expected = [dict_obj]
  1870. self.assertEqual(custom_actual, custom_expected)
  1871. # With unitary types set, strings are iterable
  1872. str_actual = list(mi.always_iterable(str_obj, base_type=None))
  1873. str_expected = list(str_obj)
  1874. self.assertEqual(str_actual, str_expected)
  1875. # base_type handles nested tuple (via isinstance).
  1876. base_type = ((dict,),)
  1877. custom_actual = list(mi.always_iterable(dict_obj, base_type=base_type))
  1878. custom_expected = [dict_obj]
  1879. self.assertEqual(custom_actual, custom_expected)
  1880. def test_iterables(self):
  1881. self.assertEqual(list(mi.always_iterable([0, 1])), [0, 1])
  1882. self.assertEqual(
  1883. list(mi.always_iterable([0, 1], base_type=list)), [[0, 1]]
  1884. )
  1885. self.assertEqual(
  1886. list(mi.always_iterable(iter('foo'))), ['f', 'o', 'o']
  1887. )
  1888. self.assertEqual(list(mi.always_iterable([])), [])
  1889. def test_none(self):
  1890. self.assertEqual(list(mi.always_iterable(None)), [])
  1891. def test_generator(self):
  1892. def _gen():
  1893. yield 0
  1894. yield 1
  1895. self.assertEqual(list(mi.always_iterable(_gen())), [0, 1])
  1896. class AdjacentTests(TestCase):
  1897. def test_typical(self):
  1898. actual = list(mi.adjacent(lambda x: x % 5 == 0, range(10)))
  1899. expected = [
  1900. (True, 0),
  1901. (True, 1),
  1902. (False, 2),
  1903. (False, 3),
  1904. (True, 4),
  1905. (True, 5),
  1906. (True, 6),
  1907. (False, 7),
  1908. (False, 8),
  1909. (False, 9),
  1910. ]
  1911. self.assertEqual(actual, expected)
  1912. def test_empty_iterable(self):
  1913. actual = list(mi.adjacent(lambda x: x % 5 == 0, []))
  1914. expected = []
  1915. self.assertEqual(actual, expected)
  1916. def test_length_one(self):
  1917. actual = list(mi.adjacent(lambda x: x % 5 == 0, [0]))
  1918. expected = [(True, 0)]
  1919. self.assertEqual(actual, expected)
  1920. actual = list(mi.adjacent(lambda x: x % 5 == 0, [1]))
  1921. expected = [(False, 1)]
  1922. self.assertEqual(actual, expected)
  1923. def test_consecutive_true(self):
  1924. """Test that when the predicate matches multiple consecutive elements
  1925. it doesn't repeat elements in the output"""
  1926. actual = list(mi.adjacent(lambda x: x % 5 < 2, range(10)))
  1927. expected = [
  1928. (True, 0),
  1929. (True, 1),
  1930. (True, 2),
  1931. (False, 3),
  1932. (True, 4),
  1933. (True, 5),
  1934. (True, 6),
  1935. (True, 7),
  1936. (False, 8),
  1937. (False, 9),
  1938. ]
  1939. self.assertEqual(actual, expected)
  1940. def test_distance(self):
  1941. actual = list(mi.adjacent(lambda x: x % 5 == 0, range(10), distance=2))
  1942. expected = [
  1943. (True, 0),
  1944. (True, 1),
  1945. (True, 2),
  1946. (True, 3),
  1947. (True, 4),
  1948. (True, 5),
  1949. (True, 6),
  1950. (True, 7),
  1951. (False, 8),
  1952. (False, 9),
  1953. ]
  1954. self.assertEqual(actual, expected)
  1955. actual = list(mi.adjacent(lambda x: x % 5 == 0, range(10), distance=3))
  1956. expected = [
  1957. (True, 0),
  1958. (True, 1),
  1959. (True, 2),
  1960. (True, 3),
  1961. (True, 4),
  1962. (True, 5),
  1963. (True, 6),
  1964. (True, 7),
  1965. (True, 8),
  1966. (False, 9),
  1967. ]
  1968. self.assertEqual(actual, expected)
  1969. def test_large_distance(self):
  1970. """Test distance larger than the length of the iterable"""
  1971. iterable = range(10)
  1972. actual = list(mi.adjacent(lambda x: x % 5 == 4, iterable, distance=20))
  1973. expected = list(zip(repeat(True), iterable))
  1974. self.assertEqual(actual, expected)
  1975. actual = list(mi.adjacent(lambda x: False, iterable, distance=20))
  1976. expected = list(zip(repeat(False), iterable))
  1977. self.assertEqual(actual, expected)
  1978. def test_zero_distance(self):
  1979. """Test that adjacent() reduces to zip+map when distance is 0"""
  1980. iterable = range(1000)
  1981. predicate = lambda x: x % 4 == 2
  1982. actual = mi.adjacent(predicate, iterable, 0)
  1983. expected = zip(map(predicate, iterable), iterable)
  1984. self.assertTrue(all(a == e for a, e in zip(actual, expected)))
  1985. def test_negative_distance(self):
  1986. """Test that adjacent() raises an error with negative distance"""
  1987. pred = lambda x: x
  1988. self.assertRaises(
  1989. ValueError, lambda: mi.adjacent(pred, range(1000), -1)
  1990. )
  1991. self.assertRaises(
  1992. ValueError, lambda: mi.adjacent(pred, range(10), -10)
  1993. )
  1994. def test_grouping(self):
  1995. """Test interaction of adjacent() with groupby_transform()"""
  1996. iterable = mi.adjacent(lambda x: x % 5 == 0, range(10))
  1997. grouper = mi.groupby_transform(iterable, itemgetter(0), itemgetter(1))
  1998. actual = [(k, list(g)) for k, g in grouper]
  1999. expected = [
  2000. (True, [0, 1]),
  2001. (False, [2, 3]),
  2002. (True, [4, 5, 6]),
  2003. (False, [7, 8, 9]),
  2004. ]
  2005. self.assertEqual(actual, expected)
  2006. def test_call_once(self):
  2007. """Test that the predicate is only called once per item."""
  2008. already_seen = set()
  2009. iterable = range(10)
  2010. def predicate(item):
  2011. self.assertNotIn(item, already_seen)
  2012. already_seen.add(item)
  2013. return True
  2014. actual = list(mi.adjacent(predicate, iterable))
  2015. expected = [(True, x) for x in iterable]
  2016. self.assertEqual(actual, expected)
  2017. class GroupByTransformTests(TestCase):
  2018. def assertAllGroupsEqual(self, groupby1, groupby2):
  2019. for a, b in zip(groupby1, groupby2):
  2020. key1, group1 = a
  2021. key2, group2 = b
  2022. self.assertEqual(key1, key2)
  2023. self.assertListEqual(list(group1), list(group2))
  2024. self.assertRaises(StopIteration, lambda: next(groupby1))
  2025. self.assertRaises(StopIteration, lambda: next(groupby2))
  2026. def test_default_funcs(self):
  2027. iterable = [(x // 5, x) for x in range(1000)]
  2028. actual = mi.groupby_transform(iterable)
  2029. expected = groupby(iterable)
  2030. self.assertAllGroupsEqual(actual, expected)
  2031. def test_valuefunc(self):
  2032. iterable = [(int(x / 5), int(x / 3), x) for x in range(10)]
  2033. # Test the standard usage of grouping one iterable using another's keys
  2034. grouper = mi.groupby_transform(
  2035. iterable, keyfunc=itemgetter(0), valuefunc=itemgetter(-1)
  2036. )
  2037. actual = [(k, list(g)) for k, g in grouper]
  2038. expected = [(0, [0, 1, 2, 3, 4]), (1, [5, 6, 7, 8, 9])]
  2039. self.assertEqual(actual, expected)
  2040. grouper = mi.groupby_transform(
  2041. iterable, keyfunc=itemgetter(1), valuefunc=itemgetter(-1)
  2042. )
  2043. actual = [(k, list(g)) for k, g in grouper]
  2044. expected = [(0, [0, 1, 2]), (1, [3, 4, 5]), (2, [6, 7, 8]), (3, [9])]
  2045. self.assertEqual(actual, expected)
  2046. # and now for something a little different
  2047. d = dict(zip(range(10), 'abcdefghij'))
  2048. grouper = mi.groupby_transform(
  2049. range(10), keyfunc=lambda x: x // 5, valuefunc=d.get
  2050. )
  2051. actual = [(k, ''.join(g)) for k, g in grouper]
  2052. expected = [(0, 'abcde'), (1, 'fghij')]
  2053. self.assertEqual(actual, expected)
  2054. def test_no_valuefunc(self):
  2055. iterable = range(1000)
  2056. def key(x):
  2057. return x // 5
  2058. actual = mi.groupby_transform(iterable, key, valuefunc=None)
  2059. expected = groupby(iterable, key)
  2060. self.assertAllGroupsEqual(actual, expected)
  2061. actual = mi.groupby_transform(iterable, key) # default valuefunc
  2062. expected = groupby(iterable, key)
  2063. self.assertAllGroupsEqual(actual, expected)
  2064. def test_reducefunc(self):
  2065. iterable = range(50)
  2066. keyfunc = lambda k: 10 * (k // 10)
  2067. valuefunc = lambda v: v + 1
  2068. reducefunc = sum
  2069. actual = list(
  2070. mi.groupby_transform(
  2071. iterable,
  2072. keyfunc=keyfunc,
  2073. valuefunc=valuefunc,
  2074. reducefunc=reducefunc,
  2075. )
  2076. )
  2077. expected = [(0, 55), (10, 155), (20, 255), (30, 355), (40, 455)]
  2078. self.assertEqual(actual, expected)
  2079. class NumericRangeTests(TestCase):
  2080. def test_basic(self):
  2081. for args, expected in [
  2082. ((4,), [0, 1, 2, 3]),
  2083. ((4.0,), [0.0, 1.0, 2.0, 3.0]),
  2084. ((1.0, 4), [1.0, 2.0, 3.0]),
  2085. ((1, 4.0), [1.0, 2.0, 3.0]),
  2086. ((1.0, 5), [1.0, 2.0, 3.0, 4.0]),
  2087. ((0, 20, 5), [0, 5, 10, 15]),
  2088. ((0, 20, 5.0), [0.0, 5.0, 10.0, 15.0]),
  2089. ((0, 10, 3), [0, 3, 6, 9]),
  2090. ((0, 10, 3.0), [0.0, 3.0, 6.0, 9.0]),
  2091. ((0, -5, -1), [0, -1, -2, -3, -4]),
  2092. ((0.0, -5, -1), [0.0, -1.0, -2.0, -3.0, -4.0]),
  2093. ((1, 2, Fraction(1, 2)), [Fraction(1, 1), Fraction(3, 2)]),
  2094. ((0,), []),
  2095. ((0.0,), []),
  2096. ((1, 0), []),
  2097. ((1.0, 0.0), []),
  2098. ((0.1, 0.30000000000000001, 0.2), [0.1]), # IEE 754 !
  2099. (
  2100. (
  2101. Decimal("0.1"),
  2102. Decimal("0.30000000000000001"),
  2103. Decimal("0.2"),
  2104. ),
  2105. [Decimal("0.1"), Decimal("0.3")],
  2106. ), # okay with Decimal
  2107. (
  2108. (
  2109. Fraction(1, 10),
  2110. Fraction(30000000000000001, 100000000000000000),
  2111. Fraction(2, 10),
  2112. ),
  2113. [Fraction(1, 10), Fraction(3, 10)],
  2114. ), # okay with Fraction
  2115. ((Fraction(2, 1),), [Fraction(0, 1), Fraction(1, 1)]),
  2116. ((Decimal('2.0'),), [Decimal('0.0'), Decimal('1.0')]),
  2117. (
  2118. (
  2119. datetime(2019, 3, 29, 12, 34, 56),
  2120. datetime(2019, 3, 29, 12, 37, 55),
  2121. timedelta(minutes=1),
  2122. ),
  2123. [
  2124. datetime(2019, 3, 29, 12, 34, 56),
  2125. datetime(2019, 3, 29, 12, 35, 56),
  2126. datetime(2019, 3, 29, 12, 36, 56),
  2127. ],
  2128. ),
  2129. ]:
  2130. actual = list(mi.numeric_range(*args))
  2131. self.assertEqual(expected, actual)
  2132. self.assertTrue(
  2133. all(type(a) is type(e) for a, e in zip(actual, expected))
  2134. )
  2135. def test_arg_count(self):
  2136. for args, message in [
  2137. ((), 'numeric_range expected at least 1 argument, got 0'),
  2138. (
  2139. (0, 1, 2, 3),
  2140. 'numeric_range expected at most 3 arguments, got 4',
  2141. ),
  2142. ]:
  2143. with self.assertRaisesRegex(TypeError, message):
  2144. mi.numeric_range(*args)
  2145. def test_zero_step(self):
  2146. for args in [
  2147. (1, 2, 0),
  2148. (
  2149. datetime(2019, 3, 29, 12, 34, 56),
  2150. datetime(2019, 3, 29, 12, 37, 55),
  2151. timedelta(minutes=0),
  2152. ),
  2153. (1.0, 2.0, 0.0),
  2154. (Decimal("1.0"), Decimal("2.0"), Decimal("0.0")),
  2155. (Fraction(2, 2), Fraction(4, 2), Fraction(0, 2)),
  2156. ]:
  2157. with self.assertRaises(ValueError):
  2158. list(mi.numeric_range(*args))
  2159. def test_bool(self):
  2160. for args, expected in [
  2161. ((1.0, 3.0, 1.5), True),
  2162. ((1.0, 2.0, 1.5), True),
  2163. ((1.0, 1.0, 1.5), False),
  2164. ((1.0, 0.0, 1.5), False),
  2165. ((3.0, 1.0, -1.5), True),
  2166. ((2.0, 1.0, -1.5), True),
  2167. ((1.0, 1.0, -1.5), False),
  2168. ((0.0, 1.0, -1.5), False),
  2169. ((Decimal("1.0"), Decimal("2.0"), Decimal("1.5")), True),
  2170. ((Decimal("1.0"), Decimal("0.0"), Decimal("1.5")), False),
  2171. ((Fraction(2, 2), Fraction(4, 2), Fraction(3, 2)), True),
  2172. ((Fraction(2, 2), Fraction(0, 2), Fraction(3, 2)), False),
  2173. (
  2174. (
  2175. datetime(2019, 3, 29),
  2176. datetime(2019, 3, 30),
  2177. timedelta(hours=1),
  2178. ),
  2179. True,
  2180. ),
  2181. (
  2182. (
  2183. datetime(2019, 3, 29),
  2184. datetime(2019, 3, 28),
  2185. timedelta(hours=1),
  2186. ),
  2187. False,
  2188. ),
  2189. ]:
  2190. self.assertEqual(expected, bool(mi.numeric_range(*args)))
  2191. def test_contains(self):
  2192. for args, expected_in, expected_not_in in [
  2193. ((10,), range(10), (0.5,)),
  2194. ((1.0, 9.9, 1.5), (1.0, 2.5, 4.0, 5.5, 7.0, 8.5), (0.9,)),
  2195. ((9.0, 1.0, -1.5), (1.5, 3.0, 4.5, 6.0, 7.5, 9.0), (0.0, 0.9)),
  2196. (
  2197. (Decimal("1.0"), Decimal("9.9"), Decimal("1.5")),
  2198. (
  2199. Decimal("1.0"),
  2200. Decimal("2.5"),
  2201. Decimal("4.0"),
  2202. Decimal("5.5"),
  2203. Decimal("7.0"),
  2204. Decimal("8.5"),
  2205. ),
  2206. (Decimal("0.9"),),
  2207. ),
  2208. (
  2209. (Fraction(0, 1), Fraction(5, 1), Fraction(1, 2)),
  2210. (Fraction(0, 1), Fraction(1, 2), Fraction(9, 2)),
  2211. (Fraction(10, 2),),
  2212. ),
  2213. (
  2214. (
  2215. datetime(2019, 3, 29),
  2216. datetime(2019, 3, 30),
  2217. timedelta(hours=1),
  2218. ),
  2219. (datetime(2019, 3, 29, 15),),
  2220. (datetime(2019, 3, 29, 15, 30),),
  2221. ),
  2222. ]:
  2223. r = mi.numeric_range(*args)
  2224. for v in expected_in:
  2225. self.assertTrue(v in r)
  2226. self.assertFalse(v not in r)
  2227. for v in expected_not_in:
  2228. self.assertFalse(v in r)
  2229. self.assertTrue(v not in r)
  2230. def test_eq(self):
  2231. for args1, args2 in [
  2232. ((0, 5, 2), (0, 6, 2)),
  2233. ((1.0, 9.9, 1.5), (1.0, 8.6, 1.5)),
  2234. ((8.5, 0.0, -1.5), (8.5, 0.7, -1.5)),
  2235. ((7.0, 0.0, 1.0), (17.0, 7.0, 0.5)),
  2236. (
  2237. (Decimal("1.0"), Decimal("9.9"), Decimal("1.5")),
  2238. (Decimal("1.0"), Decimal("8.6"), Decimal("1.5")),
  2239. ),
  2240. (
  2241. (Fraction(1, 1), Fraction(10, 1), Fraction(3, 2)),
  2242. (Fraction(1, 1), Fraction(9, 1), Fraction(3, 2)),
  2243. ),
  2244. (
  2245. (
  2246. datetime(2019, 3, 29),
  2247. datetime(2019, 3, 30),
  2248. timedelta(hours=10),
  2249. ),
  2250. (
  2251. datetime(2019, 3, 29),
  2252. datetime(2019, 3, 30, 1),
  2253. timedelta(hours=10),
  2254. ),
  2255. ),
  2256. ]:
  2257. self.assertEqual(
  2258. mi.numeric_range(*args1), mi.numeric_range(*args2)
  2259. )
  2260. for args1, args2 in [
  2261. ((0, 5, 2), (0, 7, 2)),
  2262. ((1.0, 9.9, 1.5), (1.2, 9.9, 1.5)),
  2263. ((1.0, 9.9, 1.5), (1.0, 10.3, 1.5)),
  2264. ((1.0, 9.9, 1.5), (1.0, 9.9, 1.4)),
  2265. ((8.5, 0.0, -1.5), (8.4, 0.0, -1.5)),
  2266. ((8.5, 0.0, -1.5), (8.5, -0.7, -1.5)),
  2267. ((8.5, 0.0, -1.5), (8.5, 0.0, -1.4)),
  2268. ((0.0, 7.0, 1.0), (7.0, 0.0, 1.0)),
  2269. (
  2270. (Decimal("1.0"), Decimal("10.0"), Decimal("1.5")),
  2271. (Decimal("1.0"), Decimal("10.5"), Decimal("1.5")),
  2272. ),
  2273. (
  2274. (Fraction(1, 1), Fraction(10, 1), Fraction(3, 2)),
  2275. (Fraction(1, 1), Fraction(21, 2), Fraction(3, 2)),
  2276. ),
  2277. (
  2278. (
  2279. datetime(2019, 3, 29),
  2280. datetime(2019, 3, 30),
  2281. timedelta(hours=10),
  2282. ),
  2283. (
  2284. datetime(2019, 3, 29),
  2285. datetime(2019, 3, 30, 15),
  2286. timedelta(hours=10),
  2287. ),
  2288. ),
  2289. ]:
  2290. self.assertNotEqual(
  2291. mi.numeric_range(*args1), mi.numeric_range(*args2)
  2292. )
  2293. self.assertNotEqual(mi.numeric_range(7.0), 1)
  2294. self.assertNotEqual(mi.numeric_range(7.0), "abc")
  2295. def test_get_item_by_index(self):
  2296. for args, index, expected in [
  2297. ((1, 6), 2, 3),
  2298. ((1.0, 6.0, 1.5), 0, 1.0),
  2299. ((1.0, 6.0, 1.5), 1, 2.5),
  2300. ((1.0, 6.0, 1.5), 2, 4.0),
  2301. ((1.0, 6.0, 1.5), 3, 5.5),
  2302. ((1.0, 6.0, 1.5), -1, 5.5),
  2303. ((1.0, 6.0, 1.5), -2, 4.0),
  2304. (
  2305. (Decimal("1.0"), Decimal("9.0"), Decimal("1.5")),
  2306. -1,
  2307. Decimal("8.5"),
  2308. ),
  2309. (
  2310. (Fraction(1, 1), Fraction(10, 1), Fraction(3, 2)),
  2311. 2,
  2312. Fraction(4, 1),
  2313. ),
  2314. (
  2315. (
  2316. datetime(2019, 3, 29),
  2317. datetime(2019, 3, 30),
  2318. timedelta(hours=10),
  2319. ),
  2320. 1,
  2321. datetime(2019, 3, 29, 10),
  2322. ),
  2323. ]:
  2324. self.assertEqual(expected, mi.numeric_range(*args)[index])
  2325. for args, index in [
  2326. ((1.0, 6.0, 1.5), 4),
  2327. ((1.0, 6.0, 1.5), -5),
  2328. ((6.0, 1.0, 1.5), 0),
  2329. ((6.0, 1.0, 1.5), -1),
  2330. ((Decimal("1.0"), Decimal("9.0"), Decimal("-1.5")), -1),
  2331. ((Fraction(1, 1), Fraction(2, 1), Fraction(3, 2)), 2),
  2332. (
  2333. (
  2334. datetime(2019, 3, 29),
  2335. datetime(2019, 3, 30),
  2336. timedelta(hours=10),
  2337. ),
  2338. 8,
  2339. ),
  2340. ]:
  2341. with self.assertRaises(IndexError):
  2342. mi.numeric_range(*args)[index]
  2343. def test_get_item_by_slice(self):
  2344. for args, sl, expected_args in [
  2345. ((1.0, 9.0, 1.5), slice(None, None, None), (1.0, 9.0, 1.5)),
  2346. ((1.0, 9.0, 1.5), slice(None, 1, None), (1.0, 2.5, 1.5)),
  2347. ((1.0, 9.0, 1.5), slice(None, None, 2), (1.0, 9.0, 3.0)),
  2348. ((1.0, 9.0, 1.5), slice(None, 2, None), (1.0, 4.0, 1.5)),
  2349. ((1.0, 9.0, 1.5), slice(1, 2, None), (2.5, 4.0, 1.5)),
  2350. ((1.0, 9.0, 1.5), slice(1, -1, None), (2.5, 8.5, 1.5)),
  2351. ((1.0, 9.0, 1.5), slice(10, None, 3), (9.0, 9.0, 4.5)),
  2352. ((1.0, 9.0, 1.5), slice(-10, None, 3), (1.0, 9.0, 4.5)),
  2353. ((1.0, 9.0, 1.5), slice(None, -10, 3), (1.0, 1.0, 4.5)),
  2354. ((1.0, 9.0, 1.5), slice(None, 10, 3), (1.0, 9.0, 4.5)),
  2355. (
  2356. (Decimal("1.0"), Decimal("9.0"), Decimal("1.5")),
  2357. slice(1, -1, None),
  2358. (Decimal("2.5"), Decimal("8.5"), Decimal("1.5")),
  2359. ),
  2360. (
  2361. (Fraction(1, 1), Fraction(5, 1), Fraction(3, 2)),
  2362. slice(1, -1, None),
  2363. (Fraction(5, 2), Fraction(4, 1), Fraction(3, 2)),
  2364. ),
  2365. (
  2366. (
  2367. datetime(2019, 3, 29),
  2368. datetime(2019, 3, 30),
  2369. timedelta(hours=10),
  2370. ),
  2371. slice(1, -1, None),
  2372. (
  2373. datetime(2019, 3, 29, 10),
  2374. datetime(2019, 3, 29, 20),
  2375. timedelta(hours=10),
  2376. ),
  2377. ),
  2378. ]:
  2379. self.assertEqual(
  2380. mi.numeric_range(*expected_args), mi.numeric_range(*args)[sl]
  2381. )
  2382. def test_hash(self):
  2383. for args, expected in [
  2384. ((1.0, 6.0, 1.5), hash((1.0, 5.5, 1.5))),
  2385. ((1.0, 7.0, 1.5), hash((1.0, 5.5, 1.5))),
  2386. ((1.0, 7.5, 1.5), hash((1.0, 7.0, 1.5))),
  2387. ((1.0, 1.5, 1.5), hash((1.0, 1.0, 1.5))),
  2388. ((1.5, 1.0, 1.5), hash(range(0, 0))),
  2389. ((1.5, 1.5, 1.5), hash(range(0, 0))),
  2390. (
  2391. (Decimal("1.0"), Decimal("9.0"), Decimal("1.5")),
  2392. hash((Decimal("1.0"), Decimal("8.5"), Decimal("1.5"))),
  2393. ),
  2394. (
  2395. (Fraction(1, 1), Fraction(5, 1), Fraction(3, 2)),
  2396. hash((Fraction(1, 1), Fraction(4, 1), Fraction(3, 2))),
  2397. ),
  2398. (
  2399. (
  2400. datetime(2019, 3, 29),
  2401. datetime(2019, 3, 30),
  2402. timedelta(hours=10),
  2403. ),
  2404. hash(
  2405. (
  2406. datetime(2019, 3, 29),
  2407. datetime(2019, 3, 29, 20),
  2408. timedelta(hours=10),
  2409. )
  2410. ),
  2411. ),
  2412. ]:
  2413. self.assertEqual(expected, hash(mi.numeric_range(*args)))
  2414. def test_iter_twice(self):
  2415. r1 = mi.numeric_range(1.0, 9.9, 1.5)
  2416. r2 = mi.numeric_range(8.5, 0.0, -1.5)
  2417. self.assertEqual([1.0, 2.5, 4.0, 5.5, 7.0, 8.5], list(r1))
  2418. self.assertEqual([1.0, 2.5, 4.0, 5.5, 7.0, 8.5], list(r1))
  2419. self.assertEqual([8.5, 7.0, 5.5, 4.0, 2.5, 1.0], list(r2))
  2420. self.assertEqual([8.5, 7.0, 5.5, 4.0, 2.5, 1.0], list(r2))
  2421. def test_len(self):
  2422. for args, expected in [
  2423. ((1.0, 7.0, 1.5), 4),
  2424. ((1.0, 7.01, 1.5), 5),
  2425. ((7.0, 1.0, -1.5), 4),
  2426. ((7.01, 1.0, -1.5), 5),
  2427. ((0.1, 0.30000000000000001, 0.2), 1), # IEE 754 !
  2428. (
  2429. (
  2430. Decimal("0.1"),
  2431. Decimal("0.30000000000000001"),
  2432. Decimal("0.2"),
  2433. ),
  2434. 2,
  2435. ), # works with Decimal
  2436. ((Decimal("1.0"), Decimal("9.0"), Decimal("1.5")), 6),
  2437. ((Fraction(1, 1), Fraction(5, 1), Fraction(3, 2)), 3),
  2438. (
  2439. (
  2440. datetime(2019, 3, 29),
  2441. datetime(2019, 3, 30),
  2442. timedelta(hours=10),
  2443. ),
  2444. 3,
  2445. ),
  2446. ]:
  2447. self.assertEqual(expected, len(mi.numeric_range(*args)))
  2448. def test_repr(self):
  2449. for args, *expected in [
  2450. ((7.0,), "numeric_range(0.0, 7.0)"),
  2451. ((1.0, 7.0), "numeric_range(1.0, 7.0)"),
  2452. ((7.0, 1.0, -1.5), "numeric_range(7.0, 1.0, -1.5)"),
  2453. (
  2454. (Decimal("1.0"), Decimal("9.0"), Decimal("1.5")),
  2455. (
  2456. "numeric_range(Decimal('1.0'), Decimal('9.0'), "
  2457. "Decimal('1.5'))"
  2458. ),
  2459. ),
  2460. (
  2461. (Fraction(7, 7), Fraction(10, 2), Fraction(3, 2)),
  2462. (
  2463. "numeric_range(Fraction(1, 1), Fraction(5, 1), "
  2464. "Fraction(3, 2))"
  2465. ),
  2466. ),
  2467. (
  2468. (
  2469. datetime(2019, 3, 29),
  2470. datetime(2019, 3, 30),
  2471. timedelta(hours=10),
  2472. ),
  2473. "numeric_range(datetime.datetime(2019, 3, 29, 0, 0), "
  2474. "datetime.datetime(2019, 3, 30, 0, 0), "
  2475. "datetime.timedelta(seconds=36000))",
  2476. "numeric_range(datetime.datetime(2019, 3, 29, 0, 0), "
  2477. "datetime.datetime(2019, 3, 30, 0, 0), "
  2478. "datetime.timedelta(0, 36000))",
  2479. ),
  2480. ]:
  2481. with self.subTest(args=args):
  2482. self.assertIn(repr(mi.numeric_range(*args)), expected)
  2483. def test_reversed(self):
  2484. for args, expected in [
  2485. ((7.0,), [6.0, 5.0, 4.0, 3.0, 2.0, 1.0, 0.0]),
  2486. ((1.0, 7.0), [6.0, 5.0, 4.0, 3.0, 2.0, 1.0]),
  2487. ((7.0, 1.0, -1.5), [2.5, 4.0, 5.5, 7.0]),
  2488. ((7.0, 0.9, -1.5), [1.0, 2.5, 4.0, 5.5, 7.0]),
  2489. (
  2490. (Decimal("1.0"), Decimal("5.0"), Decimal("1.5")),
  2491. [Decimal('4.0'), Decimal('2.5'), Decimal('1.0')],
  2492. ),
  2493. (
  2494. (Fraction(1, 1), Fraction(5, 1), Fraction(3, 2)),
  2495. [Fraction(4, 1), Fraction(5, 2), Fraction(1, 1)],
  2496. ),
  2497. (
  2498. (
  2499. datetime(2019, 3, 29),
  2500. datetime(2019, 3, 30),
  2501. timedelta(hours=10),
  2502. ),
  2503. [
  2504. datetime(2019, 3, 29, 20),
  2505. datetime(2019, 3, 29, 10),
  2506. datetime(2019, 3, 29),
  2507. ],
  2508. ),
  2509. ]:
  2510. self.assertEqual(expected, list(reversed(mi.numeric_range(*args))))
  2511. def test_count(self):
  2512. for args, v, c in [
  2513. ((7.0,), 0.0, 1),
  2514. ((7.0,), 0.5, 0),
  2515. ((7.0,), 6.0, 1),
  2516. ((7.0,), 7.0, 0),
  2517. ((7.0,), 10.0, 0),
  2518. (
  2519. (Decimal("1.0"), Decimal("5.0"), Decimal("1.5")),
  2520. Decimal('4.0'),
  2521. 1,
  2522. ),
  2523. (
  2524. (Fraction(1, 1), Fraction(5, 1), Fraction(3, 2)),
  2525. Fraction(5, 2),
  2526. 1,
  2527. ),
  2528. (
  2529. (
  2530. datetime(2019, 3, 29),
  2531. datetime(2019, 3, 30),
  2532. timedelta(hours=10),
  2533. ),
  2534. datetime(2019, 3, 29, 20),
  2535. 1,
  2536. ),
  2537. ]:
  2538. self.assertEqual(c, mi.numeric_range(*args).count(v))
  2539. def test_index(self):
  2540. for args, v, i in [
  2541. ((7.0,), 0.0, 0),
  2542. ((7.0,), 6.0, 6),
  2543. ((7.0, 0.0, -1.0), 7.0, 0),
  2544. ((7.0, 0.0, -1.0), 1.0, 6),
  2545. (
  2546. (Decimal("1.0"), Decimal("5.0"), Decimal("1.5")),
  2547. Decimal('4.0'),
  2548. 2,
  2549. ),
  2550. (
  2551. (Fraction(1, 1), Fraction(5, 1), Fraction(3, 2)),
  2552. Fraction(5, 2),
  2553. 1,
  2554. ),
  2555. (
  2556. (
  2557. datetime(2019, 3, 29),
  2558. datetime(2019, 3, 30),
  2559. timedelta(hours=10),
  2560. ),
  2561. datetime(2019, 3, 29, 20),
  2562. 2,
  2563. ),
  2564. ]:
  2565. self.assertEqual(i, mi.numeric_range(*args).index(v))
  2566. for args, v in [
  2567. ((0.7,), 0.5),
  2568. ((0.7,), 7.0),
  2569. ((0.7,), 10.0),
  2570. ((7.0, 0.0, -1.0), 0.5),
  2571. ((7.0, 0.0, -1.0), 0.0),
  2572. ((7.0, 0.0, -1.0), 10.0),
  2573. ((7.0, 0.0), 5.0),
  2574. ((Decimal("1.0"), Decimal("5.0"), Decimal("1.5")), Decimal('4.5')),
  2575. ((Fraction(1, 1), Fraction(5, 1), Fraction(3, 2)), Fraction(5, 3)),
  2576. (
  2577. (
  2578. datetime(2019, 3, 29),
  2579. datetime(2019, 3, 30),
  2580. timedelta(hours=10),
  2581. ),
  2582. datetime(2019, 3, 30),
  2583. ),
  2584. ]:
  2585. with self.assertRaises(ValueError):
  2586. mi.numeric_range(*args).index(v)
  2587. def test_parent_classes(self):
  2588. r = mi.numeric_range(7.0)
  2589. self.assertTrue(isinstance(r, abc.Iterable))
  2590. self.assertFalse(isinstance(r, abc.Iterator))
  2591. self.assertTrue(isinstance(r, abc.Sequence))
  2592. self.assertTrue(isinstance(r, abc.Hashable))
  2593. def test_bad_key(self):
  2594. r = mi.numeric_range(7.0)
  2595. for arg, message in [
  2596. ('a', 'numeric range indices must be integers or slices, not str'),
  2597. (
  2598. (),
  2599. 'numeric range indices must be integers or slices, not tuple',
  2600. ),
  2601. ]:
  2602. with self.assertRaisesRegex(TypeError, message):
  2603. r[arg]
  2604. def test_pickle(self):
  2605. for args in [
  2606. (7.0,),
  2607. (5.0, 7.0),
  2608. (5.0, 7.0, 3.0),
  2609. (7.0, 5.0),
  2610. (7.0, 5.0, 4.0),
  2611. (7.0, 5.0, -1.0),
  2612. (Decimal("1.0"), Decimal("5.0"), Decimal("1.5")),
  2613. (Fraction(1, 1), Fraction(5, 1), Fraction(3, 2)),
  2614. (datetime(2019, 3, 29), datetime(2019, 3, 30)),
  2615. ]:
  2616. r = mi.numeric_range(*args)
  2617. self.assertTrue(dumps(r)) # assert not empty
  2618. self.assertEqual(r, loads(dumps(r)))
  2619. class CountCycleTests(TestCase):
  2620. def test_basic(self):
  2621. expected = [
  2622. (0, 'a'),
  2623. (0, 'b'),
  2624. (0, 'c'),
  2625. (1, 'a'),
  2626. (1, 'b'),
  2627. (1, 'c'),
  2628. (2, 'a'),
  2629. (2, 'b'),
  2630. (2, 'c'),
  2631. ]
  2632. for actual in [
  2633. mi.take(9, mi.count_cycle('abc')), # n=None
  2634. list(mi.count_cycle('abc', 3)), # n=3
  2635. ]:
  2636. self.assertEqual(actual, expected)
  2637. def test_empty(self):
  2638. self.assertEqual(list(mi.count_cycle('')), [])
  2639. self.assertEqual(list(mi.count_cycle('', 2)), [])
  2640. def test_negative(self):
  2641. self.assertEqual(list(mi.count_cycle('abc', -3)), [])
  2642. class MarkEndsTests(TestCase):
  2643. def test_basic(self):
  2644. for size, expected in [
  2645. (0, []),
  2646. (1, [(True, True, '0')]),
  2647. (2, [(True, False, '0'), (False, True, '1')]),
  2648. (3, [(True, False, '0'), (False, False, '1'), (False, True, '2')]),
  2649. (
  2650. 4,
  2651. [
  2652. (True, False, '0'),
  2653. (False, False, '1'),
  2654. (False, False, '2'),
  2655. (False, True, '3'),
  2656. ],
  2657. ),
  2658. ]:
  2659. with self.subTest(size=size):
  2660. iterable = map(str, range(size))
  2661. actual = list(mi.mark_ends(iterable))
  2662. self.assertEqual(actual, expected)
  2663. class LocateTests(TestCase):
  2664. def test_default_pred(self):
  2665. iterable = [0, 1, 1, 0, 1, 0, 0]
  2666. actual = list(mi.locate(iterable))
  2667. expected = [1, 2, 4]
  2668. self.assertEqual(actual, expected)
  2669. def test_no_matches(self):
  2670. iterable = [0, 0, 0]
  2671. actual = list(mi.locate(iterable))
  2672. expected = []
  2673. self.assertEqual(actual, expected)
  2674. def test_custom_pred(self):
  2675. iterable = ['0', 1, 1, '0', 1, '0', '0']
  2676. pred = lambda x: x == '0'
  2677. actual = list(mi.locate(iterable, pred))
  2678. expected = [0, 3, 5, 6]
  2679. self.assertEqual(actual, expected)
  2680. def test_window_size(self):
  2681. iterable = ['0', 1, 1, '0', 1, '0', '0']
  2682. pred = lambda *args: args == ('0', 1)
  2683. actual = list(mi.locate(iterable, pred, window_size=2))
  2684. expected = [0, 3]
  2685. self.assertEqual(actual, expected)
  2686. def test_window_size_large(self):
  2687. iterable = [1, 2, 3, 4]
  2688. pred = lambda a, b, c, d, e: True
  2689. actual = list(mi.locate(iterable, pred, window_size=5))
  2690. expected = [0]
  2691. self.assertEqual(actual, expected)
  2692. def test_window_size_zero(self):
  2693. iterable = [1, 2, 3, 4]
  2694. pred = lambda: True
  2695. with self.assertRaises(ValueError):
  2696. list(mi.locate(iterable, pred, window_size=0))
  2697. class StripFunctionTests(TestCase):
  2698. def test_hashable(self):
  2699. iterable = list('www.example.com')
  2700. pred = lambda x: x in set('cmowz.')
  2701. self.assertEqual(list(mi.lstrip(iterable, pred)), list('example.com'))
  2702. self.assertEqual(list(mi.rstrip(iterable, pred)), list('www.example'))
  2703. self.assertEqual(list(mi.strip(iterable, pred)), list('example'))
  2704. def test_not_hashable(self):
  2705. iterable = [
  2706. list('http://'),
  2707. list('www'),
  2708. list('.example'),
  2709. list('.com'),
  2710. ]
  2711. pred = lambda x: x in [list('http://'), list('www'), list('.com')]
  2712. self.assertEqual(list(mi.lstrip(iterable, pred)), iterable[2:])
  2713. self.assertEqual(list(mi.rstrip(iterable, pred)), iterable[:3])
  2714. self.assertEqual(list(mi.strip(iterable, pred)), iterable[2:3])
  2715. def test_math(self):
  2716. iterable = [0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2]
  2717. pred = lambda x: x <= 2
  2718. self.assertEqual(list(mi.lstrip(iterable, pred)), iterable[3:])
  2719. self.assertEqual(list(mi.rstrip(iterable, pred)), iterable[:-3])
  2720. self.assertEqual(list(mi.strip(iterable, pred)), iterable[3:-3])
  2721. class IsliceExtendedTests(TestCase):
  2722. def test_all(self):
  2723. iterable = ['0', '1', '2', '3', '4', '5']
  2724. indexes = [*range(-4, 10), None]
  2725. steps = [1, 2, 3, 4, -1, -2, -3, -4]
  2726. for slice_args in product(indexes, indexes, steps):
  2727. with self.subTest(slice_args=slice_args):
  2728. actual = list(mi.islice_extended(iterable, *slice_args))
  2729. expected = iterable[slice(*slice_args)]
  2730. self.assertEqual(actual, expected, slice_args)
  2731. def test_zero_step(self):
  2732. with self.assertRaises(ValueError):
  2733. list(mi.islice_extended([1, 2, 3], 0, 1, 0))
  2734. def test_slicing(self):
  2735. iterable = map(str, count())
  2736. first_slice = mi.islice_extended(iterable)[10:]
  2737. second_slice = mi.islice_extended(first_slice)[:10]
  2738. third_slice = mi.islice_extended(second_slice)[::2]
  2739. self.assertEqual(list(third_slice), ['10', '12', '14', '16', '18'])
  2740. def test_slicing_extensive(self):
  2741. iterable = range(10)
  2742. options = (None, 1, 2, 7, -1)
  2743. for start, stop, step in product(options, options, options):
  2744. with self.subTest(slice_args=(start, stop, step)):
  2745. sliced_tuple_0 = tuple(
  2746. mi.islice_extended(iterable)[start:stop:step]
  2747. )
  2748. sliced_tuple_1 = tuple(
  2749. mi.islice_extended(iterable, start, stop, step)
  2750. )
  2751. sliced_range = tuple(iterable[start:stop:step])
  2752. self.assertEqual(sliced_tuple_0, sliced_range)
  2753. self.assertEqual(sliced_tuple_1, sliced_range)
  2754. def test_invalid_slice(self):
  2755. with self.assertRaises(TypeError):
  2756. mi.islice_extended(count())[13]
  2757. class ConsecutiveGroupsTest(TestCase):
  2758. def test_numbers(self):
  2759. iterable = [-10, -8, -7, -6, 1, 2, 4, 5, -1, 7]
  2760. actual = [list(g) for g in mi.consecutive_groups(iterable)]
  2761. expected = [[-10], [-8, -7, -6], [1, 2], [4, 5], [-1], [7]]
  2762. self.assertEqual(actual, expected)
  2763. def test_custom_ordering(self):
  2764. iterable = ['1', '10', '11', '20', '21', '22', '30', '31']
  2765. ordering = lambda x: int(x)
  2766. actual = [list(g) for g in mi.consecutive_groups(iterable, ordering)]
  2767. expected = [['1'], ['10', '11'], ['20', '21', '22'], ['30', '31']]
  2768. self.assertEqual(actual, expected)
  2769. def test_exotic_ordering(self):
  2770. iterable = [
  2771. ('a', 'b', 'c', 'd'),
  2772. ('a', 'c', 'b', 'd'),
  2773. ('a', 'c', 'd', 'b'),
  2774. ('a', 'd', 'b', 'c'),
  2775. ('d', 'b', 'c', 'a'),
  2776. ('d', 'c', 'a', 'b'),
  2777. ]
  2778. ordering = list(permutations('abcd')).index
  2779. actual = [list(g) for g in mi.consecutive_groups(iterable, ordering)]
  2780. expected = [
  2781. [('a', 'b', 'c', 'd')],
  2782. [('a', 'c', 'b', 'd'), ('a', 'c', 'd', 'b'), ('a', 'd', 'b', 'c')],
  2783. [('d', 'b', 'c', 'a'), ('d', 'c', 'a', 'b')],
  2784. ]
  2785. self.assertEqual(actual, expected)
  2786. class DifferenceTest(TestCase):
  2787. def test_normal(self):
  2788. iterable = [10, 20, 30, 40, 50]
  2789. actual = list(mi.difference(iterable))
  2790. expected = [10, 10, 10, 10, 10]
  2791. self.assertEqual(actual, expected)
  2792. def test_custom(self):
  2793. iterable = [10, 20, 30, 40, 50]
  2794. actual = list(mi.difference(iterable, add))
  2795. expected = [10, 30, 50, 70, 90]
  2796. self.assertEqual(actual, expected)
  2797. def test_roundtrip(self):
  2798. original = list(range(100))
  2799. accumulated = accumulate(original)
  2800. actual = list(mi.difference(accumulated))
  2801. self.assertEqual(actual, original)
  2802. def test_one(self):
  2803. self.assertEqual(list(mi.difference([0])), [0])
  2804. def test_empty(self):
  2805. self.assertEqual(list(mi.difference([])), [])
  2806. def test_initial(self):
  2807. original = list(range(100))
  2808. accumulated = accumulate(original, initial=100)
  2809. actual = list(mi.difference(accumulated, initial=100))
  2810. self.assertEqual(actual, original)
  2811. class SeekableTest(PeekableMixinTests, TestCase):
  2812. cls = mi.seekable
  2813. def test_exhaustion_reset(self):
  2814. iterable = [str(n) for n in range(10)]
  2815. s = mi.seekable(iterable)
  2816. self.assertEqual(list(s), iterable) # Normal iteration
  2817. self.assertEqual(list(s), []) # Iterable is exhausted
  2818. s.seek(0)
  2819. self.assertEqual(list(s), iterable) # Back in action
  2820. def test_partial_reset(self):
  2821. iterable = [str(n) for n in range(10)]
  2822. s = mi.seekable(iterable)
  2823. self.assertEqual(mi.take(5, s), iterable[:5]) # Normal iteration
  2824. s.seek(1)
  2825. self.assertEqual(list(s), iterable[1:]) # Get the rest of the iterable
  2826. def test_forward(self):
  2827. iterable = [str(n) for n in range(10)]
  2828. s = mi.seekable(iterable)
  2829. self.assertEqual(mi.take(1, s), iterable[:1]) # Normal iteration
  2830. s.seek(3) # Skip over index 2
  2831. self.assertEqual(list(s), iterable[3:]) # Result is similar to slicing
  2832. s.seek(0) # Back to 0
  2833. self.assertEqual(list(s), iterable) # No difference in result
  2834. def test_past_end(self):
  2835. iterable = [str(n) for n in range(10)]
  2836. s = mi.seekable(iterable)
  2837. self.assertEqual(mi.take(1, s), iterable[:1]) # Normal iteration
  2838. s.seek(20)
  2839. self.assertEqual(list(s), []) # Iterable is exhausted
  2840. s.seek(0) # Back to 0
  2841. self.assertEqual(list(s), iterable) # No difference in result
  2842. def test_elements(self):
  2843. iterable = map(str, count())
  2844. s = mi.seekable(iterable)
  2845. mi.take(10, s)
  2846. elements = s.elements()
  2847. self.assertEqual(
  2848. [elements[i] for i in range(10)], [str(n) for n in range(10)]
  2849. )
  2850. self.assertEqual(len(elements), 10)
  2851. mi.take(10, s)
  2852. self.assertEqual(list(elements), [str(n) for n in range(20)])
  2853. def test_maxlen(self):
  2854. iterable = map(str, count())
  2855. s = mi.seekable(iterable, maxlen=4)
  2856. self.assertEqual(mi.take(10, s), [str(n) for n in range(10)])
  2857. self.assertEqual(list(s.elements()), ['6', '7', '8', '9'])
  2858. s.seek(0)
  2859. self.assertEqual(mi.take(14, s), [str(n) for n in range(6, 20)])
  2860. self.assertEqual(list(s.elements()), ['16', '17', '18', '19'])
  2861. def test_maxlen_zero(self):
  2862. iterable = [str(x) for x in range(5)]
  2863. s = mi.seekable(iterable, maxlen=0)
  2864. self.assertEqual(list(s), iterable)
  2865. self.assertEqual(list(s.elements()), [])
  2866. def test_relative_seek(self):
  2867. iterable = [str(x) for x in range(5)]
  2868. s = mi.seekable(iterable)
  2869. s.relative_seek(2)
  2870. self.assertEqual(next(s), '2')
  2871. s.relative_seek(-2)
  2872. self.assertEqual(next(s), '1')
  2873. s.relative_seek(-2)
  2874. self.assertEqual(
  2875. next(s), '0'
  2876. ) # Seek relative to current position within the cache
  2877. s.relative_seek(-10) # Lower bound
  2878. self.assertEqual(next(s), '0')
  2879. s.relative_seek(10) # Lower bound
  2880. self.assertEqual(list(s.elements()), [str(x) for x in range(5)])
  2881. class SequenceViewTests(TestCase):
  2882. def test_init(self):
  2883. view = mi.SequenceView((1, 2, 3))
  2884. self.assertEqual(repr(view), "SequenceView((1, 2, 3))")
  2885. self.assertRaises(TypeError, lambda: mi.SequenceView({}))
  2886. def test_update(self):
  2887. seq = [1, 2, 3]
  2888. view = mi.SequenceView(seq)
  2889. self.assertEqual(len(view), 3)
  2890. self.assertEqual(repr(view), "SequenceView([1, 2, 3])")
  2891. seq.pop()
  2892. self.assertEqual(len(view), 2)
  2893. self.assertEqual(repr(view), "SequenceView([1, 2])")
  2894. def test_indexing(self):
  2895. seq = ('a', 'b', 'c', 'd', 'e', 'f')
  2896. view = mi.SequenceView(seq)
  2897. for i in range(-len(seq), len(seq)):
  2898. self.assertEqual(view[i], seq[i])
  2899. def test_slicing(self):
  2900. seq = ('a', 'b', 'c', 'd', 'e', 'f')
  2901. view = mi.SequenceView(seq)
  2902. n = len(seq)
  2903. indexes = list(range(-n - 1, n + 1)) + [None]
  2904. steps = list(range(-n, n + 1))
  2905. steps.remove(0)
  2906. for slice_args in product(indexes, indexes, steps):
  2907. i = slice(*slice_args)
  2908. self.assertEqual(view[i], seq[i])
  2909. def test_abc_methods(self):
  2910. # collections.Sequence should provide all of this functionality
  2911. seq = ('a', 'b', 'c', 'd', 'e', 'f', 'f')
  2912. view = mi.SequenceView(seq)
  2913. # __contains__
  2914. self.assertIn('b', view)
  2915. self.assertNotIn('g', view)
  2916. # __iter__
  2917. self.assertEqual(list(iter(view)), list(seq))
  2918. # __reversed__
  2919. self.assertEqual(list(reversed(view)), list(reversed(seq)))
  2920. # index
  2921. self.assertEqual(view.index('b'), 1)
  2922. # count
  2923. self.assertEqual(seq.count('f'), 2)
  2924. class RunLengthTest(TestCase):
  2925. def test_encode(self):
  2926. iterable = (int(str(n)[0]) for n in count(800))
  2927. actual = mi.take(4, mi.run_length.encode(iterable))
  2928. expected = [(8, 100), (9, 100), (1, 1000), (2, 1000)]
  2929. self.assertEqual(actual, expected)
  2930. def test_decode(self):
  2931. iterable = [('d', 4), ('c', 3), ('b', 2), ('a', 1)]
  2932. actual = ''.join(mi.run_length.decode(iterable))
  2933. expected = 'ddddcccbba'
  2934. self.assertEqual(actual, expected)
  2935. class ExactlyNTests(TestCase):
  2936. """Tests for ``exactly_n()``"""
  2937. def test_true(self):
  2938. """Iterable has ``n`` ``True`` elements"""
  2939. self.assertTrue(mi.exactly_n([True, False, True], 2))
  2940. self.assertTrue(mi.exactly_n([1, 1, 1, 0], 3))
  2941. self.assertTrue(mi.exactly_n([False, False], 0))
  2942. self.assertTrue(mi.exactly_n(range(100), 10, lambda x: x < 10))
  2943. def test_false(self):
  2944. """Iterable does not have ``n`` ``True`` elements"""
  2945. self.assertFalse(mi.exactly_n([True, False, False], 2))
  2946. self.assertFalse(mi.exactly_n([True, True, False], 1))
  2947. self.assertFalse(mi.exactly_n([False], 1))
  2948. self.assertFalse(mi.exactly_n([True], -1))
  2949. self.assertFalse(mi.exactly_n(repeat(True), 100))
  2950. def test_empty(self):
  2951. """Return ``True`` if the iterable is empty and ``n`` is 0"""
  2952. self.assertTrue(mi.exactly_n([], 0))
  2953. self.assertFalse(mi.exactly_n([], 1))
  2954. class AlwaysReversibleTests(TestCase):
  2955. """Tests for ``always_reversible()``"""
  2956. def test_regular_reversed(self):
  2957. self.assertEqual(
  2958. list(reversed(range(10))), list(mi.always_reversible(range(10)))
  2959. )
  2960. self.assertEqual(
  2961. list(reversed([1, 2, 3])), list(mi.always_reversible([1, 2, 3]))
  2962. )
  2963. self.assertEqual(
  2964. reversed([1, 2, 3]).__class__,
  2965. mi.always_reversible([1, 2, 3]).__class__,
  2966. )
  2967. def test_nonseq_reversed(self):
  2968. # Create a non-reversible generator from a sequence
  2969. with self.assertRaises(TypeError):
  2970. reversed(x for x in range(10))
  2971. self.assertEqual(
  2972. list(reversed(range(10))),
  2973. list(mi.always_reversible(x for x in range(10))),
  2974. )
  2975. self.assertEqual(
  2976. list(reversed([1, 2, 3])),
  2977. list(mi.always_reversible(x for x in [1, 2, 3])),
  2978. )
  2979. self.assertNotEqual(
  2980. reversed((1, 2)).__class__,
  2981. mi.always_reversible(x for x in (1, 2)).__class__,
  2982. )
  2983. class CircularShiftsTests(TestCase):
  2984. def test_empty(self):
  2985. # empty iterable -> empty list
  2986. self.assertEqual(list(mi.circular_shifts([])), [])
  2987. def test_simple_circular_shifts(self):
  2988. # test the a simple iterator case
  2989. self.assertEqual(
  2990. list(mi.circular_shifts(range(4))),
  2991. [(0, 1, 2, 3), (1, 2, 3, 0), (2, 3, 0, 1), (3, 0, 1, 2)],
  2992. )
  2993. def test_duplicates(self):
  2994. # test non-distinct entries
  2995. self.assertEqual(
  2996. list(mi.circular_shifts([0, 1, 0, 1])),
  2997. [(0, 1, 0, 1), (1, 0, 1, 0), (0, 1, 0, 1), (1, 0, 1, 0)],
  2998. )
  2999. def test_steps_positive(self):
  3000. actual = list(mi.circular_shifts(range(5), steps=2))
  3001. expected = [
  3002. (0, 1, 2, 3, 4),
  3003. (2, 3, 4, 0, 1),
  3004. (4, 0, 1, 2, 3),
  3005. (1, 2, 3, 4, 0),
  3006. (3, 4, 0, 1, 2),
  3007. ]
  3008. self.assertEqual(actual, expected)
  3009. def test_steps_negative(self):
  3010. actual = list(mi.circular_shifts(range(5), steps=-2))
  3011. expected = [
  3012. (0, 1, 2, 3, 4),
  3013. (3, 4, 0, 1, 2),
  3014. (1, 2, 3, 4, 0),
  3015. (4, 0, 1, 2, 3),
  3016. (2, 3, 4, 0, 1),
  3017. ]
  3018. self.assertEqual(actual, expected)
  3019. def test_steps_zero(self):
  3020. with self.assertRaises(ValueError):
  3021. list(mi.circular_shifts(range(5), steps=0))
  3022. class MakeDecoratorTests(TestCase):
  3023. def test_basic(self):
  3024. slicer = mi.make_decorator(islice)
  3025. @slicer(1, 10, 2)
  3026. def user_function(arg_1, arg_2, kwarg_1=None):
  3027. self.assertEqual(arg_1, 'arg_1')
  3028. self.assertEqual(arg_2, 'arg_2')
  3029. self.assertEqual(kwarg_1, 'kwarg_1')
  3030. return map(str, count())
  3031. it = user_function('arg_1', 'arg_2', kwarg_1='kwarg_1')
  3032. actual = list(it)
  3033. expected = ['1', '3', '5', '7', '9']
  3034. self.assertEqual(actual, expected)
  3035. def test_result_index(self):
  3036. def stringify(*args, **kwargs):
  3037. self.assertEqual(args[0], 'arg_0')
  3038. iterable = args[1]
  3039. self.assertEqual(args[2], 'arg_2')
  3040. self.assertEqual(kwargs['kwarg_1'], 'kwarg_1')
  3041. return map(str, iterable)
  3042. stringifier = mi.make_decorator(stringify, result_index=1)
  3043. @stringifier('arg_0', 'arg_2', kwarg_1='kwarg_1')
  3044. def user_function(n):
  3045. return count(n)
  3046. it = user_function(1)
  3047. actual = mi.take(5, it)
  3048. expected = ['1', '2', '3', '4', '5']
  3049. self.assertEqual(actual, expected)
  3050. def test_wrap_class(self):
  3051. seeker = mi.make_decorator(mi.seekable)
  3052. @seeker()
  3053. def user_function(n):
  3054. return map(str, range(n))
  3055. it = user_function(5)
  3056. self.assertEqual(list(it), ['0', '1', '2', '3', '4'])
  3057. it.seek(0)
  3058. self.assertEqual(list(it), ['0', '1', '2', '3', '4'])
  3059. class MapReduceTests(TestCase):
  3060. def test_default(self):
  3061. iterable = (str(x) for x in range(5))
  3062. keyfunc = lambda x: int(x) // 2
  3063. actual = sorted(mi.map_reduce(iterable, keyfunc).items())
  3064. expected = [(0, ['0', '1']), (1, ['2', '3']), (2, ['4'])]
  3065. self.assertEqual(actual, expected)
  3066. def test_valuefunc(self):
  3067. iterable = (str(x) for x in range(5))
  3068. keyfunc = lambda x: int(x) // 2
  3069. valuefunc = int
  3070. actual = sorted(mi.map_reduce(iterable, keyfunc, valuefunc).items())
  3071. expected = [(0, [0, 1]), (1, [2, 3]), (2, [4])]
  3072. self.assertEqual(actual, expected)
  3073. def test_reducefunc(self):
  3074. iterable = (str(x) for x in range(5))
  3075. keyfunc = lambda x: int(x) // 2
  3076. valuefunc = int
  3077. reducefunc = lambda value_list: reduce(mul, value_list, 1)
  3078. actual = sorted(
  3079. mi.map_reduce(iterable, keyfunc, valuefunc, reducefunc).items()
  3080. )
  3081. expected = [(0, 0), (1, 6), (2, 4)]
  3082. self.assertEqual(actual, expected)
  3083. def test_ret(self):
  3084. d = mi.map_reduce([1, 0, 2, 0, 1, 0], bool)
  3085. self.assertEqual(d, {False: [0, 0, 0], True: [1, 2, 1]})
  3086. self.assertRaises(KeyError, lambda: d[None].append(1))
  3087. class RlocateTests(TestCase):
  3088. def test_default_pred(self):
  3089. iterable = [0, 1, 1, 0, 1, 0, 0]
  3090. for it in (iterable[:], iter(iterable)):
  3091. actual = list(mi.rlocate(it))
  3092. expected = [4, 2, 1]
  3093. self.assertEqual(actual, expected)
  3094. def test_no_matches(self):
  3095. iterable = [0, 0, 0]
  3096. for it in (iterable[:], iter(iterable)):
  3097. actual = list(mi.rlocate(it))
  3098. expected = []
  3099. self.assertEqual(actual, expected)
  3100. def test_custom_pred(self):
  3101. iterable = ['0', 1, 1, '0', 1, '0', '0']
  3102. pred = lambda x: x == '0'
  3103. for it in (iterable[:], iter(iterable)):
  3104. actual = list(mi.rlocate(it, pred))
  3105. expected = [6, 5, 3, 0]
  3106. self.assertEqual(actual, expected)
  3107. def test_efficient_reversal(self):
  3108. iterable = range(9**9) # Is efficiently reversible
  3109. target = 9**9 - 2
  3110. pred = lambda x: x == target # Find-able from the right
  3111. actual = next(mi.rlocate(iterable, pred))
  3112. self.assertEqual(actual, target)
  3113. def test_window_size(self):
  3114. iterable = ['0', 1, 1, '0', 1, '0', '0']
  3115. pred = lambda *args: args == ('0', 1)
  3116. for it in (iterable, iter(iterable)):
  3117. actual = list(mi.rlocate(it, pred, window_size=2))
  3118. expected = [3, 0]
  3119. self.assertEqual(actual, expected)
  3120. def test_window_size_large(self):
  3121. iterable = [1, 2, 3, 4]
  3122. pred = lambda a, b, c, d, e: True
  3123. for it in (iterable, iter(iterable)):
  3124. actual = list(mi.rlocate(iterable, pred, window_size=5))
  3125. expected = [0]
  3126. self.assertEqual(actual, expected)
  3127. def test_window_size_zero(self):
  3128. iterable = [1, 2, 3, 4]
  3129. pred = lambda: True
  3130. for it in (iterable, iter(iterable)):
  3131. with self.assertRaises(ValueError):
  3132. list(mi.locate(iterable, pred, window_size=0))
  3133. class ReplaceTests(TestCase):
  3134. def test_basic(self):
  3135. iterable = range(10)
  3136. pred = lambda x: x % 2 == 0
  3137. substitutes = []
  3138. actual = list(mi.replace(iterable, pred, substitutes))
  3139. expected = [1, 3, 5, 7, 9]
  3140. self.assertEqual(actual, expected)
  3141. def test_count(self):
  3142. iterable = range(10)
  3143. pred = lambda x: x % 2 == 0
  3144. substitutes = []
  3145. actual = list(mi.replace(iterable, pred, substitutes, count=4))
  3146. expected = [1, 3, 5, 7, 8, 9]
  3147. self.assertEqual(actual, expected)
  3148. def test_window_size(self):
  3149. iterable = range(10)
  3150. pred = lambda *args: args == (0, 1, 2)
  3151. substitutes = []
  3152. actual = list(mi.replace(iterable, pred, substitutes, window_size=3))
  3153. expected = [3, 4, 5, 6, 7, 8, 9]
  3154. self.assertEqual(actual, expected)
  3155. def test_window_size_end(self):
  3156. iterable = range(10)
  3157. pred = lambda *args: args == (7, 8, 9)
  3158. substitutes = []
  3159. actual = list(mi.replace(iterable, pred, substitutes, window_size=3))
  3160. expected = [0, 1, 2, 3, 4, 5, 6]
  3161. self.assertEqual(actual, expected)
  3162. def test_window_size_count(self):
  3163. iterable = range(10)
  3164. pred = lambda *args: (args == (0, 1, 2)) or (args == (7, 8, 9))
  3165. substitutes = []
  3166. actual = list(
  3167. mi.replace(iterable, pred, substitutes, count=1, window_size=3)
  3168. )
  3169. expected = [3, 4, 5, 6, 7, 8, 9]
  3170. self.assertEqual(actual, expected)
  3171. def test_window_size_large(self):
  3172. iterable = range(4)
  3173. pred = lambda a, b, c, d, e: True
  3174. substitutes = [5, 6, 7]
  3175. actual = list(mi.replace(iterable, pred, substitutes, window_size=5))
  3176. expected = [5, 6, 7]
  3177. self.assertEqual(actual, expected)
  3178. def test_window_size_zero(self):
  3179. iterable = range(10)
  3180. pred = lambda *args: True
  3181. substitutes = []
  3182. with self.assertRaises(ValueError):
  3183. list(mi.replace(iterable, pred, substitutes, window_size=0))
  3184. def test_iterable_substitutes(self):
  3185. iterable = range(5)
  3186. pred = lambda x: x % 2 == 0
  3187. substitutes = iter('__')
  3188. actual = list(mi.replace(iterable, pred, substitutes))
  3189. expected = ['_', '_', 1, '_', '_', 3, '_', '_']
  3190. self.assertEqual(actual, expected)
  3191. class PartitionsTest(TestCase):
  3192. def test_types(self):
  3193. for iterable in ['abcd', ['a', 'b', 'c', 'd'], ('a', 'b', 'c', 'd')]:
  3194. with self.subTest(iterable=iterable):
  3195. actual = list(mi.partitions(iterable))
  3196. expected = [
  3197. [['a', 'b', 'c', 'd']],
  3198. [['a'], ['b', 'c', 'd']],
  3199. [['a', 'b'], ['c', 'd']],
  3200. [['a', 'b', 'c'], ['d']],
  3201. [['a'], ['b'], ['c', 'd']],
  3202. [['a'], ['b', 'c'], ['d']],
  3203. [['a', 'b'], ['c'], ['d']],
  3204. [['a'], ['b'], ['c'], ['d']],
  3205. ]
  3206. self.assertEqual(actual, expected)
  3207. def test_empty(self):
  3208. iterable = []
  3209. actual = list(mi.partitions(iterable))
  3210. expected = [[[]]]
  3211. self.assertEqual(actual, expected)
  3212. def test_order(self):
  3213. iterable = iter([3, 2, 1])
  3214. actual = list(mi.partitions(iterable))
  3215. expected = [[[3, 2, 1]], [[3], [2, 1]], [[3, 2], [1]], [[3], [2], [1]]]
  3216. self.assertEqual(actual, expected)
  3217. def test_duplicates(self):
  3218. iterable = [1, 1, 1]
  3219. actual = list(mi.partitions(iterable))
  3220. expected = [[[1, 1, 1]], [[1], [1, 1]], [[1, 1], [1]], [[1], [1], [1]]]
  3221. self.assertEqual(actual, expected)
  3222. class _FrozenMultiset(Set):
  3223. """
  3224. A helper class, useful to compare two lists without reference to the order
  3225. of elements.
  3226. FrozenMultiset represents a hashable set that allows duplicate elements.
  3227. """
  3228. def __init__(self, iterable):
  3229. self._collection = frozenset(Counter(iterable).items())
  3230. def __contains__(self, y):
  3231. """
  3232. >>> (0, 1) in _FrozenMultiset([(0, 1), (2,), (0, 1)])
  3233. True
  3234. """
  3235. return any(y == x for x, _ in self._collection)
  3236. def __iter__(self):
  3237. """
  3238. >>> sorted(_FrozenMultiset([(0, 1), (2,), (0, 1)]))
  3239. [(0, 1), (0, 1), (2,)]
  3240. """
  3241. return (x for x, c in self._collection for _ in range(c))
  3242. def __len__(self):
  3243. """
  3244. >>> len(_FrozenMultiset([(0, 1), (2,), (0, 1)]))
  3245. 3
  3246. """
  3247. return sum(c for x, c in self._collection)
  3248. def has_duplicates(self):
  3249. """
  3250. >>> _FrozenMultiset([(0, 1), (2,), (0, 1)]).has_duplicates()
  3251. True
  3252. """
  3253. return any(c != 1 for _, c in self._collection)
  3254. def __hash__(self):
  3255. return hash(self._collection)
  3256. def __repr__(self):
  3257. return "FrozenSet([{}]".format(", ".join(repr(x) for x in iter(self)))
  3258. class SetPartitionsTests(TestCase):
  3259. @staticmethod
  3260. def _normalize_partition(p):
  3261. """
  3262. Return a normalized, hashable, version of a partition using
  3263. _FrozenMultiset
  3264. """
  3265. return _FrozenMultiset(_FrozenMultiset(g) for g in p)
  3266. @staticmethod
  3267. def _normalize_partitions(ps):
  3268. """
  3269. Return a normalized set of all normalized partitions using
  3270. _FrozenMultiset
  3271. """
  3272. return _FrozenMultiset(
  3273. SetPartitionsTests._normalize_partition(p) for p in ps
  3274. )
  3275. def test_repeated(self):
  3276. it = 'aaa'
  3277. actual = mi.set_partitions(it, 2)
  3278. expected = [['a', 'aa'], ['a', 'aa'], ['a', 'aa']]
  3279. self.assertEqual(
  3280. self._normalize_partitions(expected),
  3281. self._normalize_partitions(actual),
  3282. )
  3283. def test_each_correct(self):
  3284. a = set(range(6))
  3285. for p in mi.set_partitions(a):
  3286. total = {e for g in p for e in g}
  3287. self.assertEqual(a, total)
  3288. def test_duplicates(self):
  3289. a = set(range(6))
  3290. for p in mi.set_partitions(a):
  3291. self.assertFalse(self._normalize_partition(p).has_duplicates())
  3292. def test_found_all(self):
  3293. """small example, hand-checked"""
  3294. expected = [
  3295. [[0], [1], [2, 3, 4]],
  3296. [[0], [1, 2], [3, 4]],
  3297. [[0], [2], [1, 3, 4]],
  3298. [[0], [3], [1, 2, 4]],
  3299. [[0], [4], [1, 2, 3]],
  3300. [[0], [1, 3], [2, 4]],
  3301. [[0], [1, 4], [2, 3]],
  3302. [[1], [2], [0, 3, 4]],
  3303. [[1], [3], [0, 2, 4]],
  3304. [[1], [4], [0, 2, 3]],
  3305. [[1], [0, 2], [3, 4]],
  3306. [[1], [0, 3], [2, 4]],
  3307. [[1], [0, 4], [2, 3]],
  3308. [[2], [3], [0, 1, 4]],
  3309. [[2], [4], [0, 1, 3]],
  3310. [[2], [0, 1], [3, 4]],
  3311. [[2], [0, 3], [1, 4]],
  3312. [[2], [0, 4], [1, 3]],
  3313. [[3], [4], [0, 1, 2]],
  3314. [[3], [0, 1], [2, 4]],
  3315. [[3], [0, 2], [1, 4]],
  3316. [[3], [0, 4], [1, 2]],
  3317. [[4], [0, 1], [2, 3]],
  3318. [[4], [0, 2], [1, 3]],
  3319. [[4], [0, 3], [1, 2]],
  3320. ]
  3321. actual = mi.set_partitions(range(5), 3)
  3322. self.assertEqual(
  3323. self._normalize_partitions(expected),
  3324. self._normalize_partitions(actual),
  3325. )
  3326. def test_stirling_numbers(self):
  3327. """Check against https://en.wikipedia.org/wiki/
  3328. Stirling_numbers_of_the_second_kind#Table_of_values"""
  3329. cardinality_by_k_by_n = [
  3330. [1],
  3331. [1, 1],
  3332. [1, 3, 1],
  3333. [1, 7, 6, 1],
  3334. [1, 15, 25, 10, 1],
  3335. [1, 31, 90, 65, 15, 1],
  3336. ]
  3337. for n, cardinality_by_k in enumerate(cardinality_by_k_by_n, 1):
  3338. for k, cardinality in enumerate(cardinality_by_k, 1):
  3339. self.assertEqual(
  3340. cardinality, len(list(mi.set_partitions(range(n), k)))
  3341. )
  3342. def test_no_group(self):
  3343. def helper():
  3344. list(mi.set_partitions(range(4), -1))
  3345. self.assertRaises(ValueError, helper)
  3346. def test_to_many_groups(self):
  3347. self.assertEqual([], list(mi.set_partitions(range(4), 5)))
  3348. def test_min_size(self):
  3349. it = 'abc'
  3350. actual = mi.set_partitions(it, min_size=2)
  3351. expected = [['abc']]
  3352. self.assertEqual(
  3353. self._normalize_partitions(expected),
  3354. self._normalize_partitions(actual),
  3355. )
  3356. def test_max_size(self):
  3357. it = 'abc'
  3358. actual = mi.set_partitions(it, max_size=2)
  3359. expected = [['a', 'bc'], ['ab', 'c'], ['b', 'ac'], ['a', 'b', 'c']]
  3360. self.assertEqual(
  3361. self._normalize_partitions(expected),
  3362. self._normalize_partitions(actual),
  3363. )
  3364. class TimeLimitedTests(TestCase):
  3365. def test_basic(self):
  3366. def generator():
  3367. yield 1
  3368. yield 2
  3369. sleep(0.2)
  3370. yield 3
  3371. iterable = mi.time_limited(0.1, generator())
  3372. actual = list(iterable)
  3373. expected = [1, 2]
  3374. self.assertEqual(actual, expected)
  3375. self.assertTrue(iterable.timed_out)
  3376. def test_complete(self):
  3377. iterable = mi.time_limited(2, iter(range(10)))
  3378. actual = list(iterable)
  3379. expected = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
  3380. self.assertEqual(actual, expected)
  3381. self.assertFalse(iterable.timed_out)
  3382. def test_zero_limit(self):
  3383. iterable = mi.time_limited(0, count())
  3384. actual = list(iterable)
  3385. expected = []
  3386. self.assertEqual(actual, expected)
  3387. self.assertTrue(iterable.timed_out)
  3388. def test_invalid_limit(self):
  3389. with self.assertRaises(ValueError):
  3390. list(mi.time_limited(-0.1, count()))
  3391. class OnlyTests(TestCase):
  3392. def test_defaults(self):
  3393. self.assertEqual(mi.only([]), None)
  3394. self.assertEqual(mi.only([1]), 1)
  3395. self.assertRaises(ValueError, lambda: mi.only([1, 2]))
  3396. def test_custom_value(self):
  3397. self.assertEqual(mi.only([], default='!'), '!')
  3398. self.assertEqual(mi.only([1], default='!'), 1)
  3399. self.assertRaises(ValueError, lambda: mi.only([1, 2], default='!'))
  3400. def test_custom_exception(self):
  3401. self.assertEqual(mi.only([], too_long=RuntimeError), None)
  3402. self.assertEqual(mi.only([1], too_long=RuntimeError), 1)
  3403. self.assertRaises(
  3404. RuntimeError, lambda: mi.only([1, 2], too_long=RuntimeError)
  3405. )
  3406. def test_default_exception_message(self):
  3407. self.assertRaisesRegex(
  3408. ValueError,
  3409. "Expected exactly one item in iterable, "
  3410. "but got 'foo', 'bar', and perhaps more",
  3411. lambda: mi.only(['foo', 'bar', 'baz']),
  3412. )
  3413. class IchunkedTests(TestCase):
  3414. def test_even(self):
  3415. iterable = (str(x) for x in range(10))
  3416. actual = [''.join(c) for c in mi.ichunked(iterable, 5)]
  3417. expected = ['01234', '56789']
  3418. self.assertEqual(actual, expected)
  3419. def test_odd(self):
  3420. iterable = (str(x) for x in range(10))
  3421. actual = [''.join(c) for c in mi.ichunked(iterable, 4)]
  3422. expected = ['0123', '4567', '89']
  3423. self.assertEqual(actual, expected)
  3424. def test_zero(self):
  3425. iterable = []
  3426. actual = [list(c) for c in mi.ichunked(iterable, 0)]
  3427. expected = []
  3428. self.assertEqual(actual, expected)
  3429. def test_negative(self):
  3430. iterable = count()
  3431. with self.assertRaises(ValueError):
  3432. [list(c) for c in mi.ichunked(iterable, -1)]
  3433. def test_out_of_order(self):
  3434. iterable = map(str, count())
  3435. it = mi.ichunked(iterable, 4)
  3436. chunk_1 = next(it)
  3437. chunk_2 = next(it)
  3438. self.assertEqual(''.join(chunk_2), '4567')
  3439. self.assertEqual(''.join(chunk_1), '0123')
  3440. def test_laziness(self):
  3441. def gen():
  3442. yield 0
  3443. raise RuntimeError
  3444. yield from count(1)
  3445. it = mi.ichunked(gen(), 4)
  3446. chunk = next(it)
  3447. self.assertEqual(next(chunk), 0)
  3448. self.assertRaises(RuntimeError, next, it)
  3449. def test_memory_in_order(self):
  3450. gen_numbers = []
  3451. def gen():
  3452. for gen_number in count():
  3453. gen_numbers.append(gen_number)
  3454. yield gen_number
  3455. # No items should be kept in memory when a ichunked is first called
  3456. all_chunks = mi.ichunked(gen(), 4)
  3457. self.assertEqual(gen_numbers, [])
  3458. # The first item of each chunk should be generated on chunk generation
  3459. first_chunk = next(all_chunks)
  3460. self.assertEqual(gen_numbers, [0])
  3461. # If we don't read a chunk before getting its successor, its contents
  3462. # will be cached
  3463. second_chunk = next(all_chunks)
  3464. self.assertEqual(gen_numbers, [0, 1, 2, 3, 4])
  3465. # Check if we can read in cached values
  3466. self.assertEqual(list(first_chunk), [0, 1, 2, 3])
  3467. self.assertEqual(list(second_chunk), [4, 5, 6, 7])
  3468. # Again only the most recent chunk should have an item cached
  3469. third_chunk = next(all_chunks)
  3470. self.assertEqual(len(gen_numbers), 9)
  3471. # No new item should be cached when reading past the first number
  3472. next(third_chunk)
  3473. self.assertEqual(len(gen_numbers), 9)
  3474. # we should not be able to read spent chunks
  3475. self.assertEqual(list(first_chunk), [])
  3476. self.assertEqual(list(second_chunk), [])
  3477. class DistinctCombinationsTests(TestCase):
  3478. def test_basic(self):
  3479. for iterable in [
  3480. (1, 2, 2, 3, 3, 3), # In order
  3481. range(6), # All distinct
  3482. 'abbccc', # Not numbers
  3483. 'cccbba', # Backward
  3484. 'mississippi', # No particular order
  3485. ]:
  3486. for r in range(len(iterable)):
  3487. with self.subTest(iterable=iterable, r=r):
  3488. actual = list(mi.distinct_combinations(iterable, r))
  3489. expected = list(
  3490. mi.unique_everseen(combinations(iterable, r))
  3491. )
  3492. self.assertEqual(actual, expected)
  3493. def test_negative(self):
  3494. with self.assertRaises(ValueError):
  3495. list(mi.distinct_combinations([], -1))
  3496. def test_empty(self):
  3497. self.assertEqual(list(mi.distinct_combinations([], 2)), [])
  3498. class FilterExceptTests(TestCase):
  3499. def test_no_exceptions_pass(self):
  3500. iterable = '0123'
  3501. actual = list(mi.filter_except(int, iterable))
  3502. expected = ['0', '1', '2', '3']
  3503. self.assertEqual(actual, expected)
  3504. def test_no_exceptions_raise(self):
  3505. iterable = ['0', '1', 'two', '3']
  3506. with self.assertRaises(ValueError):
  3507. list(mi.filter_except(int, iterable))
  3508. def test_raise(self):
  3509. iterable = ['0', '1' '2', 'three', None]
  3510. with self.assertRaises(TypeError):
  3511. list(mi.filter_except(int, iterable, ValueError))
  3512. def test_false(self):
  3513. # Even if the validator returns false, we pass through
  3514. validator = lambda x: False
  3515. iterable = ['0', '1', '2', 'three', None]
  3516. actual = list(mi.filter_except(validator, iterable, Exception))
  3517. expected = ['0', '1', '2', 'three', None]
  3518. self.assertEqual(actual, expected)
  3519. def test_multiple(self):
  3520. iterable = ['0', '1', '2', 'three', None, '4']
  3521. actual = list(mi.filter_except(int, iterable, ValueError, TypeError))
  3522. expected = ['0', '1', '2', '4']
  3523. self.assertEqual(actual, expected)
  3524. class MapExceptTests(TestCase):
  3525. def test_no_exceptions_pass(self):
  3526. iterable = '0123'
  3527. actual = list(mi.map_except(int, iterable))
  3528. expected = [0, 1, 2, 3]
  3529. self.assertEqual(actual, expected)
  3530. def test_no_exceptions_raise(self):
  3531. iterable = ['0', '1', 'two', '3']
  3532. with self.assertRaises(ValueError):
  3533. list(mi.map_except(int, iterable))
  3534. def test_raise(self):
  3535. iterable = ['0', '1' '2', 'three', None]
  3536. with self.assertRaises(TypeError):
  3537. list(mi.map_except(int, iterable, ValueError))
  3538. def test_multiple(self):
  3539. iterable = ['0', '1', '2', 'three', None, '4']
  3540. actual = list(mi.map_except(int, iterable, ValueError, TypeError))
  3541. expected = [0, 1, 2, 4]
  3542. self.assertEqual(actual, expected)
  3543. class MapIfTests(TestCase):
  3544. def test_without_func_else(self):
  3545. iterable = list(range(-5, 5))
  3546. actual = list(mi.map_if(iterable, lambda x: x > 3, lambda x: 'toobig'))
  3547. expected = [-5, -4, -3, -2, -1, 0, 1, 2, 3, 'toobig']
  3548. self.assertEqual(actual, expected)
  3549. def test_with_func_else(self):
  3550. iterable = list(range(-5, 5))
  3551. actual = list(
  3552. mi.map_if(
  3553. iterable, lambda x: x >= 0, lambda x: 'notneg', lambda x: 'neg'
  3554. )
  3555. )
  3556. expected = ['neg'] * 5 + ['notneg'] * 5
  3557. self.assertEqual(actual, expected)
  3558. def test_empty(self):
  3559. actual = list(mi.map_if([], lambda x: len(x) > 5, lambda x: None))
  3560. expected = []
  3561. self.assertEqual(actual, expected)
  3562. class SampleTests(TestCase):
  3563. def test_unit_case(self):
  3564. """Test against a fixed case by seeding the random module."""
  3565. # Beware that this test really just verifies random.random() behavior.
  3566. # If the algorithm is changed (e.g. to a more naive implementation)
  3567. # this test will fail, but the algorithm might be correct.
  3568. # Also, this test can pass and the algorithm can be completely wrong.
  3569. data = "abcdef"
  3570. weights = list(range(1, len(data) + 1))
  3571. seed(123)
  3572. actual = mi.sample(data, k=2, weights=weights)
  3573. expected = ['f', 'e']
  3574. self.assertEqual(actual, expected)
  3575. def test_negative(self):
  3576. data = [1, 2, 3, 4, 5]
  3577. with self.assertRaises(ValueError):
  3578. mi.sample(data, k=-1)
  3579. def test_length(self):
  3580. """Check that *k* elements are sampled."""
  3581. data = [1, 2, 3, 4, 5]
  3582. for k in [0, 3, 5, 7]:
  3583. sampled = mi.sample(data, k=k)
  3584. actual = len(sampled)
  3585. expected = min(k, len(data))
  3586. self.assertEqual(actual, expected)
  3587. def test_strict(self):
  3588. data = ['1', '2', '3', '4', '5']
  3589. self.assertEqual(set(mi.sample(data, 6, strict=False)), set(data))
  3590. with self.assertRaises(ValueError):
  3591. mi.sample(data, 6, strict=True)
  3592. def test_counts(self):
  3593. # Test with counts
  3594. seed(0)
  3595. iterable = ['red', 'blue']
  3596. counts = [4, 2]
  3597. k = 5
  3598. actual = list(mi.sample(iterable, counts=counts, k=k))
  3599. # Test without counts
  3600. seed(0)
  3601. decoded_iterable = (['red'] * 4) + (['blue'] * 2)
  3602. expected = list(mi.sample(decoded_iterable, k=k))
  3603. self.assertEqual(actual, expected)
  3604. def test_counts_all(self):
  3605. actual = Counter(mi.sample('uwxyz', 35, counts=(1, 0, 4, 10, 20)))
  3606. expected = Counter({'u': 1, 'x': 4, 'y': 10, 'z': 20})
  3607. self.assertEqual(actual, expected)
  3608. def test_sampling_entire_iterable(self):
  3609. """If k=len(iterable), the sample contains the original elements."""
  3610. data = ["a", 2, "a", 4, (1, 2, 3)]
  3611. actual = set(mi.sample(data, k=len(data)))
  3612. expected = set(data)
  3613. self.assertEqual(actual, expected)
  3614. def test_scale_invariance_of_weights(self):
  3615. """The probability of choosing element a_i is w_i / sum(weights).
  3616. Scaling weights should not change the probability or outcome."""
  3617. data = "abcdef"
  3618. weights = list(range(1, len(data) + 1))
  3619. seed(123)
  3620. first_sample = mi.sample(data, k=2, weights=weights)
  3621. # Scale the weights and sample again
  3622. weights_scaled = [w / 1e10 for w in weights]
  3623. seed(123)
  3624. second_sample = mi.sample(data, k=2, weights=weights_scaled)
  3625. self.assertEqual(first_sample, second_sample)
  3626. def test_invariance_under_permutations_unweighted(self):
  3627. """The order of the data should not matter. This is a stochastic test,
  3628. but it will fail in less than 1 / 10_000 cases."""
  3629. # Create a data set and a reversed data set
  3630. data = list(range(100))
  3631. data_rev = list(reversed(data))
  3632. # Sample each data set 10 times
  3633. data_means = [mean(mi.sample(data, k=50)) for _ in range(10)]
  3634. data_rev_means = [mean(mi.sample(data_rev, k=50)) for _ in range(10)]
  3635. # The difference in the means should be low, i.e. little bias
  3636. difference_in_means = abs(mean(data_means) - mean(data_rev_means))
  3637. # The observed largest difference in 10,000 simulations was 5.09599
  3638. self.assertTrue(difference_in_means < 5.1)
  3639. def test_invariance_under_permutations_weighted(self):
  3640. """The order of the data should not matter. This is a stochastic test,
  3641. but it will fail in less than 1 / 10_000 cases."""
  3642. # Create a data set and a reversed data set
  3643. data = list(range(1, 101))
  3644. data_rev = list(reversed(data))
  3645. # Sample each data set 10 times
  3646. data_means = [
  3647. mean(mi.sample(data, k=50, weights=data)) for _ in range(10)
  3648. ]
  3649. data_rev_means = [
  3650. mean(mi.sample(data_rev, k=50, weights=data_rev))
  3651. for _ in range(10)
  3652. ]
  3653. # The difference in the means should be low, i.e. little bias
  3654. difference_in_means = abs(mean(data_means) - mean(data_rev_means))
  3655. # The observed largest difference in 10,000 simulations was 4.337999
  3656. self.assertTrue(difference_in_means < 4.4)
  3657. class IsSortedTests(TestCase):
  3658. def test_basic(self):
  3659. for iterable, kwargs, expected in [
  3660. ([], {}, True),
  3661. ([1], {}, True),
  3662. ([1, 2, 3], {}, True),
  3663. ([1, 1, 2, 3], {}, True),
  3664. ([1, 10, 2, 3], {}, False),
  3665. ([3, float('nan'), 1, 2], {}, True),
  3666. (['1', '10', '2', '3'], {}, True),
  3667. (['1', '10', '2', '3'], {'key': int}, False),
  3668. ([1, 2, 3], {'reverse': True}, False),
  3669. ([1, 1, 2, 3], {'reverse': True}, False),
  3670. ([1, 10, 2, 3], {'reverse': True}, False),
  3671. (['3', '2', '10', '1'], {'reverse': True}, True),
  3672. (['3', '2', '10', '1'], {'key': int, 'reverse': True}, False),
  3673. # strict
  3674. ([], {'strict': True}, True),
  3675. ([1], {'strict': True}, True),
  3676. ([1, 1], {'strict': True}, False),
  3677. ([1, 2, 3], {'strict': True}, True),
  3678. ([1, 1, 2, 3], {'strict': True}, False),
  3679. ([1, 10, 2, 3], {'strict': True}, False),
  3680. (['1', '10', '2', '3'], {'strict': True}, True),
  3681. (['1', '10', '2', '3', '3'], {'strict': True}, False),
  3682. (['1', '10', '2', '3'], {'strict': True, 'key': int}, False),
  3683. ([1, 2, 3], {'strict': True, 'reverse': True}, False),
  3684. ([1, 1, 2, 3], {'strict': True, 'reverse': True}, False),
  3685. ([1, 10, 2, 3], {'strict': True, 'reverse': True}, False),
  3686. (['3', '2', '10', '1'], {'strict': True, 'reverse': True}, True),
  3687. (
  3688. ['3', '2', '10', '10', '1'],
  3689. {'strict': True, 'reverse': True},
  3690. False,
  3691. ),
  3692. (
  3693. ['3', '2', '10', '1'],
  3694. {'strict': True, 'key': int, 'reverse': True},
  3695. False,
  3696. ),
  3697. # We'll do the same weird thing as Python here
  3698. (['nan', 0, 'nan', 0], {'key': float}, True),
  3699. ([0, 'nan', 0, 'nan'], {'key': float}, True),
  3700. (['nan', 0, 'nan', 0], {'key': float, 'reverse': True}, True),
  3701. ([0, 'nan', 0, 'nan'], {'key': float, 'reverse': True}, True),
  3702. ([0, 'nan', 0, 'nan'], {'strict': True, 'key': float}, True),
  3703. (
  3704. ['nan', 0, 'nan', 0],
  3705. {'strict': True, 'key': float, 'reverse': True},
  3706. True,
  3707. ),
  3708. ]:
  3709. key = kwargs.get('key', None)
  3710. reverse = kwargs.get('reverse', False)
  3711. strict = kwargs.get('strict', False)
  3712. with self.subTest(
  3713. iterable=iterable, key=key, reverse=reverse, strict=strict
  3714. ):
  3715. mi_result = mi.is_sorted(
  3716. iter(iterable), key=key, reverse=reverse, strict=strict
  3717. )
  3718. sorted_iterable = sorted(iterable, key=key, reverse=reverse)
  3719. if strict:
  3720. sorted_iterable = list(mi.unique_justseen(sorted_iterable))
  3721. py_result = iterable == sorted_iterable
  3722. self.assertEqual(mi_result, expected)
  3723. self.assertEqual(mi_result, py_result)
  3724. class CallbackIterTests(TestCase):
  3725. def _target(self, cb=None, exc=None, wait=0):
  3726. total = 0
  3727. for i, c in enumerate('abc', 1):
  3728. total += i
  3729. if wait:
  3730. sleep(wait)
  3731. if cb:
  3732. cb(i, c, intermediate_total=total)
  3733. if exc:
  3734. raise exc('error in target')
  3735. return total
  3736. def test_basic(self):
  3737. func = lambda callback=None: self._target(cb=callback, wait=0.02)
  3738. with mi.callback_iter(func, wait_seconds=0.01) as it:
  3739. # Execution doesn't start until we begin iterating
  3740. self.assertFalse(it.done)
  3741. # Consume everything
  3742. self.assertEqual(
  3743. list(it),
  3744. [
  3745. ((1, 'a'), {'intermediate_total': 1}),
  3746. ((2, 'b'), {'intermediate_total': 3}),
  3747. ((3, 'c'), {'intermediate_total': 6}),
  3748. ],
  3749. )
  3750. # After consuming everything the future is done and the
  3751. # result is available.
  3752. self.assertTrue(it.done)
  3753. self.assertEqual(it.result, 6)
  3754. # This examines the internal state of the ThreadPoolExecutor. This
  3755. # isn't documented, so may break in future Python versions.
  3756. self.assertTrue(it._executor._shutdown)
  3757. def test_callback_kwd(self):
  3758. with mi.callback_iter(self._target, callback_kwd='cb') as it:
  3759. self.assertEqual(
  3760. list(it),
  3761. [
  3762. ((1, 'a'), {'intermediate_total': 1}),
  3763. ((2, 'b'), {'intermediate_total': 3}),
  3764. ((3, 'c'), {'intermediate_total': 6}),
  3765. ],
  3766. )
  3767. def test_partial_consumption(self):
  3768. func = lambda callback=None: self._target(cb=callback)
  3769. with mi.callback_iter(func) as it:
  3770. self.assertEqual(next(it), ((1, 'a'), {'intermediate_total': 1}))
  3771. self.assertTrue(it._executor._shutdown)
  3772. def test_abort(self):
  3773. func = lambda callback=None: self._target(cb=callback, wait=0.1)
  3774. with mi.callback_iter(func) as it:
  3775. self.assertEqual(next(it), ((1, 'a'), {'intermediate_total': 1}))
  3776. with self.assertRaises(mi.AbortThread):
  3777. it.result
  3778. def test_no_result(self):
  3779. func = lambda callback=None: self._target(cb=callback)
  3780. with mi.callback_iter(func) as it:
  3781. with self.assertRaises(RuntimeError):
  3782. it.result
  3783. def test_exception(self):
  3784. func = lambda callback=None: self._target(cb=callback, exc=ValueError)
  3785. with mi.callback_iter(func) as it:
  3786. self.assertEqual(
  3787. next(it),
  3788. ((1, 'a'), {'intermediate_total': 1}),
  3789. )
  3790. with self.assertRaises(ValueError):
  3791. it.result
  3792. class WindowedCompleteTests(TestCase):
  3793. """Tests for ``windowed_complete()``"""
  3794. def test_basic(self):
  3795. actual = list(mi.windowed_complete([1, 2, 3, 4, 5], 3))
  3796. expected = [
  3797. ((), (1, 2, 3), (4, 5)),
  3798. ((1,), (2, 3, 4), (5,)),
  3799. ((1, 2), (3, 4, 5), ()),
  3800. ]
  3801. self.assertEqual(actual, expected)
  3802. def test_zero_length(self):
  3803. actual = list(mi.windowed_complete([1, 2, 3], 0))
  3804. expected = [
  3805. ((), (), (1, 2, 3)),
  3806. ((1,), (), (2, 3)),
  3807. ((1, 2), (), (3,)),
  3808. ((1, 2, 3), (), ()),
  3809. ]
  3810. self.assertEqual(actual, expected)
  3811. def test_wrong_length(self):
  3812. seq = [1, 2, 3, 4, 5]
  3813. for n in (-10, -1, len(seq) + 1, len(seq) + 10):
  3814. with self.subTest(n=n):
  3815. with self.assertRaises(ValueError):
  3816. list(mi.windowed_complete(seq, n))
  3817. def test_every_partition(self):
  3818. every_partition = lambda seq: chain(
  3819. *map(partial(mi.windowed_complete, seq), range(len(seq)))
  3820. )
  3821. seq = 'ABC'
  3822. actual = list(every_partition(seq))
  3823. expected = [
  3824. ((), (), ('A', 'B', 'C')),
  3825. (('A',), (), ('B', 'C')),
  3826. (('A', 'B'), (), ('C',)),
  3827. (('A', 'B', 'C'), (), ()),
  3828. ((), ('A',), ('B', 'C')),
  3829. (('A',), ('B',), ('C',)),
  3830. (('A', 'B'), ('C',), ()),
  3831. ((), ('A', 'B'), ('C',)),
  3832. (('A',), ('B', 'C'), ()),
  3833. ]
  3834. self.assertEqual(actual, expected)
  3835. class AllUniqueTests(TestCase):
  3836. def test_basic(self):
  3837. for iterable, expected in [
  3838. ([], True),
  3839. ([1, 2, 3], True),
  3840. ([1, 1], False),
  3841. ([1, 2, 3, 1], False),
  3842. ([1, 2, 3, '1'], True),
  3843. ]:
  3844. with self.subTest(args=(iterable,)):
  3845. self.assertEqual(mi.all_unique(iterable), expected)
  3846. def test_non_hashable(self):
  3847. self.assertEqual(mi.all_unique([[1, 2], [3, 4]]), True)
  3848. self.assertEqual(mi.all_unique([[1, 2], [3, 4], [1, 2]]), False)
  3849. def test_partially_hashable(self):
  3850. self.assertEqual(mi.all_unique([[1, 2], [3, 4], (5, 6)]), True)
  3851. self.assertEqual(
  3852. mi.all_unique([[1, 2], [3, 4], (5, 6), [1, 2]]), False
  3853. )
  3854. self.assertEqual(
  3855. mi.all_unique([[1, 2], [3, 4], (5, 6), (5, 6)]), False
  3856. )
  3857. def test_key(self):
  3858. iterable = ['A', 'B', 'C', 'b']
  3859. self.assertEqual(mi.all_unique(iterable, lambda x: x), True)
  3860. self.assertEqual(mi.all_unique(iterable, str.lower), False)
  3861. def test_infinite(self):
  3862. self.assertEqual(mi.all_unique(mi.prepend(3, count())), False)
  3863. class NthProductTests(TestCase):
  3864. def test_basic(self):
  3865. iterables = ['ab', 'cdef', 'ghi']
  3866. for index, expected in enumerate(product(*iterables)):
  3867. actual = mi.nth_product(index, *iterables)
  3868. self.assertEqual(actual, expected)
  3869. def test_long(self):
  3870. actual = mi.nth_product(1337, range(101), range(22), range(53))
  3871. expected = (1, 3, 12)
  3872. self.assertEqual(actual, expected)
  3873. def test_negative(self):
  3874. iterables = ['abc', 'de', 'fghi']
  3875. for index, expected in enumerate(product(*iterables)):
  3876. actual = mi.nth_product(index - 24, *iterables)
  3877. self.assertEqual(actual, expected)
  3878. def test_invalid_index(self):
  3879. with self.assertRaises(IndexError):
  3880. mi.nth_product(24, 'ab', 'cde', 'fghi')
  3881. class NthCombinationWithReplacementTests(TestCase):
  3882. def test_basic(self):
  3883. iterable = 'abcdefg'
  3884. r = 4
  3885. for index, expected in enumerate(
  3886. combinations_with_replacement(iterable, r)
  3887. ):
  3888. actual = mi.nth_combination_with_replacement(iterable, r, index)
  3889. self.assertEqual(actual, expected)
  3890. def test_long(self):
  3891. actual = mi.nth_combination_with_replacement(range(90), 4, 2000000)
  3892. expected = (22, 65, 68, 81)
  3893. self.assertEqual(actual, expected)
  3894. def test_invalid_r(self):
  3895. for r in (-1, 3):
  3896. with self.assertRaises(ValueError):
  3897. mi.nth_combination_with_replacement([], r, 0)
  3898. def test_invalid_index(self):
  3899. with self.assertRaises(IndexError):
  3900. mi.nth_combination_with_replacement('abcdefg', 3, -85)
  3901. class ValueChainTests(TestCase):
  3902. def test_empty(self):
  3903. actual = list(mi.value_chain())
  3904. expected = []
  3905. self.assertEqual(actual, expected)
  3906. def test_simple(self):
  3907. actual = list(mi.value_chain(1, 2.71828, False, 'foo'))
  3908. expected = [1, 2.71828, False, 'foo']
  3909. self.assertEqual(actual, expected)
  3910. def test_more(self):
  3911. actual = list(mi.value_chain(b'bar', [1, 2, 3], 4, {'key': 1}))
  3912. expected = [b'bar', 1, 2, 3, 4, 'key']
  3913. self.assertEqual(actual, expected)
  3914. def test_empty_lists(self):
  3915. actual = list(mi.value_chain(1, 2, [], [3, 4]))
  3916. expected = [1, 2, 3, 4]
  3917. self.assertEqual(actual, expected)
  3918. def test_complex(self):
  3919. obj = object()
  3920. actual = list(
  3921. mi.value_chain(
  3922. (1, (2, (3,))),
  3923. ['foo', ['bar', ['baz']], 'tic'],
  3924. {'key': {'foo': 1}},
  3925. obj,
  3926. )
  3927. )
  3928. expected = [1, (2, (3,)), 'foo', ['bar', ['baz']], 'tic', 'key', obj]
  3929. self.assertEqual(actual, expected)
  3930. class ProductIndexTests(TestCase):
  3931. def test_basic(self):
  3932. iterables = ['ab', 'cdef', 'ghi']
  3933. first_index = {}
  3934. for index, element in enumerate(product(*iterables)):
  3935. actual = mi.product_index(element, *iterables)
  3936. expected = first_index.setdefault(element, index)
  3937. self.assertEqual(actual, expected)
  3938. def test_multiplicity(self):
  3939. iterables = ['ab', 'bab', 'cab']
  3940. first_index = {}
  3941. for index, element in enumerate(product(*iterables)):
  3942. actual = mi.product_index(element, *iterables)
  3943. expected = first_index.setdefault(element, index)
  3944. self.assertEqual(actual, expected)
  3945. def test_long(self):
  3946. actual = mi.product_index((1, 3, 12), range(101), range(22), range(53))
  3947. expected = 1337
  3948. self.assertEqual(actual, expected)
  3949. def test_invalid_empty(self):
  3950. with self.assertRaises(ValueError):
  3951. mi.product_index('', 'ab', 'cde', 'fghi')
  3952. def test_invalid_small(self):
  3953. with self.assertRaises(ValueError):
  3954. mi.product_index('ac', 'ab', 'cde', 'fghi')
  3955. def test_invalid_large(self):
  3956. with self.assertRaises(ValueError):
  3957. mi.product_index('achi', 'ab', 'cde', 'fghi')
  3958. def test_invalid_match(self):
  3959. with self.assertRaises(ValueError):
  3960. mi.product_index('axf', 'ab', 'cde', 'fghi')
  3961. class CombinationIndexTests(TestCase):
  3962. def test_r_less_than_n(self):
  3963. iterable = 'abcdefg'
  3964. r = 4
  3965. first_index = {}
  3966. for index, element in enumerate(combinations(iterable, r)):
  3967. actual = mi.combination_index(element, iterable)
  3968. expected = first_index.setdefault(element, index)
  3969. self.assertEqual(actual, expected)
  3970. def test_r_equal_to_n(self):
  3971. iterable = 'abcd'
  3972. r = len(iterable)
  3973. first_index = {}
  3974. for index, element in enumerate(combinations(iterable, r=r)):
  3975. actual = mi.combination_index(element, iterable)
  3976. expected = first_index.setdefault(element, index)
  3977. self.assertEqual(actual, expected)
  3978. def test_multiplicity(self):
  3979. iterable = 'abacba'
  3980. r = 3
  3981. first_index = {}
  3982. for index, element in enumerate(combinations(iterable, r)):
  3983. actual = mi.combination_index(element, iterable)
  3984. expected = first_index.setdefault(element, index)
  3985. self.assertEqual(actual, expected)
  3986. def test_null(self):
  3987. actual = mi.combination_index(tuple(), [])
  3988. expected = 0
  3989. self.assertEqual(actual, expected)
  3990. def test_long(self):
  3991. actual = mi.combination_index((2, 12, 35, 126), range(180))
  3992. expected = 2000000
  3993. self.assertEqual(actual, expected)
  3994. def test_invalid_order(self):
  3995. with self.assertRaises(ValueError):
  3996. mi.combination_index(tuple('acb'), 'abcde')
  3997. def test_invalid_large(self):
  3998. with self.assertRaises(ValueError):
  3999. mi.combination_index(tuple('abcdefg'), 'abcdef')
  4000. def test_invalid_match(self):
  4001. with self.assertRaises(ValueError):
  4002. mi.combination_index(tuple('axe'), 'abcde')
  4003. class CombinationWithReplacementIndexTests(TestCase):
  4004. def test_r_less_than_n(self):
  4005. iterable = 'abcdefg'
  4006. r = 4
  4007. first_index = {}
  4008. for index, element in enumerate(
  4009. combinations_with_replacement(iterable, r)
  4010. ):
  4011. actual = mi.combination_with_replacement_index(element, iterable)
  4012. expected = first_index.setdefault(element, index)
  4013. self.assertEqual(actual, expected)
  4014. def test_r_equal_to_n(self):
  4015. iterable = 'abcd'
  4016. r = len(iterable)
  4017. first_index = {}
  4018. for index, element in enumerate(
  4019. combinations_with_replacement(iterable, r=r)
  4020. ):
  4021. actual = mi.combination_with_replacement_index(element, iterable)
  4022. expected = first_index.setdefault(element, index)
  4023. self.assertEqual(actual, expected)
  4024. def test_multiplicity(self):
  4025. iterable = 'abacba'
  4026. r = 3
  4027. first_index = {}
  4028. for index, element in enumerate(
  4029. combinations_with_replacement(iterable, r)
  4030. ):
  4031. actual = mi.combination_with_replacement_index(element, iterable)
  4032. expected = first_index.setdefault(element, index)
  4033. self.assertEqual(actual, expected)
  4034. def test_null(self):
  4035. actual = mi.combination_with_replacement_index(tuple(), [])
  4036. expected = 0
  4037. self.assertEqual(actual, expected)
  4038. def test_long(self):
  4039. actual = mi.combination_with_replacement_index(
  4040. (22, 65, 68, 81), range(90)
  4041. )
  4042. expected = 2000000
  4043. self.assertEqual(actual, expected)
  4044. def test_invalid_order(self):
  4045. with self.assertRaises(ValueError):
  4046. mi.combination_with_replacement_index(tuple('acb'), 'abcde')
  4047. def test_invalid_large(self):
  4048. with self.assertRaises(ValueError):
  4049. mi.combination_with_replacement_index(tuple('abcdefg'), 'abcdef')
  4050. def test_invalid_match(self):
  4051. with self.assertRaises(ValueError):
  4052. mi.combination_with_replacement_index(tuple('axe'), 'abcde')
  4053. class PermutationIndexTests(TestCase):
  4054. def test_r_less_than_n(self):
  4055. iterable = 'abcdefg'
  4056. r = 4
  4057. first_index = {}
  4058. for index, element in enumerate(permutations(iterable, r)):
  4059. actual = mi.permutation_index(element, iterable)
  4060. expected = first_index.setdefault(element, index)
  4061. self.assertEqual(actual, expected)
  4062. def test_r_equal_to_n(self):
  4063. iterable = 'abcd'
  4064. first_index = {}
  4065. for index, element in enumerate(permutations(iterable)):
  4066. actual = mi.permutation_index(element, iterable)
  4067. expected = first_index.setdefault(element, index)
  4068. self.assertEqual(actual, expected)
  4069. def test_multiplicity(self):
  4070. iterable = 'abacba'
  4071. r = 3
  4072. first_index = {}
  4073. for index, element in enumerate(permutations(iterable, r)):
  4074. actual = mi.permutation_index(element, iterable)
  4075. expected = first_index.setdefault(element, index)
  4076. self.assertEqual(actual, expected)
  4077. def test_null(self):
  4078. actual = mi.permutation_index(tuple(), [])
  4079. expected = 0
  4080. self.assertEqual(actual, expected)
  4081. def test_long(self):
  4082. actual = mi.permutation_index((2, 12, 35, 126), range(180))
  4083. expected = 11631678
  4084. self.assertEqual(actual, expected)
  4085. def test_invalid_large(self):
  4086. with self.assertRaises(ValueError):
  4087. mi.permutation_index(tuple('abcdefg'), 'abcdef')
  4088. def test_invalid_match(self):
  4089. with self.assertRaises(ValueError):
  4090. mi.permutation_index(tuple('axe'), 'abcde')
  4091. class CountableTests(TestCase):
  4092. def test_empty(self):
  4093. iterable = []
  4094. it = mi.countable(iterable)
  4095. self.assertEqual(it.items_seen, 0)
  4096. self.assertEqual(list(it), [])
  4097. def test_basic(self):
  4098. iterable = '0123456789'
  4099. it = mi.countable(iterable)
  4100. self.assertEqual(it.items_seen, 0)
  4101. self.assertEqual(next(it), '0')
  4102. self.assertEqual(it.items_seen, 1)
  4103. self.assertEqual(''.join(it), '123456789')
  4104. self.assertEqual(it.items_seen, 10)
  4105. class ChunkedEvenTests(TestCase):
  4106. """Tests for ``chunked_even()``"""
  4107. def test_0(self):
  4108. self._test_finite('', 3, [])
  4109. def test_1(self):
  4110. self._test_finite('A', 1, [['A']])
  4111. def test_4(self):
  4112. self._test_finite('ABCD', 3, [['A', 'B'], ['C', 'D']])
  4113. def test_5(self):
  4114. self._test_finite('ABCDE', 3, [['A', 'B', 'C'], ['D', 'E']])
  4115. def test_6(self):
  4116. self._test_finite('ABCDEF', 3, [['A', 'B', 'C'], ['D', 'E', 'F']])
  4117. def test_7(self):
  4118. self._test_finite(
  4119. 'ABCDEFG', 3, [['A', 'B', 'C'], ['D', 'E'], ['F', 'G']]
  4120. )
  4121. def _test_finite(self, seq, n, expected):
  4122. # Check with and without `len()`
  4123. self.assertEqual(list(mi.chunked_even(seq, n)), expected)
  4124. self.assertEqual(list(mi.chunked_even(iter(seq), n)), expected)
  4125. def test_infinite(self):
  4126. for n in range(1, 5):
  4127. k = 0
  4128. def count_with_assert():
  4129. for i in count():
  4130. # Look-ahead should be less than n^2
  4131. self.assertLessEqual(i, n * k + n * n)
  4132. yield i
  4133. ls = mi.chunked_even(count_with_assert(), n)
  4134. while k < 2:
  4135. self.assertEqual(next(ls), list(range(k * n, (k + 1) * n)))
  4136. k += 1
  4137. def test_evenness(self):
  4138. for N in range(1, 50):
  4139. for n in range(1, N + 2):
  4140. lengths = []
  4141. items = []
  4142. for l in mi.chunked_even(range(N), n):
  4143. L = len(l)
  4144. self.assertLessEqual(L, n)
  4145. self.assertGreaterEqual(L, 1)
  4146. lengths.append(L)
  4147. items.extend(l)
  4148. self.assertEqual(items, list(range(N)))
  4149. self.assertLessEqual(max(lengths) - min(lengths), 1)
  4150. class ZipBroadcastTests(TestCase):
  4151. def test_zip(self):
  4152. for objects, zipped, strict_ok in [
  4153. # Empty
  4154. ([], [], True),
  4155. # One argument
  4156. ([1], [(1,)], True),
  4157. ([[1]], [(1,)], True),
  4158. ([[1, 2]], [(1,), (2,)], True),
  4159. # All scalars
  4160. ([1, 2], [(1, 2)], True),
  4161. ([1, 2, 3], [(1, 2, 3)], True),
  4162. # Iterables with length = 0
  4163. ([[], 1], [], True),
  4164. ([1, []], [], True),
  4165. ([[], []], [], True),
  4166. ([[], 1, 2], [], True),
  4167. ([[], 1, []], [], True),
  4168. ([1, [], 2], [], True),
  4169. ([1, [], []], [], True),
  4170. ([[], [], 1], [], True),
  4171. ([[], [], []], [], True),
  4172. # Iterables with length = 1
  4173. ([1, [2]], [(1, 2)], True),
  4174. ([[1], 2], [(1, 2)], True),
  4175. ([[1], [2]], [(1, 2)], True),
  4176. ([1, [2], 3], [(1, 2, 3)], True),
  4177. ([1, [2], [3]], [(1, 2, 3)], True),
  4178. ([[1], 2, 3], [(1, 2, 3)], True),
  4179. ([[1], 2, [3]], [(1, 2, 3)], True),
  4180. ([[1], [2], 3], [(1, 2, 3)], True),
  4181. ([[1], [2], [3]], [(1, 2, 3)], True),
  4182. # Iterables with length > 1
  4183. ([1, [2, 3]], [(1, 2), (1, 3)], True),
  4184. ([[1, 2], 3], [(1, 3), (2, 3)], True),
  4185. ([[1, 2], [3, 4]], [(1, 3), (2, 4)], True),
  4186. ([1, [2, 3], 4], [(1, 2, 4), (1, 3, 4)], True),
  4187. ([1, [2, 3], [4, 5]], [(1, 2, 4), (1, 3, 5)], True),
  4188. ([[1, 2], 3, 4], [(1, 3, 4), (2, 3, 4)], True),
  4189. ([[1, 2], 3, [4, 5]], [(1, 3, 4), (2, 3, 5)], True),
  4190. ([[1, 2], [3, 4], 5], [(1, 3, 5), (2, 4, 5)], True),
  4191. ([[1, 2], [3, 4], [5, 6]], [(1, 3, 5), (2, 4, 6)], True),
  4192. # Iterables with different lengths
  4193. ([[], [1]], [], False),
  4194. ([[1], []], [], False),
  4195. ([[1], [2, 3]], [(1, 2)], False),
  4196. ([[1, 2], [3]], [(1, 3)], False),
  4197. ([[1, 2], [3], [4]], [(1, 3, 4)], False),
  4198. ([[1], [2, 3], [4]], [(1, 2, 4)], False),
  4199. ([[1], [2], [3, 4]], [(1, 2, 3)], False),
  4200. ([[1], [2, 3], [4, 5]], [(1, 2, 4)], False),
  4201. ([[1, 2], [3], [4, 5]], [(1, 3, 4)], False),
  4202. ([[1, 2], [3, 4], [5]], [(1, 3, 5)], False),
  4203. ([1, [2, 3], [4, 5, 6]], [(1, 2, 4), (1, 3, 5)], False),
  4204. ([[1, 2], 3, [4, 5, 6]], [(1, 3, 4), (2, 3, 5)], False),
  4205. ([1, [2, 3, 4], [5, 6]], [(1, 2, 5), (1, 3, 6)], False),
  4206. ([[1, 2, 3], 4, [5, 6]], [(1, 4, 5), (2, 4, 6)], False),
  4207. ([[1, 2], [3, 4, 5], 6], [(1, 3, 6), (2, 4, 6)], False),
  4208. ([[1, 2, 3], [4, 5], 6], [(1, 4, 6), (2, 5, 6)], False),
  4209. # Infinite
  4210. ([count(), 1, [2]], [(0, 1, 2)], False),
  4211. ([count(), 1, [2, 3]], [(0, 1, 2), (1, 1, 3)], False),
  4212. # Miscellaneous
  4213. (['a', [1, 2], [3, 4, 5]], [('a', 1, 3), ('a', 2, 4)], False),
  4214. ]:
  4215. # Truncate by default
  4216. with self.subTest(objects=objects, strict=False, zipped=zipped):
  4217. self.assertEqual(list(mi.zip_broadcast(*objects)), zipped)
  4218. # Raise an exception for strict=True
  4219. with self.subTest(objects=objects, strict=True, zipped=zipped):
  4220. if strict_ok:
  4221. self.assertEqual(
  4222. list(mi.zip_broadcast(*objects, strict=True)),
  4223. zipped,
  4224. )
  4225. else:
  4226. with self.assertRaises(ValueError):
  4227. list(mi.zip_broadcast(*objects, strict=True))
  4228. def test_scalar_types(self):
  4229. # Default: str and bytes are treated as scalar
  4230. self.assertEqual(
  4231. list(mi.zip_broadcast('ab', [1, 2, 3])),
  4232. [('ab', 1), ('ab', 2), ('ab', 3)],
  4233. )
  4234. self.assertEqual(
  4235. list(mi.zip_broadcast(b'ab', [1, 2, 3])),
  4236. [(b'ab', 1), (b'ab', 2), (b'ab', 3)],
  4237. )
  4238. # scalar_types=None allows str and bytes to be treated as iterable
  4239. self.assertEqual(
  4240. list(mi.zip_broadcast('abc', [1, 2, 3], scalar_types=None)),
  4241. [('a', 1), ('b', 2), ('c', 3)],
  4242. )
  4243. # Use a custom type
  4244. self.assertEqual(
  4245. list(mi.zip_broadcast({'a': 'b'}, [1, 2, 3], scalar_types=dict)),
  4246. [({'a': 'b'}, 1), ({'a': 'b'}, 2), ({'a': 'b'}, 3)],
  4247. )
  4248. class UniqueInWindowTests(TestCase):
  4249. def test_invalid_n(self):
  4250. with self.assertRaises(ValueError):
  4251. list(mi.unique_in_window([], 0))
  4252. def test_basic(self):
  4253. for iterable, n, expected in [
  4254. (range(9), 10, list(range(9))),
  4255. (range(20), 10, list(range(20))),
  4256. ([1, 2, 3, 4, 4, 4], 1, [1, 2, 3, 4, 4, 4]),
  4257. ([1, 2, 3, 4, 4, 4], 2, [1, 2, 3, 4]),
  4258. ([1, 2, 3, 4, 4, 4], 3, [1, 2, 3, 4]),
  4259. ([1, 2, 3, 4, 4, 4], 4, [1, 2, 3, 4]),
  4260. ([1, 2, 3, 4, 4, 4], 5, [1, 2, 3, 4]),
  4261. (
  4262. [0, 1, 1, 1, 1, 1, 1, 1, 1, 0, 2, 3, 4, 2],
  4263. 2,
  4264. [0, 1, 0, 2, 3, 4, 2],
  4265. ),
  4266. ]:
  4267. with self.subTest(expected=expected):
  4268. actual = list(mi.unique_in_window(iterable, n))
  4269. self.assertEqual(actual, expected)
  4270. def test_key(self):
  4271. iterable = [0, 1, 3, 4, 5, 6, 7, 8, 9]
  4272. n = 3
  4273. key = lambda x: x // 3
  4274. actual = list(mi.unique_in_window(iterable, n, key=key))
  4275. expected = [0, 3, 6, 9]
  4276. self.assertEqual(actual, expected)
  4277. class StrictlyNTests(TestCase):
  4278. def test_basic(self):
  4279. iterable = ['a', 'b', 'c', 'd']
  4280. n = 4
  4281. actual = list(mi.strictly_n(iter(iterable), n))
  4282. expected = iterable
  4283. self.assertEqual(actual, expected)
  4284. def test_too_short_default(self):
  4285. iterable = ['a', 'b', 'c', 'd']
  4286. n = 5
  4287. with self.assertRaises(ValueError) as exc:
  4288. list(mi.strictly_n(iter(iterable), n))
  4289. self.assertEqual(
  4290. 'Too few items in iterable (got 4)', exc.exception.args[0]
  4291. )
  4292. def test_too_long_default(self):
  4293. iterable = ['a', 'b', 'c', 'd']
  4294. n = 3
  4295. with self.assertRaises(ValueError) as cm:
  4296. list(mi.strictly_n(iter(iterable), n))
  4297. self.assertEqual(
  4298. 'Too many items in iterable (got at least 4)',
  4299. cm.exception.args[0],
  4300. )
  4301. def test_too_short_custom(self):
  4302. call_count = 0
  4303. def too_short(item_count):
  4304. nonlocal call_count
  4305. call_count += 1
  4306. iterable = ['a', 'b', 'c', 'd']
  4307. n = 6
  4308. actual = []
  4309. for item in mi.strictly_n(iter(iterable), n, too_short=too_short):
  4310. actual.append(item)
  4311. expected = ['a', 'b', 'c', 'd']
  4312. self.assertEqual(actual, expected)
  4313. self.assertEqual(call_count, 1)
  4314. def test_too_long_custom(self):
  4315. import logging
  4316. iterable = ['a', 'b', 'c', 'd']
  4317. n = 2
  4318. too_long = lambda item_count: logging.warning(
  4319. 'Picked the first %s items', n
  4320. )
  4321. with self.assertLogs(level='WARNING') as cm:
  4322. actual = list(mi.strictly_n(iter(iterable), n, too_long=too_long))
  4323. self.assertEqual(actual, ['a', 'b'])
  4324. self.assertIn('Picked the first 2 items', cm.output[0])
  4325. class DuplicatesEverSeenTests(TestCase):
  4326. def test_basic(self):
  4327. for iterable, expected in [
  4328. ([], []),
  4329. ([1, 2, 3], []),
  4330. ([1, 1], [1]),
  4331. ([1, 2, 1, 2], [1, 2]),
  4332. ([1, 2, 3, '1'], []),
  4333. ]:
  4334. with self.subTest(args=(iterable,)):
  4335. self.assertEqual(
  4336. list(mi.duplicates_everseen(iterable)), expected
  4337. )
  4338. def test_non_hashable(self):
  4339. self.assertEqual(list(mi.duplicates_everseen([[1, 2], [3, 4]])), [])
  4340. self.assertEqual(
  4341. list(mi.duplicates_everseen([[1, 2], [3, 4], [1, 2]])), [[1, 2]]
  4342. )
  4343. def test_partially_hashable(self):
  4344. self.assertEqual(
  4345. list(mi.duplicates_everseen([[1, 2], [3, 4], (5, 6)])), []
  4346. )
  4347. self.assertEqual(
  4348. list(mi.duplicates_everseen([[1, 2], [3, 4], (5, 6), [1, 2]])),
  4349. [[1, 2]],
  4350. )
  4351. self.assertEqual(
  4352. list(mi.duplicates_everseen([[1, 2], [3, 4], (5, 6), (5, 6)])),
  4353. [(5, 6)],
  4354. )
  4355. def test_key_hashable(self):
  4356. iterable = 'HEheHEhe'
  4357. self.assertEqual(list(mi.duplicates_everseen(iterable)), list('HEhe'))
  4358. self.assertEqual(
  4359. list(mi.duplicates_everseen(iterable, str.lower)),
  4360. list('heHEhe'),
  4361. )
  4362. def test_key_non_hashable(self):
  4363. iterable = [[1, 2], [3, 0], [5, -2], [5, 6]]
  4364. self.assertEqual(
  4365. list(mi.duplicates_everseen(iterable, lambda x: x)), []
  4366. )
  4367. self.assertEqual(
  4368. list(mi.duplicates_everseen(iterable, sum)), [[3, 0], [5, -2]]
  4369. )
  4370. def test_key_partially_hashable(self):
  4371. iterable = [[1, 2], (1, 2), [1, 2], [5, 6]]
  4372. self.assertEqual(
  4373. list(mi.duplicates_everseen(iterable, lambda x: x)), [[1, 2]]
  4374. )
  4375. self.assertEqual(
  4376. list(mi.duplicates_everseen(iterable, list)), [(1, 2), [1, 2]]
  4377. )
  4378. class DuplicatesJustSeenTests(TestCase):
  4379. def test_basic(self):
  4380. for iterable, expected in [
  4381. ([], []),
  4382. ([1, 2, 3, 3, 2, 2], [3, 2]),
  4383. ([1, 1], [1]),
  4384. ([1, 2, 1, 2], []),
  4385. ([1, 2, 3, '1'], []),
  4386. ]:
  4387. with self.subTest(args=(iterable,)):
  4388. self.assertEqual(
  4389. list(mi.duplicates_justseen(iterable)), expected
  4390. )
  4391. def test_non_hashable(self):
  4392. self.assertEqual(list(mi.duplicates_justseen([[1, 2], [3, 4]])), [])
  4393. self.assertEqual(
  4394. list(
  4395. mi.duplicates_justseen(
  4396. [[1, 2], [3, 4], [3, 4], [3, 4], [1, 2]]
  4397. )
  4398. ),
  4399. [[3, 4], [3, 4]],
  4400. )
  4401. def test_partially_hashable(self):
  4402. self.assertEqual(
  4403. list(mi.duplicates_justseen([[1, 2], [3, 4], (5, 6)])), []
  4404. )
  4405. self.assertEqual(
  4406. list(
  4407. mi.duplicates_justseen(
  4408. [[1, 2], [3, 4], (5, 6), [1, 2], [1, 2]]
  4409. )
  4410. ),
  4411. [[1, 2]],
  4412. )
  4413. self.assertEqual(
  4414. list(
  4415. mi.duplicates_justseen(
  4416. [[1, 2], [3, 4], (5, 6), (5, 6), (5, 6)]
  4417. )
  4418. ),
  4419. [(5, 6), (5, 6)],
  4420. )
  4421. def test_key_hashable(self):
  4422. iterable = 'HEheHHHhEheeEe'
  4423. self.assertEqual(list(mi.duplicates_justseen(iterable)), list('HHe'))
  4424. self.assertEqual(
  4425. list(mi.duplicates_justseen(iterable, str.lower)),
  4426. list('HHheEe'),
  4427. )
  4428. def test_key_non_hashable(self):
  4429. iterable = [[1, 2], [3, 0], [5, -2], [5, 6], [1, 2]]
  4430. self.assertEqual(
  4431. list(mi.duplicates_justseen(iterable, lambda x: x)), []
  4432. )
  4433. self.assertEqual(
  4434. list(mi.duplicates_justseen(iterable, sum)), [[3, 0], [5, -2]]
  4435. )
  4436. def test_key_partially_hashable(self):
  4437. iterable = [[1, 2], (1, 2), [1, 2], [5, 6], [1, 2]]
  4438. self.assertEqual(
  4439. list(mi.duplicates_justseen(iterable, lambda x: x)), []
  4440. )
  4441. self.assertEqual(
  4442. list(mi.duplicates_justseen(iterable, list)), [(1, 2), [1, 2]]
  4443. )
  4444. def test_nested(self):
  4445. iterable = [[[1, 2], [1, 2]], [5, 6], [5, 6]]
  4446. self.assertEqual(list(mi.duplicates_justseen(iterable)), [[5, 6]])
  4447. class ClassifyUniqueTests(TestCase):
  4448. def test_basic(self):
  4449. self.assertEqual(
  4450. list(mi.classify_unique('mississippi')),
  4451. [
  4452. ('m', True, True),
  4453. ('i', True, True),
  4454. ('s', True, True),
  4455. ('s', False, False),
  4456. ('i', True, False),
  4457. ('s', True, False),
  4458. ('s', False, False),
  4459. ('i', True, False),
  4460. ('p', True, True),
  4461. ('p', False, False),
  4462. ('i', True, False),
  4463. ],
  4464. )
  4465. def test_non_hashable(self):
  4466. self.assertEqual(
  4467. list(mi.classify_unique([[1, 2], [3, 4], [3, 4], [1, 2]])),
  4468. [
  4469. ([1, 2], True, True),
  4470. ([3, 4], True, True),
  4471. ([3, 4], False, False),
  4472. ([1, 2], True, False),
  4473. ],
  4474. )
  4475. def test_partially_hashable(self):
  4476. self.assertEqual(
  4477. list(
  4478. mi.classify_unique(
  4479. [[1, 2], [3, 4], (5, 6), (5, 6), (3, 4), [1, 2]]
  4480. )
  4481. ),
  4482. [
  4483. ([1, 2], True, True),
  4484. ([3, 4], True, True),
  4485. ((5, 6), True, True),
  4486. ((5, 6), False, False),
  4487. ((3, 4), True, True),
  4488. ([1, 2], True, False),
  4489. ],
  4490. )
  4491. def test_key_hashable(self):
  4492. iterable = 'HEheHHHhEheeEe'
  4493. self.assertEqual(
  4494. list(mi.classify_unique(iterable)),
  4495. [
  4496. ('H', True, True),
  4497. ('E', True, True),
  4498. ('h', True, True),
  4499. ('e', True, True),
  4500. ('H', True, False),
  4501. ('H', False, False),
  4502. ('H', False, False),
  4503. ('h', True, False),
  4504. ('E', True, False),
  4505. ('h', True, False),
  4506. ('e', True, False),
  4507. ('e', False, False),
  4508. ('E', True, False),
  4509. ('e', True, False),
  4510. ],
  4511. )
  4512. self.assertEqual(
  4513. list(mi.classify_unique(iterable, str.lower)),
  4514. [
  4515. ('H', True, True),
  4516. ('E', True, True),
  4517. ('h', True, False),
  4518. ('e', True, False),
  4519. ('H', True, False),
  4520. ('H', False, False),
  4521. ('H', False, False),
  4522. ('h', False, False),
  4523. ('E', True, False),
  4524. ('h', True, False),
  4525. ('e', True, False),
  4526. ('e', False, False),
  4527. ('E', False, False),
  4528. ('e', False, False),
  4529. ],
  4530. )
  4531. def test_key_non_hashable(self):
  4532. iterable = [[1, 2], [3, 0], [5, -2], [5, 6], [1, 2]]
  4533. self.assertEqual(
  4534. list(mi.classify_unique(iterable, lambda x: x)),
  4535. [
  4536. ([1, 2], True, True),
  4537. ([3, 0], True, True),
  4538. ([5, -2], True, True),
  4539. ([5, 6], True, True),
  4540. ([1, 2], True, False),
  4541. ],
  4542. )
  4543. self.assertEqual(
  4544. list(mi.classify_unique(iterable, sum)),
  4545. [
  4546. ([1, 2], True, True),
  4547. ([3, 0], False, False),
  4548. ([5, -2], False, False),
  4549. ([5, 6], True, True),
  4550. ([1, 2], True, False),
  4551. ],
  4552. )
  4553. def test_key_partially_hashable(self):
  4554. iterable = [[1, 2], (1, 2), [1, 2], [5, 6], [1, 2]]
  4555. self.assertEqual(
  4556. list(mi.classify_unique(iterable, lambda x: x)),
  4557. [
  4558. ([1, 2], True, True),
  4559. ((1, 2), True, True),
  4560. ([1, 2], True, False),
  4561. ([5, 6], True, True),
  4562. ([1, 2], True, False),
  4563. ],
  4564. )
  4565. self.assertEqual(
  4566. list(mi.classify_unique(iterable, list)),
  4567. [
  4568. ([1, 2], True, True),
  4569. ((1, 2), False, False),
  4570. ([1, 2], False, False),
  4571. ([5, 6], True, True),
  4572. ([1, 2], True, False),
  4573. ],
  4574. )
  4575. def test_vs_unique_everseen(self):
  4576. input = 'AAAABBBBCCDAABBB'
  4577. output = [e for e, j, u in mi.classify_unique(input) if u]
  4578. self.assertEqual(output, ['A', 'B', 'C', 'D'])
  4579. self.assertEqual(list(mi.unique_everseen(input)), output)
  4580. def test_vs_unique_everseen_key(self):
  4581. input = 'aAbACCc'
  4582. output = [e for e, j, u in mi.classify_unique(input, str.lower) if u]
  4583. self.assertEqual(output, list('abC'))
  4584. self.assertEqual(list(mi.unique_everseen(input, str.lower)), output)
  4585. def test_vs_unique_justseen(self):
  4586. input = 'AAAABBBCCDABB'
  4587. output = [e for e, j, u in mi.classify_unique(input) if j]
  4588. self.assertEqual(output, list('ABCDAB'))
  4589. self.assertEqual(list(mi.unique_justseen(input)), output)
  4590. def test_vs_unique_justseen_key(self):
  4591. input = 'AABCcAD'
  4592. output = [e for e, j, u in mi.classify_unique(input, str.lower) if j]
  4593. self.assertEqual(output, list('ABCAD'))
  4594. self.assertEqual(list(mi.unique_justseen(input, str.lower)), output)
  4595. def test_vs_duplicates_everseen(self):
  4596. input = [1, 2, 1, 2]
  4597. output = [e for e, j, u in mi.classify_unique(input) if not u]
  4598. self.assertEqual(output, [1, 2])
  4599. self.assertEqual(list(mi.duplicates_everseen(input)), output)
  4600. def test_vs_duplicates_everseen_key(self):
  4601. input = 'HEheHEhe'
  4602. output = [
  4603. e for e, j, u in mi.classify_unique(input, str.lower) if not u
  4604. ]
  4605. self.assertEqual(output, list('heHEhe'))
  4606. self.assertEqual(
  4607. list(mi.duplicates_everseen(input, str.lower)), output
  4608. )
  4609. def test_vs_duplicates_justseen(self):
  4610. input = [1, 2, 3, 3, 2, 2]
  4611. output = [e for e, j, u in mi.classify_unique(input) if not j]
  4612. self.assertEqual(output, [3, 2])
  4613. self.assertEqual(list(mi.duplicates_justseen(input)), output)
  4614. def test_vs_duplicates_justseen_key(self):
  4615. input = 'HEheHHHhEheeEe'
  4616. output = [
  4617. e for e, j, u in mi.classify_unique(input, str.lower) if not j
  4618. ]
  4619. self.assertEqual(output, list('HHheEe'))
  4620. self.assertEqual(
  4621. list(mi.duplicates_justseen(input, str.lower)), output
  4622. )
  4623. class LongestCommonPrefixTests(TestCase):
  4624. def test_basic(self):
  4625. iterables = [[1, 2], [1, 2, 3], [1, 2, 4]]
  4626. self.assertEqual(list(mi.longest_common_prefix(iterables)), [1, 2])
  4627. def test_iterators(self):
  4628. iterables = iter([iter([1, 2]), iter([1, 2, 3]), iter([1, 2, 4])])
  4629. self.assertEqual(list(mi.longest_common_prefix(iterables)), [1, 2])
  4630. def test_no_iterables(self):
  4631. iterables = []
  4632. self.assertEqual(list(mi.longest_common_prefix(iterables)), [])
  4633. def test_empty_iterables_only(self):
  4634. iterables = [[], [], []]
  4635. self.assertEqual(list(mi.longest_common_prefix(iterables)), [])
  4636. def test_includes_empty_iterables(self):
  4637. iterables = [[1, 2], [1, 2, 3], [1, 2, 4], []]
  4638. self.assertEqual(list(mi.longest_common_prefix(iterables)), [])
  4639. def test_non_hashable(self):
  4640. # See https://github.com/more-itertools/more-itertools/issues/603
  4641. iterables = [[[1], [2]], [[1], [2], [3]], [[1], [2], [4]]]
  4642. self.assertEqual(list(mi.longest_common_prefix(iterables)), [[1], [2]])
  4643. def test_prefix_contains_elements_of_the_first_iterable(self):
  4644. iterables = [[[1], [2]], [[1], [2], [3]], [[1], [2], [4]]]
  4645. prefix = list(mi.longest_common_prefix(iterables))
  4646. self.assertIs(prefix[0], iterables[0][0])
  4647. self.assertIs(prefix[1], iterables[0][1])
  4648. self.assertIsNot(prefix[0], iterables[1][0])
  4649. self.assertIsNot(prefix[1], iterables[1][1])
  4650. self.assertIsNot(prefix[0], iterables[2][0])
  4651. self.assertIsNot(prefix[1], iterables[2][1])
  4652. def test_infinite_iterables(self):
  4653. prefix = mi.longest_common_prefix([count(), count()])
  4654. self.assertEqual(next(prefix), 0)
  4655. self.assertEqual(next(prefix), 1)
  4656. self.assertEqual(next(prefix), 2)
  4657. def test_contains_infinite_iterables(self):
  4658. iterables = [[0, 1, 2], count()]
  4659. self.assertEqual(list(mi.longest_common_prefix(iterables)), [0, 1, 2])
  4660. class IequalsTests(TestCase):
  4661. def test_basic(self):
  4662. self.assertTrue(mi.iequals("abc", iter("abc")))
  4663. self.assertTrue(mi.iequals(range(3), [0, 1, 2]))
  4664. self.assertFalse(mi.iequals("abc", [0, 1, 2]))
  4665. def test_no_iterables(self):
  4666. self.assertTrue(mi.iequals())
  4667. def test_one_iterable(self):
  4668. self.assertTrue(mi.iequals("abc"))
  4669. def test_more_than_two_iterable(self):
  4670. self.assertTrue(mi.iequals("abc", iter("abc"), ['a', 'b', 'c']))
  4671. self.assertFalse(mi.iequals("abc", iter("abc"), ['a', 'b', 'd']))
  4672. def test_order_matters(self):
  4673. self.assertFalse(mi.iequals("abc", "acb"))
  4674. def test_not_equal_lengths(self):
  4675. self.assertFalse(mi.iequals("abc", "ab"))
  4676. self.assertFalse(mi.iequals("abc", "bc"))
  4677. self.assertFalse(mi.iequals("aaa", "aaaa"))
  4678. def test_empty_iterables(self):
  4679. self.assertTrue(mi.iequals([], ""))
  4680. def test_none_is_not_a_sentinel(self):
  4681. # See https://stackoverflow.com/a/900444
  4682. self.assertFalse(mi.iequals([1, 2], [1, 2, None]))
  4683. self.assertFalse(mi.iequals([1, 2], [None, 1, 2]))
  4684. def test_not_identical_but_equal(self):
  4685. self.assertTrue([1, True], [1.0, complex(1, 0)])
  4686. class ConstrainedBatchesTests(TestCase):
  4687. def test_basic(self):
  4688. zen = [
  4689. 'Beautiful is better than ugly',
  4690. 'Explicit is better than implicit',
  4691. 'Simple is better than complex',
  4692. 'Complex is better than complicated',
  4693. 'Flat is better than nested',
  4694. 'Sparse is better than dense',
  4695. 'Readability counts',
  4696. ]
  4697. for size, expected in (
  4698. (
  4699. 34,
  4700. [
  4701. (zen[0],),
  4702. (zen[1],),
  4703. (zen[2],),
  4704. (zen[3],),
  4705. (zen[4],),
  4706. (zen[5],),
  4707. (zen[6],),
  4708. ],
  4709. ),
  4710. (
  4711. 61,
  4712. [
  4713. (zen[0], zen[1]),
  4714. (zen[2],),
  4715. (zen[3], zen[4]),
  4716. (zen[5], zen[6]),
  4717. ],
  4718. ),
  4719. (
  4720. 90,
  4721. [
  4722. (zen[0], zen[1], zen[2]),
  4723. (zen[3], zen[4], zen[5]),
  4724. (zen[6],),
  4725. ],
  4726. ),
  4727. (
  4728. 124,
  4729. [(zen[0], zen[1], zen[2], zen[3]), (zen[4], zen[5], zen[6])],
  4730. ),
  4731. (
  4732. 150,
  4733. [(zen[0], zen[1], zen[2], zen[3], zen[4]), (zen[5], zen[6])],
  4734. ),
  4735. (
  4736. 177,
  4737. [(zen[0], zen[1], zen[2], zen[3], zen[4], zen[5]), (zen[6],)],
  4738. ),
  4739. ):
  4740. with self.subTest(size=size):
  4741. actual = list(mi.constrained_batches(iter(zen), size))
  4742. self.assertEqual(actual, expected)
  4743. def test_max_count(self):
  4744. iterable = ['1', '1', '12345678', '12345', '12345']
  4745. max_size = 10
  4746. max_count = 2
  4747. actual = list(mi.constrained_batches(iterable, max_size, max_count))
  4748. expected = [('1', '1'), ('12345678',), ('12345', '12345')]
  4749. self.assertEqual(actual, expected)
  4750. def test_strict(self):
  4751. iterable = ['1', '123456789', '1']
  4752. size = 8
  4753. with self.assertRaises(ValueError):
  4754. list(mi.constrained_batches(iterable, size))
  4755. actual = list(mi.constrained_batches(iterable, size, strict=False))
  4756. expected = [('1',), ('123456789',), ('1',)]
  4757. self.assertEqual(actual, expected)
  4758. def test_get_len(self):
  4759. class Record(tuple):
  4760. def total_size(self):
  4761. return sum(len(x) for x in self)
  4762. record_3 = Record(('1', '23'))
  4763. record_5 = Record(('1234', '1'))
  4764. record_10 = Record(('1', '12345678', '1'))
  4765. record_2 = Record(('1', '1'))
  4766. iterable = [record_3, record_5, record_10, record_2]
  4767. self.assertEqual(
  4768. list(
  4769. mi.constrained_batches(
  4770. iterable, 10, get_len=lambda x: x.total_size()
  4771. )
  4772. ),
  4773. [(record_3, record_5), (record_10,), (record_2,)],
  4774. )
  4775. def test_bad_max(self):
  4776. with self.assertRaises(ValueError):
  4777. list(mi.constrained_batches([], 0))
  4778. class GrayProductTests(TestCase):
  4779. def test_basic(self):
  4780. self.assertEqual(
  4781. tuple(mi.gray_product(('a', 'b', 'c'), range(1, 3))),
  4782. (("a", 1), ("b", 1), ("c", 1), ("c", 2), ("b", 2), ("a", 2)),
  4783. )
  4784. out = mi.gray_product(('foo', 'bar'), (3, 4, 5, 6), ['quz', 'baz'])
  4785. self.assertEqual(next(out), ('foo', 3, 'quz'))
  4786. self.assertEqual(
  4787. list(out),
  4788. [
  4789. ('bar', 3, 'quz'),
  4790. ('bar', 4, 'quz'),
  4791. ('foo', 4, 'quz'),
  4792. ('foo', 5, 'quz'),
  4793. ('bar', 5, 'quz'),
  4794. ('bar', 6, 'quz'),
  4795. ('foo', 6, 'quz'),
  4796. ('foo', 6, 'baz'),
  4797. ('bar', 6, 'baz'),
  4798. ('bar', 5, 'baz'),
  4799. ('foo', 5, 'baz'),
  4800. ('foo', 4, 'baz'),
  4801. ('bar', 4, 'baz'),
  4802. ('bar', 3, 'baz'),
  4803. ('foo', 3, 'baz'),
  4804. ],
  4805. )
  4806. self.assertEqual(tuple(mi.gray_product()), ((),))
  4807. self.assertEqual(tuple(mi.gray_product((1, 2))), ((1,), (2,)))
  4808. def test_errors(self):
  4809. with self.assertRaises(ValueError):
  4810. list(mi.gray_product((1, 2), ()))
  4811. with self.assertRaises(ValueError):
  4812. list(mi.gray_product((1, 2), (2,)))
  4813. def test_vs_product(self):
  4814. iters = (
  4815. ("a", "b"),
  4816. range(3, 6),
  4817. [None, None],
  4818. {"i", "j", "k", "l"},
  4819. "XYZ",
  4820. )
  4821. self.assertEqual(
  4822. sorted(product(*iters)), sorted(mi.gray_product(*iters))
  4823. )
  4824. class PartialProductTests(TestCase):
  4825. def test_no_iterables(self):
  4826. self.assertEqual(tuple(mi.partial_product()), ((),))
  4827. def test_empty_iterable(self):
  4828. self.assertEqual(tuple(mi.partial_product('AB', '', 'CD')), ())
  4829. def test_one_iterable(self):
  4830. # a single iterable should pass through
  4831. self.assertEqual(
  4832. tuple(mi.partial_product('ABCD')),
  4833. (
  4834. ('A',),
  4835. ('B',),
  4836. ('C',),
  4837. ('D',),
  4838. ),
  4839. )
  4840. def test_two_iterables(self):
  4841. self.assertEqual(
  4842. list(mi.partial_product('ABCD', [1])),
  4843. [('A', 1), ('B', 1), ('C', 1), ('D', 1)],
  4844. )
  4845. expected = [
  4846. ('A', 1),
  4847. ('B', 1),
  4848. ('C', 1),
  4849. ('D', 1),
  4850. ('D', 2),
  4851. ('D', 3),
  4852. ('D', 4),
  4853. ]
  4854. self.assertEqual(
  4855. list(mi.partial_product('ABCD', [1, 2, 3, 4])), expected
  4856. )
  4857. def test_basic(self):
  4858. ones = [1, 2, 3]
  4859. tens = [10, 20, 30, 40, 50]
  4860. hundreds = [100, 200]
  4861. expected = [
  4862. (1, 10, 100),
  4863. (2, 10, 100),
  4864. (3, 10, 100),
  4865. (3, 20, 100),
  4866. (3, 30, 100),
  4867. (3, 40, 100),
  4868. (3, 50, 100),
  4869. (3, 50, 200),
  4870. ]
  4871. actual = list(mi.partial_product(ones, tens, hundreds))
  4872. self.assertEqual(actual, expected)
  4873. def test_uneven_length_iterables(self):
  4874. # this is also the docstring example
  4875. expected = [
  4876. ('A', 'C', 'D'),
  4877. ('B', 'C', 'D'),
  4878. ('B', 'C', 'E'),
  4879. ('B', 'C', 'F'),
  4880. ]
  4881. self.assertEqual(list(mi.partial_product('AB', 'C', 'DEF')), expected)
  4882. class IterateTests(TestCase):
  4883. def test_basic(self) -> None:
  4884. result = list(islice(mi.iterate(lambda x: 2 * x, start=1), 10))
  4885. expected = [1, 2, 4, 8, 16, 32, 64, 128, 256, 512]
  4886. self.assertEqual(result, expected)
  4887. def test_func_controls_iteration_stop(self) -> None:
  4888. def func(num):
  4889. if num > 100:
  4890. raise StopIteration
  4891. return num * 2
  4892. result = list(islice(mi.iterate(func, start=1), 10))
  4893. expected = [1, 2, 4, 8, 16, 32, 64, 128]
  4894. self.assertEqual(result, expected)
  4895. class TakewhileInclusiveTests(TestCase):
  4896. def test_basic(self) -> None:
  4897. result = list(mi.takewhile_inclusive(lambda x: x < 5, [1, 4, 6, 4, 1]))
  4898. expected = [1, 4, 6]
  4899. self.assertEqual(result, expected)
  4900. def test_empty_iterator(self) -> None:
  4901. result = list(mi.takewhile_inclusive(lambda x: True, []))
  4902. expected = []
  4903. self.assertEqual(result, expected)
  4904. def test_collatz_sequence(self) -> None:
  4905. is_even = lambda n: n % 2 == 0
  4906. start = 11
  4907. result = list(
  4908. mi.takewhile_inclusive(
  4909. lambda n: n != 1,
  4910. mi.iterate(
  4911. lambda n: n // 2 if is_even(n) else 3 * n + 1, start
  4912. ),
  4913. )
  4914. )
  4915. expected = [11, 34, 17, 52, 26, 13, 40, 20, 10, 5, 16, 8, 4, 2, 1]
  4916. self.assertEqual(result, expected)
  4917. class OuterProductTests(TestCase):
  4918. def test_basic(self) -> None:
  4919. greetings = ['Hello', 'Goodbye']
  4920. names = ['Alice', 'Bob', 'Carol']
  4921. greet = lambda greeting, name: f'{greeting}, {name}!'
  4922. result = list(mi.outer_product(greet, greetings, names))
  4923. expected = [
  4924. ('Hello, Alice!', 'Hello, Bob!', 'Hello, Carol!'),
  4925. ('Goodbye, Alice!', 'Goodbye, Bob!', 'Goodbye, Carol!'),
  4926. ]
  4927. self.assertEqual(result, expected)
  4928. class IterSuppressTests(TestCase):
  4929. class Producer:
  4930. def __init__(self, exc, die_early=False):
  4931. self.exc = exc
  4932. self.pos = 0
  4933. self.die_early = die_early
  4934. def __iter__(self):
  4935. if self.die_early:
  4936. raise self.exc
  4937. return self
  4938. def __next__(self):
  4939. ret = self.pos
  4940. if self.pos >= 5:
  4941. raise self.exc
  4942. self.pos += 1
  4943. return ret
  4944. def test_no_error(self):
  4945. iterator = range(5)
  4946. actual = list(mi.iter_suppress(iterator, RuntimeError))
  4947. expected = [0, 1, 2, 3, 4]
  4948. self.assertEqual(actual, expected)
  4949. def test_raises_error(self):
  4950. iterator = self.Producer(ValueError)
  4951. with self.assertRaises(ValueError):
  4952. list(mi.iter_suppress(iterator, RuntimeError))
  4953. def test_suppression(self):
  4954. iterator = self.Producer(ValueError)
  4955. actual = list(mi.iter_suppress(iterator, RuntimeError, ValueError))
  4956. expected = [0, 1, 2, 3, 4]
  4957. self.assertEqual(actual, expected)
  4958. def test_early_suppression(self):
  4959. iterator = self.Producer(ValueError, die_early=True)
  4960. actual = list(mi.iter_suppress(iterator, RuntimeError, ValueError))
  4961. expected = []
  4962. self.assertEqual(actual, expected)
  4963. class FilterMapTests(TestCase):
  4964. def test_no_iterables(self):
  4965. actual = list(mi.filter_map(lambda _: None, []))
  4966. expected = []
  4967. self.assertEqual(actual, expected)
  4968. def test_filter(self):
  4969. actual = list(mi.filter_map(lambda _: None, [1, 2, 3]))
  4970. expected = []
  4971. self.assertEqual(actual, expected)
  4972. def test_map(self):
  4973. actual = list(mi.filter_map(lambda x: x + 1, [1, 2, 3]))
  4974. expected = [2, 3, 4]
  4975. self.assertEqual(actual, expected)
  4976. def test_filter_map(self):
  4977. actual = list(
  4978. mi.filter_map(
  4979. lambda x: int(x) if x.isnumeric() else None,
  4980. ['1', 'a', '2', 'b', '3'],
  4981. )
  4982. )
  4983. expected = [1, 2, 3]
  4984. self.assertEqual(actual, expected)
  4985. class PowersetOfSetsTests(TestCase):
  4986. def test_simple(self):
  4987. iterable = [0, 1, 2]
  4988. actual = list(mi.powerset_of_sets(iterable))
  4989. expected = [set(), {0}, {1}, {2}, {0, 1}, {0, 2}, {1, 2}, {0, 1, 2}]
  4990. self.assertEqual(actual, expected)
  4991. def test_hash_count(self):
  4992. hash_count = 0
  4993. class Str(str):
  4994. def __hash__(true_self):
  4995. nonlocal hash_count
  4996. hash_count += 1
  4997. return super.__hash__(true_self)
  4998. iterable = map(Str, 'ABBBCDD')
  4999. self.assertEqual(len(list(mi.powerset_of_sets(iterable))), 128)
  5000. self.assertLessEqual(hash_count, 14)
  5001. class JoinMappingTests(TestCase):
  5002. def test_basic(self):
  5003. salary_map = {'e1': 12, 'e2': 23, 'e3': 34}
  5004. dept_map = {'e1': 'eng', 'e2': 'sales', 'e3': 'eng'}
  5005. service_map = {'e1': 5, 'e2': 9, 'e3': 2}
  5006. field_to_map = {
  5007. 'salary': salary_map,
  5008. 'dept': dept_map,
  5009. 'service': service_map,
  5010. }
  5011. expected = {
  5012. 'e1': {'salary': 12, 'dept': 'eng', 'service': 5},
  5013. 'e2': {'salary': 23, 'dept': 'sales', 'service': 9},
  5014. 'e3': {'salary': 34, 'dept': 'eng', 'service': 2},
  5015. }
  5016. self.assertEqual(dict(mi.join_mappings(**field_to_map)), expected)
  5017. def test_empty(self):
  5018. self.assertEqual(dict(mi.join_mappings()), {})
  5019. class DiscreteFourierTransformTests(TestCase):
  5020. def test_basic(self):
  5021. # Example calculation from:
  5022. # https://en.wikipedia.org/wiki/Discrete_Fourier_transform#Example
  5023. xarr = [1, 2 - 1j, -1j, -1 + 2j]
  5024. Xarr = [2, -2 - 2j, -2j, 4 + 4j]
  5025. self.assertTrue(all(map(cmath.isclose, mi.dft(xarr), Xarr)))
  5026. self.assertTrue(all(map(cmath.isclose, mi.idft(Xarr), xarr)))
  5027. def test_roundtrip(self):
  5028. for _ in range(1_000):
  5029. N = randrange(35)
  5030. xarr = [complex(random(), random()) for i in range(N)]
  5031. Xarr = list(mi.dft(xarr))
  5032. assert all(map(cmath.isclose, mi.idft(Xarr), xarr))
  5033. class DoubleStarMapTests(TestCase):
  5034. def test_construction(self):
  5035. iterable = [{'price': 1.23}, {'price': 42}, {'price': 0.1}]
  5036. actual = list(mi.doublestarmap('{price:.2f}'.format, iterable))
  5037. expected = ['1.23', '42.00', '0.10']
  5038. self.assertEqual(actual, expected)
  5039. def test_identity(self):
  5040. iterable = [{'x': 1}, {'x': 2}, {'x': 3}]
  5041. actual = list(mi.doublestarmap(lambda x: x, iterable))
  5042. expected = [1, 2, 3]
  5043. self.assertEqual(actual, expected)
  5044. def test_adding(self):
  5045. iterable = [{'a': 1, 'b': 2}, {'a': 3, 'b': 4}]
  5046. actual = list(mi.doublestarmap(lambda a, b: a + b, iterable))
  5047. expected = [3, 7]
  5048. self.assertEqual(actual, expected)
  5049. def test_mismatch_function_smaller(self):
  5050. iterable = [{'a': 1, 'b': 2}, {'a': 3, 'b': 4}]
  5051. with self.assertRaises(TypeError):
  5052. list(mi.doublestarmap(lambda a: a, iterable))
  5053. def test_mismatch_function_different(self):
  5054. iterable = [{'a': 1}, {'a': 2}]
  5055. with self.assertRaises(TypeError):
  5056. list(mi.doublestarmap(lambda x: x, iterable))
  5057. def test_mismatch_function_larger(self):
  5058. iterable = [{'a': 1}, {'a': 2}]
  5059. with self.assertRaises(TypeError):
  5060. list(mi.doublestarmap(lambda a, b: a + b, iterable))
  5061. def test_no_mapping(self):
  5062. iterable = [1, 2, 3, 4]
  5063. with self.assertRaises(TypeError):
  5064. list(mi.doublestarmap(lambda x: x, iterable))
  5065. def test_empty(self):
  5066. actual = list(mi.doublestarmap(lambda x: x, []))
  5067. expected = []
  5068. self.assertEqual(actual, expected)