U
    b-                     @   sN  d dl Zd dlZd dlZd dlmZ ej	ddddgddggddgddgdd	gd
fdddgdej
gddggdddgdddgddd	gd
fgdd Zej	ddddgddggddgddgdd	gd
fdddgdej
gddgej
dggdddd	gdddd	gddd	d	gd
fgdd Zej	ddddgddgddgd	dgdfdddej
gdddgdddgd	ddgdfgdd Zej	dddddej
gejddgddgdfddddej
gejdddgddej
gdfgdd Zej	d dejd!d"gddgd#d$fdejd!d"d%gddej
gd#d$fgd&d' Zej	d(d)d*d+ Zej	d,ded d-dd.d.d/gied dfded0dd.d.d/gied1fdejjd2d3d4d5gd6d7gd8dd.d.d/giejjd2d3d4gd6d7gd8fded d-dd.d.d/d/gied d-fded0dd.d.d/d/gied0fdejjd2d3d4d5gd6d7gd8dd.d.d/d/giejjd2d3d4d5gd6d7gd8fgd9d: Zej	ddddgddggddgddgd	d	gd
fdddgdej
gddggdddgdddgd	dd	gd
fgd;d< Zejjej	d=ed>ed?fed@edAfed>ed?fgej	dBddCdgfddCddDgfgdEdF Zej	dGejddddej
gdHdHdIdJgdKdLd d/d d giddej
gdMdNejddddej
gdHdHdIdJgdKdLd d/d giddOdNejdddddPgdHdHdIdJgdKdLd d/d d giddQdNejdddddPgdHdHdIdJgdKdLd d/d d giddRdNgdSdT ZdUdV ZdS )W    Nzdropna, tuples, outputsTABg      *@Q^@g     ^@      ?cdeF皙(@      m@g      (@c           	      C   s   dddddgd|dddgdddddgddddd	gg}t j|d
ddddgd}|jd
dg| d }t jj|tdd}| s|jddtj	gdd}t j||d}t
|| d S )Nr   r      r
   r   r   {      r   abr   r   r	   columnsdropnaabnameslevelindexpd	DataFramegroupbysum
MultiIndexfrom_tupleslist
set_levelsnpnantmassert_frame_equal)	r   tuplesoutputsnulls_fixturedf_listdfgroupedmiexpected r1   L/tmp/pip-unpacked-wheel-ck39h295/pandas/tests/groupby/test_groupby_dropna.py:test_groupby_dropna_multi_index_dataframe_nan_in_one_group   s    r3   g*@g     @m@c           
   	   C   s   dddddgd|dddgdddddg|dddd	gd|ddd	gg}t j|d
ddddgd}|jd
dg| d }t jj|tdd}| s|ddtj	gddtj	gg}t j||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+   Znulls_fixture2r,   r-   r.   r/   r0   r1   r1   r2   ;test_groupby_dropna_multi_index_dataframe_nan_in_two_groups3   s    r4   zdropna, idx, outputs)r   r   r   c                 C   s~   ddddgd dddgddddgdddd	gg}t j|d
dddgd}|jd
| d }t j|t j|dd
dd}t|| d S )Nr   r   r
   r   r   r   r   r   r   r   r   r   r   r   r   object)dtypenamer   )r   r   r   r    Indexr'   r(   )r   idxr*   r,   r-   r.   r0   r1   r1   r2   *test_groupby_dropna_normal_index_dataframe_   s    



