a
    Pf[                     @   s  d dl Zd dlZd dlmZ d dlmZ d dlZd dl	m
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g dg dg d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g dg dg d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g dg dg d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g d&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g d.d"d#ejgd,d-fgd/d0 Zejd1d2d3d4 Zejd5ed d6ed7ejjd8d9gd:d;gd<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gg dg d?g d@dfgdAdB ZejjejdCedDedEfe dFe dGfe!dDe!dEfgejdHddId$gfdg dJfgdKdL Z"ejdMej#dd"d"d#ejgg dNdOdPg dQid"d#ejgdRdSej#dd"d"d#ejgg dNdOdPg dTiddUdSej#dg dVg dNdOdPg dQiddWdSej#dg dVg dNdOdPg dQiddXdSgdYdZ Z$ejd[dd"gd"d#ggejd\d"gd"d#ggejd]ddgd^d_ Z%d`da Z&dbdc Z'ejdde(deejdfddgdhdidjdkdldmdndodndpdqdrej#dsejj)edtdudvdwdxdygejdzddgd{d| Z*ejdzddgejdfe+dgd}d~ Z,ejdg ddd Z-dd Z.ejdddgdd Z/dd Z0dd Z1dS )    N)pa_version_under10p1)na_value_for_dtype)get_groupby_method_argszdropna, tuples, outputsTAB      *@Q^@     ^@      ?cdeF)r   皙(@r   )r         m@r	   )r         (@r
   c           	      C   s   g dd|dddgg dg dg}t j|g dd	}|jd
dg| d }t jj|tdd}| sx|jddtj	gdd}t j||d}t
|| d S )Nr   r      r   r   r   r   r   r   r   r   r   {      r   r   r   r   r
   abr   r   r   columnsr   r   dropnaabnamesr   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 r;   a/var/www/ai-form-bot/venv/lib/python3.9/site-packages/pandas/tests/groupby/test_groupby_dropna.py:test_groupby_dropna_multi_index_dataframe_nan_in_one_group   s    r=   r   )r   g*@r   r
   )r   g     @m@r	   r
   )r   r   r
   r
   c           
   	   C   s   g dd|dddgg d|dddd	gd|ddd	gg}t j|g d
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   r3   r4   r5   nulls_fixture2r6   r7   r8   r9   r:   r;   r;   r<   ;test_groupby_dropna_multi_index_dataframe_nan_in_two_groups8   s    r?   zdropna, idx, outputs)r   r   r   )r   r   r   )r	   r   r   )r
   r   r   c                 C   sj   g dg dg dg dg}t j|g dd}|jd| d }t j|t j|d	dd
