U
    >vf;                     @   s   d dl Z d dlZd dlZd dlmZ d dlmZmZm	Z	m
Z
mZmZ d dlmZ d dlmZ ejdd ZG dd dZG d	d
 d
ZdS )    N)SettingWithCopyError)	DataFrameIndex
IndexSlice
MultiIndexSeriesconcat)BDayc               	   C   s   t dddddgdddd	d
gdddddgg} tddgddgdddgdddggdddgdddgdddgdddggdddd gd!}t| |td"d#S )$NEJYr鷯|?5^ݿ6?粝K?-?T?9#?9]?46<R?gQg^)gxֿgI&?g=U@axbqı.n$@      4@g      >@         r         onetwothreefourlevelscodesnamesABCDEindexcolumns)nparrayr   r   list)arrr)    r/   G/tmp/pip-unpacked-wheel-vdrwu74i/pandas/tests/frame/indexing/test_xs.pyfour_level_index_dataframe   s    "
r1   c                   @   s4   e Zd Zdd Zdd Zdd Zdd Zd	d
 ZdS )TestXSc              	   C   st  |  }|jd }||}| D ]>\}}t|rNt|| | sbtq$||| | ks$tq$dddddddd	}	t|	}
|
d}|jtj	kst|d
 dkst|d dkstt
jttdd ||jd t   W 5 Q R X |jd
dd}|d
 }t|| |jd
dd}d|d d < |r^t|d
 |d
  |dk rptn|dk sptd S )Nr   r   r   )12r3   r4   3)r3   r4   r5   )ABr6   r7   z Timestamp('1999-12-31 00:00:00')matchr   )axis)copyr)   xsitemsr+   isnanAssertionErrorr   dtypeZobject_pytestraisesKeyErrorreescaper	   tmassert_series_equalall)selfZfloat_frameZdatetime_frameusing_copy_on_writeZfloat_frame_origidxr<   itemvalueZ	test_dataframeZseriesexpectedr/   r/   r0   test_xs'   s6    



  zTestXS.test_xsc                 C   s   t dgd}d|d< d|d< d|d< d	|d
< d|d< |d}tdddd	dgtddd}t|| t dddgd}|d}tg dtjd}t|| d S )Nr   r)   g      ?r6   foor7   g       @CbarDg      @Er'   )r)   namer   r   c)rW   r@   )r   r<   r   r-   rF   rG   r+   Zfloat64)rI   dfr<   expresultrO   r/   r/   r0   test_xs_cornerM   s    

zTestXS.test_xs_cornerc                 C   sD   t tjdddddddgd}|d}|jd }t|| d S )Nr   r   r   rX   r   rQ   )r   r+   randomrandnr<   ilocrF   rG   )rI   rY   ZcrossrZ   r/   r/   r0   test_xs_duplicates`   s     

zTestXS.test_xs_duplicatesc                 C   s   t dddddddddddddd	d
dg}|jdddd}|d d }t|| |jdd	dgdd}t|| d S )Nsatsun)r   r   Z
strawberry
        )dayflavourZsalesyearrh   rg   rf   F)level
drop_levelr   )re   ra   )r   	set_indexr<   rF   assert_frame_equalrI   rY   r[   rO   r/   r/   r0   test_xs_keep_levelg   s    zTestXS.test_xs_keep_levelc              	   C   s   t tdddtdtdd}| }|rdt  d|dd d < W 5 Q R X t	|| nt|rd}t
jt|d d|dd d < W 5 Q R X |ddk rtn(d|dd d < |ddk std S )	Nr   r   r   r(      r   zC\nA value is trying to be set on a copy of a slice from a DataFramer8   )r   r+   ZarangeZreshaperanger;   rF   Zraises_chained_assignment_errorr<   rl   rA   rB   r   anyr?   rH   )rI   using_array_managerrJ   dmdf_origmsgr/   r/   r0   test_xs_vieww   s    $
zTestXS.test_xs_viewN)__name__
__module____qualname__rP   r\   r`   rn   rv   r/   r/   r/   r0   r2   &   s
   &r2   c                   @   s2  e Zd Zdd Zdd Zdd Zdd Zd	d
 Zdd Ze	j
