B @`A@s"ddlmZddlZddlZddlZddlmZddlmZ ddl m m Z ejddgdddZejejdfejdfejdfejdfejdfejdfejdfejdfejdfejdfejdfejdfg dd d Zejejed ejfejed ejfejd ejfejdejfejdejfejdejfejdejfejd ejfejdejfejdejfejdejfejdejfejd ejfejdejfejdejfejdejfejdejfejd ejfejdejfejdejfejdejfejdejfgdddZGdddZGdddZ dS))datetimeN)iNaTTF)paramscCs|jS)N)param)requestr:/tmp/pip-unpacked-wheel-q9tj5l6a/pandas/tests/test_take.py writeable sr cCs|jS)N)r)rrrrdtype_can_hold_nasr g@y@@cCs|jS)N)r)rrrrdtype_fill_out_dtype)src@seZdZedZddZddZddZdd Z d d Z d d Z ddZ ddZ ddZddZddZddZddZddZdS)TestTakez Incompatible type for fill_valuec Cs|\}}tjddd|}||j_ddddg}tjd|d}tj|||d| |}t ||ddddg}tjd|d}|rtj|||d| |}tj |d<t ||n:t jt|jd tj|||dWdQRX|j ||ddS) Nrr )dtype)out)match)nprandomrandintastypeflagsr emptyalgostake_1dtaketmassert_almost_equalnanpytestraises TypeError fill_error) selfr r r can_hold_nadataindexerrexpectedrrrtest_1d_with_outKs$      zTestTake.test_1d_with_outcCs|\}}}tjddd|}ddddg}tj|||d}|dddg|dddgks`t|d|kspt|j|ks~tddddg}tj|||d}|ddddg||kst|j|kstdS)Nrrrr r) fill_valuer) rrrrrrallAssertionErrorr)r(rrr. out_dtyper*r+resultrrrtest_1d_fill_nonnags  $  zTestTake.test_1d_fill_nonnac Cs|\}}tjddd|}||j_ddddg}tjd|d}tjd|d}tj|||ddtj|||dd|j |dd } |j |dd } t || t || dddd g}tjd|d}tjd|d}|rRtj|||ddtj|||dd|j |dd } |j |dd } tj | d ddf<tj | ddd f<t || t || nZxXt ||gD]H\} } tjt|jd tj||| | dWdQRX|j || | dq`WdS) Nrr)rr )rr)r)r4r)raxis)r5rr)r)rrrrrr rrtake_ndr r!r"r# enumerater$r%r&r') r(r r rr)r*r+out0out1 expected0 expected1irrrrtest_2d_with_outws8     zTestTake.test_2d_with_outcCs|\}}}tjddd|}ddddg}tj||d|d}|dddgddf|dddgddfksrt|dddf|kst|j|ksttj||d|d}|dddddgf|dddddgfkst|dddf|kst|j|kstddddg}tj||d|d}|ddddgddf||ddfks^t|j|ksnttj||d|d}|ddddddgf|dd|fkst|j|kstdS)Nrr)r4rr r)r5r.r) rrrrrr6r/r0r)r(rrr.r1r*r+r2rrrtest_2d_fill_nonnas$  44 22zTestTake.test_2d_fill_nonnac CsB|\}}tjddd|}ddddg}tjd|d}tjd|d}tjd|d}tj|||ddtj|||ddtj|||dd|j|dd } |j|dd } |j|dd } t || t || t || dddd g}tjd|d}tjd|d}tjd|d}|rtj|||ddtj|||ddtj|||dd|j|dd } |j|dd } |j|dd } tj | d ddddf<tj | ddd ddf<tj | ddddd f<t || t || t || n\xZt |||gD]H\} } t j t|jd tj||| | dWdQRX|j|| | dqWdS) Nrr)r4rrr )rrr)r)r4rr)rr5)r5rr)r)rrrrrrr6r r!r"r#r7r$r%r&r')r(r rr)r*r+r8r9Zout2r:r;Z expected2r<rrrrtest_3d_with_outsH       zTestTake.test_3d_with_outcCs|\}}}tjddd|}ddddg}tj||d|d}|dddgddddf|dddgddddfks~t|dddddf|kst|j|ksttj||d|d}|dddddgddf|dddddgddfkst|dddddf|ks&t|j|ks6ttj||d|d}|dddddddgf|dddddddgfkst|dddddf|kst|j|kstddddg}tj||d|d}|ddddgddddf||ddddfkst|j|ks*ttj||d|d}|ddddddgddf|dd|ddfkszt|j|ksttj||d|d}|ddddddddgf|dddd|fkst|j|kstdS)Nrr)r4rrr r)r5r.r) rrrrrr6r/r0r)r(rrr.r1r*r+r2rrrtest_3d_fill_nonnas2  @"B$B$ >>>zTestTake.test_3d_fill_nonnacCsPtjdtj}ddddg}t||}||}tj|d<t ||dS)N r rrr) rrrandnrfloat32rrr r#r!r")r(arrr+r2r,rrrtest_1d_other_dtypes s     zTestTake.test_1d_other_dtypescCstjddtj}ddddg}tj||dd}|j|dd}tj|d<t ||tj||dd}|j|dd}tj|dddf<t ||dS) NrAr4r rrrr)r5) rrrBrrCrr6r r#r!r")r(rDr+r2r,rrrtest_2d_other_dtypess   zTestTake.test_2d_other_dtypescCsltjdddgtd}t|ddddg}|ddddg}t||t|dddg}|jtj ksht dS)Nrr )rrr) rarrayboolrrr r!assert_numpy_array_equalrobject_r0)r(rDr2r,rrr test_1d_bool&s  zTestTake.test_1d_boolcCstjdddgdddgdddggtd}t|ddddg}|jddddgdd}t||tj|ddddgdd}|jddddgdd}t||t|dddg}|jtj kst dS)Nrr )rr)r5r) rrGrHrr6r r!rIrrJr0)r(rDr2r,rrr test_2d_bool0s&  zTestTake.test_2d_boolcCs$tjddtj}dddddg}tj||dd}t|}tj||d|dt |||j |dd}tj |ddgddf<t ||tj t ||jdfd d }tj|||d tj||dd}t|}tj||d|dt |||j |dd}tj |ddddgf<t ||dS) Nrrrrrr )r5)r5rrC)r)r)rrrBrrCrr6 empty_liker!r"r r#rlenshape)r(rDr+r2result2r,rrrrtest_2d_float32>s$     zTestTake.test_2d_float32c Cstjdddd}|jdd}ddd d d g}tj||dd }t|}tj||d|d t|||j |dd }t |tj dd gddf<t||tj||dt dd d d}t|}tj|||dt dd d dt|||j |dd }t dd d |dd gddf<t||tj||d d }t|}tj||d |d t|||j |d d }t |tj dddd gf<t||tj||d t dd d d}t|}tj|||d t dd d dt|||j |d d }t dd d |dddd gf<t||dS)NiiY)r4rlh]zdatetime64[ns])rrrrr )r5)r5rri)r5r.)rr5r.) rrrviewrr6rMr!r"r rint64r)r(rDr+r2rPr,rrrtest_2d_datetime64ZsB            zTestTake.test_2d_datetime64cCstddd}t|ddg}tdddgdd d gg}t||tj|ddgd dd }tdddgdddgg}t||dS) N rrrrr r rA T) allow_fillr.)rarangereshaperr rGr!rI)r(rDr2r,rrrtest_take_axis_0s zTestTake.test_take_axis_0c Cstddd}tj|ddgdd}tddgdd gd d gd d gg}t||tj|ddgdddd}tddgddgd dgd dgg}t||tj t ddtj|ddgddddWdQRXdS)NrUrrrrr )r5rr4rVrWT)r5rXr.zindices are out-of-bounds)r) rrYrZrr rGr!rIr$r% IndexError)r(rDr2r,rrrtest_take_axis_1s" " zTestTake.test_take_axis_1N)__name__ __module__ __qualname__recompiler'r-r3r=r>r?r@rErFrKrLrQrTr[r_rrrrrGs (3"  0 rc@sFeZdZddZddZejdddgdd Zd d Z d d Z dS)TestExtensionTakec Csvtddg}d}tjt|dtj|ddgddWdQRXd}tjt|dtj|ddgd dWdQRXdS) Nr rzindices are out-of-bounds)rrT)rXz2index 2 is out of bounds for( axis 0 with)? size 2F)rrGr$r%r^rr )r(rDmsgrrrtest_bounds_check_largesz)TestExtensionTake.test_bounds_check_largec Cstjdddgtjd}dddg}d}tjt|d tj||d d WdQRXt||}tjdddgtjd}t ||dS) Nr rr)rrrz7'indices' contains values less than allowed \(-2 < -1\))rT)rX) rrGrSr$r% ValueErrorrr r!rI)r(rDr+rfr2r,rrrtest_bounds_check_smalls  z)TestExtensionTake.test_bounds_check_smallrXTFc Cs`tjgtjd}tj|g|d}t||d}tjt |dtj|dg|dWdQRXdS)N)r)rXzHcannot do a non-empty take from an empty axes.|indices are out-of-bounds)rr) rrGrSrr r!rIr$r%r^)r(rXrDr2rfrrrtest_take_emptys  z!TestExtensionTake.test_take_emptycCs:tjtgddgddd}tddg}t||dS)NrTg)rXr.)rr rrGr!rI)r(r2r,rrrtest_take_na_emptysz$TestExtensionTake.test_take_na_emptycCs8dddg}t|ddg}tddg}t||dS)Nr rrr)rr rrGr!rI)r(rDr2r,rrrtest_take_coerces_lists z(TestExtensionTake.test_take_coerces_listN) r`rarbrgrjr$markZ parametrizerkrlrmrrrrres   re)!rrcZnumpyrr$Z pandas._libsrZpandas._testingZ_testingr!Zpandas.core.algorithmscoreZ algorithmsrZfixturer Zfloat64rCZuint64Zuint32Zuint16Zuint8rSZint32Zint16Zint8rJZbool_r Z complex128rrrerrrrs^                      `