d}t|| d S )N)r   r   r   r   )Nr   r   r   )r   r   r   r   )r   r   r   r
   )r   r   r   r   r   r   r   objectdtypenamer$   )r'   r(   r)   r*   Indexr1   r2   )r   idxr4   r6   r7   r8   r:   r;   r;   r<   *test_groupby_dropna_normal_index_dataframed   s    rF   zdropna, idx, expectedr   r      r$   )rG   rG   rG   c                 C   s4   t jg d|d}|jd| d }t|| d S )N)r      rG   rG   r$   r   )r#   r   )r'   Seriesr)   r*   r1   assert_series_equal)r   rE   r:   serresultr;   r;   r<    test_groupby_dropna_series_level   s    rM   zdropna, expected     @j@     u@	Max Speedr%   rC   )rN   rO         4@c                 C   sD   t jg dg ddd}|jdddtjg| d }t|| d S )N)g     `x@rO   g      >@rR   )FalconrS   ParrotrT   rP   rQ   r   r   r   )r'   rI   r)   r/   r0   meanr1   rJ   )r   r:   rK   rL   r;   r;   r<   test_groupby_dropna_series_by   s    rV   r   )FTc                 C   s:   t g dg dd}|jd| d}|jj| ks6J d S )Nr   r   r   Nr   rH   rG   Nr   r   r   r   )r'   r(   r)   _grouperr   )r   r7   gbr;   r;   r<   test_grouper_dropna_propagation   s    r\   r%      abcd)r   rH   )Rr   numcolr    c                 C   s   dddd| rt jndgi}tjg dg d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   rH   r   rW   rX   rY   r$   r   r   rQ   )r/   r0   r'   r(   r)   	transformlenr1   r2   rI   rJ   )r   r%   Zexpected_datar7   r[   rL   r:   r;   r;   r<   +test_groupby_dataframe_slice_then_transform   s    

rd   )r   r   r	   )r
   r   r
   c           	      C   s   g dg dg dg dg}t j|g dd}dd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   Nr   r   r   r   r   r   r   r*   maxminr   r   r   r   r   r    r   r   r"   r$   )r'   r(   r)   aggr+   r,   r-   r.   r/   r0   r1   r2   )	r   r3   r4   r6   r7   Zagg_dictr8   r9   r:   r;   r;   r<   -test_groupby_dropna_multi_index_dataframe_agg   s    rh   zdatetime1, datetime2z
2020-01-01z
2020-02-01z-2 daysz-1 dayszdropna, valuesr   )r   rG      c           
   	   C   s~   t g d||||||gd}| r.||g}n||tjg}|jd| dddi}t jd|it j|ddd}	t||	 d S )	N)r   rH   rG   r]      ri   )valuesdtrl   r   rk   r*   rC   r$   )	r'   r(   r/   r0   r)   rg   rD   r1   r2   )
r   rk   Z	datetime1Z	datetime2Zunique_nulls_fixtureZunique_nulls_fixture2r7   Zindexesr8   r:   r;   r;   r<   &test_groupby_dropna_datetime_like_data   s"    
rn   z#dropna, data, selected_data, levels)
   ro         )groupsrk   rk   )r   r   r   r   Zdropna_false_has_nan)id)r   r   r   Zdropna_true_has_nan)r   r   r   r   Zdropna_false_no_nanZdropna_true_no_nanc                 C   s   t |}|jd| d}d}tjt|d |dd }W d    n1 sN0    Y  tt|d |d }t j	j
|dd gd}	| s|r|	j|dd	}	t j||	d
}
t||
 d S )Nrr   r   z7DataFrameGroupBy.apply operated on the grouping columnsmatchc                 S   s   t dtt| iS )Nrk   )r'   r(   rangerc   )grpr;   r;   r<   <lambda>I      z@test_groupby_apply_with_dropna_for_multi_index.<locals>.<lambda>rk   r    r"   r$   )r'   r(   r)   r1   assert_produces_warningDeprecationWarningapplytuplezipr+   r,   r.   r2   )r   dataZselected_dataZlevelsr7   r[   msgrL   Z	mi_tuplesr9   r:   r;   r;   r<   .test_groupby_apply_with_dropna_for_multi_index  s    &
,r   input_indexkeysseriesc                 C   s   t dtjgddgddgd}||}|r8|d }n | ddgkrX|dgkrX|dg }| d urj|| }|j|dd	}|r|d }| }t|| d S )
Nr   rH   rG   r   r   r   r   r   r   Fr   )	r'   r(   r/   r0   	set_indexr)   r*   r1   assert_equal)r   r   r   objr:   r[   rL   r;   r;   r<   )test_groupby_dropna_with_multiindex_inputV  s$    



r   c               	   C   s   dt jddt jgg d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	 sJ t| dd	 ddgksJ d S )Ng1g2)r   r   rH   rG   r]   )groupr   r   Fr   r   rH   rB   rG   r   r]   )r/   r0   r'   r(   r)   indicesZintparrayr~   rk   r1   Zassert_numpy_array_equalisnanr-   r   )r   r7   r8   rL   rB   r:   Zresult_valuesZexpected_valuesr;   r;   r<   test_groupby_nan_includedt  s    
r   c                  C   sT   t jtjddggg dd} | ddg} | jddgdd }| }t|| d S )	Nr   r   r   r   r   r   Fr   )	r'   r(   r/   r0   r   r)   firstr1   r2   )r7   rL   r:   r;   r;   r<   &test_groupby_drop_nan_with_multi_index  s
    r   sequence_indexQ   rB   ZUInt8ZInt8ZUInt16ZInt16ZUInt32ZInt32ZUInt64ZInt64ZFloat32ZFloat64categorystringstring[pyarrow]zpyarrow is not installed)reason)Zmarksdatetime64[ns]	period[d]zSparse[float]test_seriesc                    s  d  fddtdD }|dv r4ddtjdn&|d	v rLd
dtjdnddtjdttjfdd|D |dg dd}|jddd|dd}|r|d }|	 }i }t
|D ]\}	}
||
d|	 ||
< q|dkr
tjfdd|D |d jjdd}n^t|trJ|drJtjtjfdd|D |ddd}ntjfdd|D |dd}tj| |dd d}|s| }|s| }|d ur|dr|d ||d< t|| d S )N c                    s(   g | ] }d ddd d|  d  qS )xyz)r   r   rH   rG   r;   ).0k)r   r;   r<   
<listcomp>  ry   z(test_no_sort_keep_na.<locals>.<listcomp>r]   )r   r   r   r   )r   r   r   )r   r   z
2016-01-01z
2017-01-01r   rH   c                    s   g | ]} | qS r;   r;   r   labeluniquesr;   r<   r     ry   r   )r   r   rH   rG   )keyr   r   F)r   sortas_indexobservedr   r   r   c                    s   g | ]} | qS r;   r;   r   r   r   r;   r<   r     ry   rm   ZSparsec                    s   g | ]} | qS r;   r;   r   r   r;   r<   r     ry   c                    s   g | ]} | qS r;   r;   r   r   r;   r<   r     ry   rA   )r%   rC   rB   )joinrv   r'   ZNAr/   r0   r(   rI   r)   r*   	enumerategetZCategoricalIndexcat
categories
isinstancestr
startswithrD   r   rk   Zto_framereset_indexastyper1   r   )r   rB   r   r   sequencer7   r[   rL   ZsummedrE   r   r%   r:   r;   )r   r   r<   test_no_sort_keep_na  sN    #

r   c                 C   s   t dddgi}t j||g|d}|r0|d n|}|j|d| d}| }	t t|jg}
t jddgi|
d}|rt	|	|d  nt
|	| d S )	Nr   r   rH   r   F)r   r   rG   r$   )r'   r(   rI   r)   r*   rD   r   rB   r1   rJ   r2   )r   rB   r5   r>   r   r7   rr   r   r[   rL   r%   r:   r;   r;   r<   test_null_is_null_for_dtype  s    r   
index_kind)rv   singlemultic                 C   sR  t t jdjg dddd }ttj|g ddtdd}|	 }tj|g d	d
d
|d< |dkrxdg}n`|dkrdg}|d}|d}n<ddg}|d |d< |ddg}|d |d< |ddg}t| |}	t| |}
| dkr$|dkr$|	d j|df}	|
d j|df}
|j|d|||d}|s| dv rtjtdd t|| |	  W d    n1 s|0    Y  d S |j|||dd}t|| |
  }|d jd