dddgdgfgdd Zdd Ze	j
ddddd eddgddfdddd edgddfgdd  Ze	j
d!d"d d#d gd$d% Ze	j
d!d&d d'd gd(d) Zd*d+ Zd,d- Zd.d/ Zd0d1 Zd2d3 Zd4d5 Zd6d7 Zd8d9 Zd:S );TestXSWithMultiIndexc              	   C   s   ddddddddgddddddddgg}t t| }tj|ddgd	}ttjd
ddddg|d}|jdddd}|j	d d dgf }t
|| d S )NrT   bazrR   quxr   r    firstsecondr&   r      r6   r7   rS   r(   )r   rT   )r~   r}   r   ri   r:   r   )r-   zipr   from_tuplesr   r+   r]   r^   r<   r_   rF   rl   )rI   ZarraysZtuplesr)   rY   r[   rO   r/   r/   r0   test_xs_doc_example   s    z(TestXSWithMultiIndex.test_xs_doc_examplec                 C   sx   t dd}td}tj||gddgd}ttjt|d|dd	d
g}|j	ddd}|j
dd d f }t|| d S )Nia2ie2abcdedateZsecidr   r   XYZri   )rp   r-   r   Zfrom_productr   r+   r]   r^   lenr<   locrF   rl   )rI   datesidsr)   rY   r[   rO   r/   r/   r0   test_xs_integer_key   s    
 z(TestXSWithMultiIndex.test_xs_integer_keyc                 C   sL   |}|j ddd}||jddk }tddddgd	d
|_t|| d S )Nr    r~   r   r   rR   rT   r{   r|   r}   rW   )r<   r)   Zget_level_valuesr   rF   rl   rI    multiindex_dataframe_random_datarY   r[   rO   r/   r/   r0   test_xs_level   s
    z"TestXSWithMultiIndex.test_xs_levelc                 C   s   t jdd}tdddgdddgd	d
dggdddgdddgdddggd}t||d}t|dd dgdggd}|jd	dd}t|| d S )Nr   r   r   pr   r   r   yrX   rzr   r   r   r$   r%   rQ   r   r+   r]   r^   r   r   r<   rF   rl   )rI   r.   r)   rY   rO   r[   r/   r/   r0   test_xs_level_eq_2   s    z'TestXSWithMultiIndex.test_xs_level_eq_2c              	   C   sf   |}|  }|jddd}|r,d|d d < n*d}tjt|d d|d d < W 5 Q R X t|| d S )Nr    r~   r   rc   AA value is trying to be set on a copy of a slice from a DataFramer8   r;   r<   rA   rB   r   rF   rl   )rI   r   rJ   rY   rt   r[   ru   r/   r/   r0   test_xs_setting_with_copy_error   s    z4TestXSWithMultiIndex.test_xs_setting_with_copy_errorc              	   C   sj   |}|  }|jdddgd}|r0d|d d < n*d}tjt|d d|d d < W 5 Q R X t|| d S )Nr   r   r   r"   r   rc   r   r8   r   )rI   r1   rJ   rY   rt   r[   ru   r/   r/   r0   (test_xs_setting_with_copy_error_multiple   s    z=TestXSWithMultiIndex.test_xs_setting_with_copy_error_multiplez
key, level)r   r~   r   r~   c                 C   sv   |}t |gd }|jjdks"tt |jdddgd }t|trX|jt||d}n|j||d}t	|| d S )Nr   Fr   r~   r   )
r   r)   Z	is_uniquer?   r<   
isinstancer-   tuplerF   rl   )rI   keyri   r   rN   rY   rO   r[   r/   r/   r0   test_xs_with_duplicates   s    
z,TestXSWithMultiIndex.test_xs_with_duplicatesc                 C   sv   dddddddg}t |dd	d
gddd	g}t d
ddddgitdtjddgd	dd}|jddd}t|| d S )N)r   r   r   )r   Zbbcder   )r   Zyzcde   )r   xbcde   )r   N   )r   zbcder   )r   ybcder   Za1Za2Zcntr*   r   r   r   r   r   r   r   rQ   r   r   )r   rk   r   r+   nanr<   rF   rl   )rI   accrY   rO   r[   r/   r/   r0   test_xs_missing_values_in_index   s    	z4TestXSWithMultiIndex.test_xs_missing_values_in_indexzkey, level, exp_arr, exp_indexr   lvl0c                 C   s   | d d ddf S )Nr   r   r/   r   r/   r/   r0   <lambda>      zTestXSWithMultiIndex.<lambda>rT   rR   lvl1r   c                 C   s   | d d ddf S )Nr   r   r/   r   r/   r/   r0   r     r   c           
      C   s   t jdd}tddgddddggddd	d	gdd	d
