U
    sVccO                     @   s   d dl Z d dlZd dlZd dlmZ d dlZd dlmZm	Z	m
Z
mZmZmZ d dlmZ ejdddgdd Zejd	d
dgdd Zdd ZG dd dZdS )    N)PerformanceWarning)	DataFrameDatetimeIndexIndex
MultiIndexSeries	Timestampzmsg,labels,level)zlabels \[4\] not found in level   a)zlabels \[7\] not found in level   bc              	   C   s   t jdddgdddggddgd	}td
ddg|d}td
ddg|d}tjt| d |j||d W 5 Q R X tjt| d |j||d W 5 Q R X d S )N         r	         r
   r   names
         indexmatchlevel)r   from_arraysr   r   pytestraisesKeyErrordrop)msglabelsr   misdf r'   H/tmp/pip-unpacked-wheel-xj8nt62q/pandas/tests/frame/methods/test_drop.py0test_drop_raise_exception_if_labels_not_in_level   s    	"r)   zlabels,level)r	   r
   )r   r   c                 C   s   t jdddgdddggddgd	}td
ddg|d}td
ddg|d}|j| |dd}t|| |j| |dd}t|| d S )Nr   r   r   r	   r   r   r
   r   r   r   r   r   r   ignore)r   errors)r   r   r   r   r!   tmZassert_series_equalassert_frame_equal)r#   r   r$   r%   r&   Z
expected_sZexpected_dfr'   r'   r(   test_drop_errors_ignore'   s    "r.   c               	   C   sl   t tjdddddgtjddddd	} | jd
dddg  } tj	t
dd | ddg W 5 Q R X d S )Nr   r   r
   r   cZ2012H)freqZperiodscolumnsr   r   r   not found in axisr   )r   nprandomrandnpd
date_rangeiloccopyr   r   r    r!   )r&   r'   r'   r(   9test_drop_with_non_unique_datetime_index_and_invalid_keys5   s    r<   c                
   @   sF  e Zd Zdd Zdd Zdd Zdd Zd	d
dgd	d
dggZej	
deeddgdeeddgdeeddgddd	geeddgddd	ggdd Zej	
dd	d
dgd	d	d
ggej	
dg d	gd
ggdd Zej	
dd	d
dgd	d
d
ggej	
dd	dgddggdd Zej	
dg i eg eg ddeg eg gdd Zd d! Zd"d# Zd$d% Zd&d' Zd(d) Zej	
d*eegd+d, Zd-d. Zej	
d/d0d1d2d3gej	
d4d5d6gd7d8 Zd9d: Zej	
d;d<d<ggd=d> Zd?d@ Z dAdB Z!dCdD Z"dEdF Z#dGdH Z$ej	
dIddgdfdgdJfgdKdL Z%dJS )MTestDataFrameDropc              	   C   s6  t dddgdddgdddggdd	d
gdddgd}d\|j_|j_|d	}|jddd}| |  }}|jd	dd}|d kst|jdddd}|d kst||||fD ]$}|jjdkst|jjdkstqt|jdddgkstd}tj	t
|d |dg W 5 Q R X tj	t
|d |jdgdd W 5 Q R X |jdgdd}	tdd	d
gdd}
t|	j|
 |jd	dgdd}	tdd
gdd}
t|	j|
 |jdgddd}	tdddgdd}
t|	j|
 |jddgddd}	tddgdd}
t|	j|
 |jg dd}	tdd	d
gdd}
t|	j|
 d S )Nr   r   r   r	   r   r   r   r
   r   r/   defr   r3   )firstsecondaxisT)inplacerE   rF   rB   rC   z\['g'\] not found in axisr   gr*   r+   namerE   r+   )r   r   rK   r3   r!   r;   AssertionErrorlistr   r   r    r   r,   assert_index_equal)selfr&   Zdf_dropped_bZdf_dropped_eZdf_inplace_bZdf_inplace_ereturn_valueobjr"   Zdroppedexpectedr'   r'   r(   test_drop_namesF   sJ    