g|d< |dkr|d jd
g|d< |r|dkr|ddg}n
|d}n6|dkrF| dkrF|jdd}|dkrF|jdd}| dv r|dkr|d j }|dkrdd |D }tj|g dd|d< ndd |D }||d< | dkr|jddid}|r|d d }|s|dks| dkrd }nt}d }tj||d t|| |	 }W d    n1 s80    Y  t|| d S )!NrH   r   rH   N   sizer   rH   rG   r   rp   r   r   )r   rH   rG   r]   r]   r   rv   r   Zx2Zcorrwithr   r   Fr   r   r   r   )idxminidxmaxz(empty group due to unobserved categoriesrt   T)r   r   r   r   r   )r   r   r   c                 S   s   g | ]}|d krt jn|qS )r]   r/   r0   r   r;   r;   r<   r   6  ry   z-test_categorical_reducers.<locals>.<listcomp>c                 S   s$   g | ]}|d krt jt jfn|qS ))r]   r]   r   r   r;   r;   r<   r   9  ry   z*A grouping .* was excluded from the result)r/   appendrandomdefault_rngchoicer'   r(   Categoricalrv   copyfillnar   r   Zdropr)   pytestZraises
ValueErrorgetattrr   r   Zremove_categoriesrk   tolistrenameFutureWarningr1   rz   r   )Zreduction_funcr   r   r   r   rk   r7   Z	df_filledr   argsZargs_filled	gb_keepnaZ	gb_filledr:   warnr   rL   r;   r;   r<   test_categorical_reducers  s|    "