dggddgd}t||d}|j||d	d}t|||d}	t||	 d S )Nr   r   r   rT   rR   ZhelloZworldr   r   r   r   r   r   r#   r   r   r   )
rI   r   ri   Zexp_arrZ	exp_indexr.   r)   rY   r[   rO   r/   r/   r0   test_xs_named_levels_axis_eq_1  s    	z3TestXSWithMultiIndex.test_xs_named_levels_axis_eq_1indexerc                 C   s   | j dddgdS )Nr   r   r"   r   r<   rY   r/   r/   r0   r     r   c                 C   s   |  dj dddS )Nr   r   r"   r   r   r   r/   r/   r0   r     r   c                 C   s`   |}dddddgg}t dgdggdgdggd	d
gd}t||tdd}||}t|| d S )Nr   r   r   r   r   r   r   r   r    r!   r#   r'   r(   r   r   r-   rF   rl   rI   r   r1   rY   Zexpected_valuesZexpected_indexrO   r[   r/   r/   r0   test_xs_level_multiple  s    
 
   z+TestXSWithMultiIndex.test_xs_level_multiplec                 C   s   | j dddS )Nr   r   r   r   r   r/   r/   r0   r   ,  r   c                 C   s
   |  dS )Nr   r   r   r/   r/   r0   r   ,  r   c                 C   s   |}dddddgdddd	d
gg}t ddgddgddggddgddgddggdddgd}t||tdd}||}t|| d S )Nr
   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r    r!   r"   r#   r'   r(   r   r   r/   r/   r0   test_xs_level0+  s       z#TestXSWithMultiIndex.test_xs_level0c                 C   s*   |}| dj}|jd }t|| d S )NrT   r    r   )r<   valuesrF   Zassert_almost_equalr   r/   r/   r0   test_xs_values@  s    
z#TestXSWithMultiIndex.test_xs_valuesc                 C   s(   |}| d}|jd }t|| d S )Nr   )r<   r   rF   rG   r   r/   r/   r0   test_xs_loc_equalityF  s    

z)TestXSWithMultiIndex.test_xs_loc_equalityc              	   C   s   t dddgddggddddddgddddddggd}ttjdd	|d
}|tkr\|d }|jdd  d}|t	dd d f }t
|| |jt	dd d f  }t
|| d S )N)rR   rT   r   )rR   r{   r   )rR   r|   r   r   r   r   r      r   rQ   )r   r   r+   r]   r^   r   r_   Z	droplevelr<   r   rF   Zassert_equalr   )rI   frame_or_seriesr)   objrO   r[   r/   r/   r0   +test_xs_IndexSlice_argument_not_implementedL  s    z@TestXSWithMultiIndex.test_xs_IndexSlice_argument_not_implementedc              	   C   sR   t ddddgi}|tkr"|d }d}tjt|d |jddd	 W 5 Q R X d S )
Nr6   r   r   r   zIndex must be a MultiIndexr8   r   asr   )r   r   rA   rB   	TypeErrorr<   )rI   r   r   ru   r/   r/   r0   test_xs_levels_raises`  s    z*TestXSWithMultiIndex.test_xs_levels_raisesc                 C   sp   t jdddgddgd}tddd	gg|d
}|jdddd}tddggt jddgddgdd
}t|| d S )N)r   r   )r   r   )r   r   Zlevel1Zlevel2r   r   r   r   r   r   Fr:   rj   )r   r   r   r<   rF   rl   )rI   mirY   r[   rO   r/   r/   r0   "test_xs_multiindex_droplevel_falsei  s      z7TestXSWithMultiIndex.test_xs_multiindex_droplevel_falsec                 C   sL   t dddggtdddgd}|jdddd	}t ddgi}t|| d S )
Nr   r   r   r   r   rX   r   Fr   )r   r   r<   rF   rl   rm   r/   r/   r0   test_xs_droplevel_falsex  s    z,TestXSWithMultiIndex.test_xs_droplevel_falsec                 C   s  t dddggtdddgd}|jdddd	}t|jd d d
f j|jd d d
f js^td|jd< |r|t ddgi}nt ddgi}t	|| t dddggtdddgd}|jdddd	}d|jd< |rt ddgi}n"|rt ddgi}nt ddgi}t	|| d S )Nr   r   r   r   r   rX   r   Fr   r   )r   r   g      @)
r   r   r<   r+   Zshares_memoryr_   Z_valuesr?   rF   rl   )rI   rr   rJ   rY   r[   rO   r/   r/   r0   test_xs_droplevel_false_view  s"    0

z1TestXSWithMultiIndex.test_xs_droplevel_false_viewc              	   C   sZ   t dddg}tdddgddd	gg|d
}tjtdd |jdddd W 5 Q R X d S )N)r   mr   )r   nr   )r   orX   r   r   r   r   r   r   r   r   r8   )r   r   F)rj   r:   )r   r   r   rA   rB   rC   r<   )rI   r   rY   r/   r/   r0   $test_xs_list_indexer_droplevel_false  s    z9TestXSWithMultiIndex.test_xs_list_indexer_droplevel_falseN)rw   rx   ry   r   r   r   r   r   r   rA   markZparametrizer   r   r   r   r   r   r   r   r   r   r   r   r   r   r/   r/   r/   r0   rz      sH   


 
	!rz   )rD   Znumpyr+   rA   Zpandas.errorsr   Zpandasr   r   r   r   r   r   Zpandas._testingZ_testingrF   Zpandas.tseries.offsetsr	   Zfixturer1   r2   rz   r/   r/   r/   r0   <module>   s    
h