r:   zdropna, idx, expectedr   r      r   c                 C   s8   t jddddg|d}|jd| d }t|| d S )Nr      r;   r   r   )r   r   )r   Seriesr   r    r'   assert_series_equal)r   r9   r0   serresultr1   r1   r2    test_groupby_dropna_series_level~   s    rA   zdropna, expectedg     @j@     u@	Max Speedr   r7         4@c                 C   sL   t jddddgddddgdd}|jd	d
d	tjg| d }t|| d S )Ng     `x@rB   g      >@rE   ZFalconZParrotrC   rD   r   r   r   )r   r=   r   r%   r&   Zmeanr'   r>   )r   r0   r?   r@   r1   r1   r2   test_groupby_dropna_series_by   s    

rF   r   )FTc                 C   sB   t dddd gdddd gd}|jd| d}|jj| ks>td S )Nr   r   r<   r;   r   r   r   r   )r   r   r   Zgrouperr   AssertionError)r   r-   gbr1   r1   r2   test_grouper_dropna_propagation   s     rJ   z/dropna,input_index,expected_data,expected_index   r<   r   Zabcdabc)r   R)r   r   )r<   rM   )r<   r   numcolr   c                 C   s   t jdddd gdddd gd|d}|jd| d}|t}t j||d}t|| |d	g t}t j||d}t|| |d	 t}t j|d	 |d	d
}t|| d S )Nr   r   r<   r;   rG   r   r   r   r   rD   )	r   r   r   Z	transformlenr'   r(   r=   r>   )r   Zinput_indexZexpected_dataZexpected_indexr-   rI   r@   r0   r1   r1   r2   +test_groupby_dataframe_slice_then_transform   s    "$
rQ   c           	      C   s   dddddgdd dddgdddddgddddd	gg}t j|d
ddddgd}ttdd}|jd
dg| d|}t jj|tdd}| s|j	ddt
jgdd}t j||d}t|| d S )Nr   r   r   r
   r   r   r   r   r   r   r   r   r   r	   r   minr   r   r   r   r   r   )r   r   r    maxr   aggr!   r"   r#   r$   r%   r&   r'   r(   )	r   r)   r*   r,   r-   Zagg_dictr.   r/   r0   r1   r1   r2   -test_groupby_dropna_multi_index_dataframe_agg   s    rU   zdatetime1, datetime2z
2020-01-01z
2020-02-01z-2 daysz-1 dayszdropna, valuesr      c           
   	   C   s   t ddddddg||||||gd}| r6||g}n||tjg}|jd| d	d
ti}t jd
|it j|ddd}	t	||	 d S )Nr   r<   r;   rK      rV   )valuesdtrY   r   rX   )r7   r   )
r   r   r%   r&   r   rT   r    r8   r'   r(   )
r   rX   Z	datetime1Z	datetime2Zunique_nulls_fixtureZunique_nulls_fixture2r-   Zindexesr.   r0   r1   r1   r2   &test_groupby_dropna_datetime_like_data	  s"    
rZ   z#dropna, data, selected_data, levels
         )groupsrX   rX   Zdropna_false_has_nan)idZdropna_true_has_nanr   Zdropna_false_no_nanZdropna_true_no_nanc           
      C   s   t |}|jd| d}|dd }tt|d |d }t jj|dd gd}| sf|rf|j|dd}t j||d}	t	
||	 d S )	Nr^   r   c                 S   s   t dtt| iS )NrX   )r   r   rangerP   )grpr1   r1   r2   <lambda>X      z@test_groupby_apply_with_dropna_for_multi_index.<locals>.<lambda>rX   r   r   r   )r   r   r   applytuplezipr!   r"   r$   r'   r(   )
r   dataZselected_datalevelsr-   rI   r@   Z	mi_tuplesr/   r0   r1   r1   r2   .test_groupby_apply_with_dropna_for_multi_index0  s    &
ri   c               	   C   s   dt jddt jgdddddgd} t| }|jd	d
d}|j}t j}dt jddg|ddt jdg|dt jt jddg|di}t|	 |	 D ]\}}t
|| qt t| d stt| dd ddgkstd S )NZg1Zg2r   r   r<   r;   rK   )groupr   rj   Fr   )r6   )r%   r&   r   r   r   indicesZintparrayrf   rX   r'   Zassert_numpy_array_equalisnanr#   keysrH   )rg   r-   r.   r@   r6   r0   Zresult_valuesZexpected_valuesr1   r1   r2   test_groupby_nan_includede  s     "
   ro   ) Znumpyr%   ZpytestZpandasr   Zpandas._testingZ_testingr'   markZparametrizer&   r3   r4   r:   r=   rA   rF   rJ   Z
RangeIndexr#   r!   r"   rQ   rU   Zarm_slowZ	TimestampZ	TimedeltaZPeriodrZ   paramri   ro   r1   r1   r1   r2   <module>   s@  




 

$



$
  
&
 
 

"
#