.




.r   c                 C   sX  |dkr"d}|  tjj|dd ttjdjg dddd }t	
t	j|g d	d
tdd}t||}||d   d }	|dkrttt|	}
nx|dkr|r|r|d jddd }q|d jddd }n|jd |	jd  d  }t|	|g }
nt|	|| }
t	
d|
i}|jdd|||d}|jdd||d}d}|dkrtjt|d t|d| }W d    n1 s0    Y  nt||| }t||| }t||d   j |j D ]0\}}|jdkr||j|< n||j|df< q|dkr4||d  ||@   d7  < |dvrH| d}t!|| d S )Nr   zBGH#49651 fillna may incorrectly reorders results when dropna=FalseF)r   strictrH   r   r   r   r   r   rp   r   r   r   ZcumcountZngroupr   r   r   r   T)r   r   r   zLThe default fill_method='ffill' in DataFrameGroupBy.pct_change is deprecated
pct_changert   )Zrankdiffr   shiftZint64)"Zapplymarkerr   markZxfailr/   r   r   r   r   r'   r(   r   rv   r   Zisnullr-   rc   Znuniqueilocr%   r   r)   r1   rz   r   r~   r   rk   ZravelndimZnotnullger   r   )requestZtransformation_funcr   r   r   r   rk   r7   r   Znull_group_valuesZnull_group_dataZna_groupZnull_group_resultr   Z	gb_dropnarL   r:   r   valuer;   r;   r<   test_categorical_transformersM  sR    "



0
"

r   methodheadtailc           
      C   s   t jdg dd}ttj|g ddtt|d}|j	dd|||d	}t
||  }| d
krt|d d d }|dk|dk dk@ |dk|dk dk@ B |d k|d k dk@ B }| d
kr|d d d }|| }	t||	 d S )NrH   r   rq   r   r   r   r   Fr   r   r   rj   )r/   r   r   r   r'   r(   r   rv   rc   r)   r   Zcumsumr1   r2   )
r   r   r   r   rk   r7   r[   rL   maskr:   r;   r;   r<   test_categorical_head_tail  s$    r   c                  C   st   t jdg dd} ttj| g ddtt| d}|j	dddd	}|
d
d }| }t|| d S )NrH   r   rq   r   r   r   r   Fr   r   c                 S   s   |   S Nr*   r   r;   r;   r<   rx     ry   z&test_categorical_agg.<locals>.<lambda>)r/   r   r   r   r'   r(   r   rv   rc   r)   rg   r*   r1   r2   rk   r7   r[   rL   r:   r;   r;   r<   test_categorical_agg  s    r   c                  C   sv   t jdg dd} ttj| g ddtt| d}|j	dddd	}|
d
d }|
d}t|| d S )NrH   r   rq   r   r   r   r   Fr   c                 S   s   |   S r   r   r   r;   r;   r<   rx     ry   z,test_categorical_transform.<locals>.<lambda>r*   )r/   r   r   r   r'   r(   r   rv   rc   r)   rb   r1   r2   r   r;   r;   r<   test_categorical_transform  s    
r   )2numpyr/   r   Zpandas.compat.pyarrowr   Zpandas.core.dtypes.missingr   Zpandasr'   Zpandas._testingZ_testingr1   Zpandas.tests.groupbyr   r   Zparametrizer0   r=   r?   rF   rI   rM   rV   r\   Z
RangeIndexr-   r+   Zfrom_productrd   rh   Zarm_slow	TimestampZ	TimedeltaZPeriodrn   paramr   r   r   r   rv   Zskipifr   r@   r   r   r   r   r   r   r;   r;   r;   r<   <module>   sr  

 

$






 




#

8
S<