z!TestDataFrameDrop.test_drop_namesc              	   C   s>  t ddddgddddgd}t|jddd|d	g  t|jdd	gd
d|g   t|jdddgdd|jdgd d f  t|jddgdd|jddgd d f  tjtdd |d W 5 Q R X tjtdd |jddd W 5 Q R X tjtdd |ddg W 5 Q R X tjtdd |jddgdd W 5 Q R X tjtdd |jdddgdd W 5 Q R X t|jddd| t|jddgdd|jdddgd d f  t|jdddd| t|jddgddd|d	g  t tt	t
dt
ddtddddgd}t|jddd|dg  t|jdd
d|d  t|g | |tdddg}td|_t|jddd|jdgd d f  t|jddgdd|jg d d f  t tjddtdd}||jdk  }|j||jdk jd d!}|d ks.tt|| d S )"Nr   r   r   r	   r   )ABrU   rD   rV   r3   r   z\[5\] not found in axisr   r   z\['C'\] not found in axisCz#\['C', 'D', 'F'\] not found in axisDFr*   rI   rL   abcr
   r   r3   XYZrowsr   T)r#   rF   )r   r,   r-   r!   locr   r   r    rN   ziprange	set_indexr   r3   r5   r6   r7   r   r   rM   )rP   simpleZnu_dfr&   rS   rQ   r'   r'   r(   	test_dropu   sR    ,,   
&(zTestDataFrameDrop.test_dropc              	   C   s   t jdddgddgd}tddd	gg|d
}|j s:ttddddgddddgdddd	ggd}|jdddgdd}| }|j rtt	|| |j
ddd}tt |j
ddd}W 5 Q R X t	|| d S )N)r
    )b1c1)b2c2r   r/   r   r   r   r	   r\   r
   r>   rf   rg   rh   ri   )r3   data)r   r3   valuesrD   )r   from_tuplesr   r3   Z_is_lexsortedrM   Zpivot_tablereset_indexr,   r-   r!   assert_produces_warningr   )rP   Zlexsorted_miZlexsorted_dfZnot_lexsorted_dfrS   resultr'   r'   r(   "test_drop_multiindex_not_lexsorted   s,     
   z4TestDataFrameDrop.test_drop_multiindex_not_lexsortedc              	   C   s  t dddgdddgdddggdd	d
gdddgd}|d}|jdd}t|| |jddd}|jdd}t|| |jddd}|jdd}t|| |jdgdd}|jdgd}t|| |jdgddjdgdd}|jdgdgd}t|| d}tjt|d |jdd	d W 5 Q R X tjt|d |jdd	d W 5 Q R X d}tjt|d |jdd W 5 Q R X d S )Nr   r   r   r	   r   r   r   r
   r   r/   r>   r?   r@   rA   r   rD   r\   )r#   rE   r   z2Cannot specify both 'labels' and 'index'/'columns'r   )r#   r   )r#   r3   z>Need to specify at least one of 'labels', 'index' or 'columns')r   r!   r,   r-   r   r   
ValueError)rP   r&   Zres1Zres2r"   r'   r'   r(   test_drop_api_equivalence   s8    
z+TestDataFrameDrop.test_drop_api_equivalencer   r   r   actualr
   rj   r   r   r   c              	   C   s   t |jtrdnd }td}tjt|d |jd|dd W 5 Q R X tjt|d |j	jd|dd W 5 Q R X |jdd|dd}t
|| |j	jdd|dd}t
|j	| d S )	Nr   z"['c'] not found in axis"r   r/   r   rE   r   r*   )rE   r   r+   )
isinstancer   r   reescaper   r   r    r!   Tr,   r-   )rP   rs   r   r"   Zexpected_no_errr'   r'   r(   "test_raise_on_drop_duplicate_index   s    
z4TestDataFrameDrop.test_raise_on_drop_duplicate_indexr   drop_labelsc                    s8    fdd|D }t |d }t|t |d d S )Nc                    s   g | ]}| kr|qS r'   r'   ).0ir{   r'   r(   
<listcomp>	  s      z:TestDataFrameDrop.test_drop_empty_list.<locals>.<listcomp>r   r   r!   r,   r-   )rP   r   r{   Zexpected_indexframer'   r~   r(   test_drop_empty_list  s    z&TestDataFrameDrop.test_drop_empty_listr	   r   c              	   C   s.   t jtdd t|d| W 5 Q R X d S )Nr4   r   r   )r   r   r    r   r!   )rP   r   r{   r'   r'   r(   test_drop_non_empty_list  s    z*TestDataFrameDrop.test_drop_non_empty_listempty_listlikezdatetime64[ns]Zdtypec                 C   sP   ddgddgd}t dt dg}t||d}| }||}t|| d S )Nr   r   onetwo)Zcolumn_aZcolumn_bz
2021-01-01r   )r   r   r;   r!   r,   r-   )rP   r   rj   r   r&   rS   ro   r'   r'   r(   2test_drop_empty_listlike_non_unique_datetime_index  s    
zDTestDataFrameDrop.test_drop_empty_listlike_non_unique_datetime_indexc                 C   s:  ddddddgddddddgdd	d
dddgg}t t| }t|}ttjdd|d}|jddd}|jdgdd}t	
|| |jdgdd}|jdgdd}|jdgdd}t	
|| |jddd}|jdgdd}t	
|| |jdgdd}|jddd}|jdddd}|jddgdd}t	
|| d S )Nr
   toproutine1routine2re   ODresult1Zresult2wxwyr	   r   r\   r   rD   )r
   re   re   )r   r   r   )r   r   r   ru   )r   r   re   )r   r   re   )sortedr`   r   rl   r   r5   r6   r7   r!   r,   r-   )rP   ZarraysZtuplesr   r&   ro   rS   r'   r'   r(   test_mixed_depth_drop*  s2    
 z'TestDataFrameDrop.test_mixed_depth_dropc                 C   s   t ddddgtjdddgddddgd	d
ddgddddg }|jddd}t d
d	gdgtjddtjdfgdddgdd}t	|| d S )Nr   r           g      ?g       @r
   r   r/   r   r   r   r	   rU   rV   rW   rX   rU   rV   rW   r   rX   )r   r   r   r   r2   )
r   r5   nanrb   Z
sort_indexr!   r   rl   r,   r-   rP   r&   ro   rS   r'   r'   r(   $test_drop_multiindex_other_level_nanK  s(    


 z6TestDataFrameDrop.test_drop_multiindex_other_level_nanc                 C   s   t ddddgddddgdddd	gdddd
gddddgdddd
gddddgddddgddddgg	ddddgd}|d }|j|dk }|dddg}|j|jdd }||j	|j  }|j|_t
|| d S )Nzx-axr
   g      ?g333333?zz-czr/   g@gffffff@zx-br   gffffff@g@zy-ayzz-bg @var1Zvar2Zvar3Zvar4r\   r   r   r   )r   groupbysizer_   rb   r!   r   rm   r   isinr,   r-   )rP   r&   Zgrp_sizeZdrop_idxZidfro   rS   r'   r'   r(   test_drop_nonuniquec  s(    









z%TestDataFrameDrop.test_drop_nonuniquec                 C   s   |}|j ddgdd}|jddddd	g }t|| |j d
gdd}|jdddd	ddg }t|| |jj ddgddd}|jddddd	g j}t|| |jj d
gddd}|jdddd	ddg j}t|| d S )NbarZquxrB   r   r   r   r   r   r   r   rC   r   r   	   )rE   r   )r!   r:   r,   r-   ry   )rP   Z multiindex_dataframe_random_datar   ro   rS   r'   r'   r(   test_drop_level  s    z!TestDataFrameDrop.test_drop_levelc                 C   s   t dddddgdd}tddd	d	d
g}ttdddtd|d}||d< |jddd}t	d	}|j
jdksxt|j|dd}|j|dk }t|| d S )Nr   r   r	   r   idrJ   Z201603231400Z201603231500Z201603231600Z201603231700r   abr2   ZtstampT)appendFr   )r   r8   Zto_datetimer   r5   ZarangeZreshaperN   rb   r   r   Z	is_uniquerM   r!   r_   r,   r-   )rP   idxZidxdtr&   tsro   rS   r'   r'   r(   "test_drop_level_nonunique_datetime  s"    	 z4TestDataFrameDrop.test_drop_level_nonunique_datetimeboxc           
      C   s   t ddd}t ddd}tj||dd}|dgt| |d}||}t d	dd}tj||dd}|dgt| |d}	t||	 d S )
Nz
2017-10-29zEurope/Berlin)tzz2017-10-29 04:00:00Z15min)r1   r   rt   z2017-10-29 00:15:00)r   r8   r9   lenr!   r,   Zassert_equal)
rP   r   startendr   rj   ro   Zexpected_startZexpected_idxrS   r'   r'   r(   'test_drop_tz_aware_timestamp_across_dst  s    
z9TestDataFrameDrop.test_drop_tz_aware_timestamp_across_dstc                 C   sd   t jddddddgddddddggddgd}ttjdd|d	}|d
g}|jjdks`t	d S )Nr   r   r   r   r   r   r   r   r   )r   r   )r   r   )
r   r   r   r5   r6   r7   r!   r   r   rM   )rP   r   r&   ro   r'   r'   r(   test_drop_preserve_names  s     z*TestDataFrameDrop.test_drop_preserve_names	operation__iadd____isub____imul____ipow__rF   FTc              	   C   s   t dtdi}| }td|d< |d }td F |rN|jdd|d n|jdd|d}t||d t|| W 5 Q R X d S )Nr   r   r   r   rG   )r   ra   r;   r,   rn   r!   getattrr-   )rP   r   rF   r&   rS   r   r'   r'   r(   test_inplace_drop_and_operation  s    z1TestDataFrameDrop.test_inplace_drop_and_operationc                 C   sd   t dddgdddgg}tdddg|d}|jdd}tdgt dgdggd}t|| d S )	Nr   r   r}   jr   r   r   r   )r   r   r   r!   r,   r-   )rP   r$   r&   ro   rS   r'   r'   r(   $test_drop_with_non_unique_multiindex  s
    z6TestDataFrameDrop.test_drop_with_non_unique_multiindexindexerr
   r
   c                 C   sj   t ddgddgg}tdtt|i|d}|jdgd}tdddgit ddgd}t|| d S )	Nr
   r   r   r   r   r   r   )r   r
   )	r   from_productr   ra   r   r!   rl   r,   r-   )rP   r   r   r&   ro   rS   r'   r'   r(   *test_drop_tuple_with_non_unique_multiindex  s    
 z<TestDataFrameDrop.test_drop_tuple_with_non_unique_multiindexc                 C   s~   t dddgdddgdddggdddgd}|jdgdd}t dgdgdggdgd}t|| |jddd}t|| d S )Nr   r   g      @r   r
   r\   rD   r   r   r'   r'   r(    test_drop_with_duplicate_columns  s     z2TestDataFrameDrop.test_drop_with_duplicate_columnsc              	   C   s~   t tjdtjdtjddddddgd}|jdd	d	gd	d
}|jddd	dd	gd	d
}|jdd	d
}t|| d S )Nr   r
   r   r/   r>   r?   r   r   r   rD   r   rW   )r   r5   r6   r7   Ztaker!   r,   r-   )rP   r&   rS   Zdf2ro   r'   r'   r(   !test_drop_with_duplicate_columns2  s    


	z3TestDataFrameDrop.test_drop_with_duplicate_columns2c              	   C   s\   t ddddgi}d}tjt|d |dd}W 5 Q R X t dddgd}t|| d S )	Nr
   r   r   r   ztIn a future version of pandas all arguments of DataFrame\.drop except for the argument 'labels' will be keyword-onlyr   r   r   )r   r,   rn   FutureWarningr!   r-   )rP   r&   r"   ro   rS   r'   r'   r(   test_drop_pos_args_deprecation  s    z0TestDataFrameDrop.test_drop_pos_args_deprecationc                 C   sf   t ddddgi}|j}|jdgddd t|jtg dd || 8 }t|jtg dd d S )	Nr
   r   r   r   TrG   objectr   )r   r
   r!   r,   rO   r3   r   Zmean)rP   r&   r
   r'   r'   r(   .test_drop_inplace_no_leftover_column_reference  s    z@TestDataFrameDrop.test_drop_inplace_no_leftover_column_referencec              	   C   sH   t ttdtdgd}tjtdd |jddd W 5 Q R X d S )Nr   r   zlabels \[5\] not found in levelr   r   r   r   )r   r   r   ra   r   r   r    r!   )rP   r&   r'   r'   r(   (test_drop_level_missing_label_multiindex  s    z:TestDataFrameDrop.test_drop_level_missing_label_multiindexz
idx, levelNc                 C   sd   t dddtjgdd|d|}|jtdtjg|d}t dgdd|d|}t|| d S )Nr   r   d   )r
   r   r   r   )r   r8   ZNArb   r!   r   r,   r-   )rP   Zany_numeric_ea_dtyper   r   r&   ro   rS   r'   r'   r(   test_drop_index_ea_dtype  s     
 z*TestDataFrameDrop.test_drop_index_ea_dtype)&__name__
__module____qualname__rT   rd   rp   rr   rj   r   markparametrizer   rb   rz   r   r   r5   arrayr   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r'   r'   r'   r(   r=   E   sn   /3'	


!

 



	r=   )rw   Znumpyr5   r   Zpandas.errorsr   Zpandasr8   r   r   r   r   r   r   Zpandas._testingZ_testingr,   r   r   r)   r.   r<   r=   r'   r'   r'   r(   <module>   s"    

