a
    PfP                     @   sp  d Z ddlZddlm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mZ ddlmZ ddlmZ dd Zd	d
 Zdd Zdd Zdd Zdd Zdd Zdd Zdd Z d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*d+ Z*d,d- Z+d.d/ Z,ej()d0e-ed1d2d3d4ife.ed1g d2g d3d4ifd5d6 ed1d2d3d4ifd7d6 ed1g d2g d3d4ifgd8d9 Z/ej()d0e-ed2d3gd:d;gd1d<fe.eg d2g d3gd:d;gd1d<fd=d6 ed2d3gd:d;gd1d<fd>d6 eg d2g d3gd:d;gd1d<fgd?d@ Z0dAdB Z1dCdD Z2dEdF Z3ej()dGe4e5e6e7e8e9e:e;gdHdI Z<dS )Jz
test all other .agg behavior
    N)partial)SpecificationError)	DataFrameIndex
MultiIndexPeriodIndexSeries
date_rangeperiod_range)pprint_thingc                  C   s   t tjddtjddg dg dd} | d}dd }tjtd	d
 |	|g W d    n1 sx0    Y  tjtd	d
 |	| W d    n1 s0    Y  d S )N      )ar   br   r   )onetwor   r   r   )Zdata1Zdata2key1key2r   c                 S   s   |   |   S N)maxmin)Zarr r   b/var/www/ai-form-bot/venv/lib/python3.9/site-packages/pandas/tests/groupby/aggregate/test_other.pypeak_to_peak)   s    z5test_agg_partial_failure_raises.<locals>.peak_to_peakzunsupported operand typematch)
r   nprandomdefault_rngstandard_normalgroupbypytestraises	TypeErroragg)dfgroupedr   r   r   r   test_agg_partial_failure_raises   s    
*r'   c                  C   s   g dg dg dg} t dd | D dd | D dd | D d}d	d | D } t d
d | D dd | D dd | D d}|d |d   |d< |dd}|d |d   |d< |dd}t|t|ksJ d S )N)   
2012-01-01      ?)r   z
2012-01-02g       @)   Ng      @c                 S   s   g | ]}|d  qS r   r   .0xr   r   r   
<listcomp>8       z,test_agg_datetimes_mixed.<locals>.<listcomp>c                 S   s   g | ]}|d  qS r(   r   r-   r   r   r   r0   9   r1   c                 S   s   g | ]}|d  qS r   r   r-   r   r   r   r0   :   r1   )keydatevaluec                 S   s<   g | ]4}|d  |d r,t j|d d nd|d gqS )r   r(   z%Y-%m-%dNr   )dtdatetimestrptimer5   )r.   rowr   r   r   r0   >   s
    c                 S   s   g | ]}|d  qS r,   r   r-   r   r   r   r0   I   r1   c                 S   s   g | ]}|d  qS r2   r   r-   r   r   r   r0   J   r1   c                 S   s   g | ]}|d  qS r3   r   r-   r   r   r   r0   K   r1   r6   weightsr5   sum)r   r<   r    	aggregatelen)dataZdf1Zdf2Zgb1Zgb2r   r   r   test_agg_datetimes_mixed3   s*    	r@   c                  C   s   t dddd} ttjdd| d}|jdd	 }t|j	t
sHJ t d
ddd}ttjdt||d}ttjdt||d}t||d}||j	j}t| d S )Nz2012-1-1Mr+   )freqperiodsr   )r+   r   indexr   )levelz1999-01r   )startrC   rB   )s1s2)r
   r   r   r   r   r   r    r<   
isinstancerE   r   r   r>   	from_dictmonthlist)Zprngr%   rsrE   rH   rI   r&   r   r   r   test_agg_period_indexX   s    rO   c                  C   s  t g dtddddd} d | jg ddf< | jg d	 d
}| d
}t| | t|d| t|ddi| t	|j
 |d  t	|j
d|d  | jg d d
}| d
}t| | t|d| t|ddi| t	|j
 |d  t	|j
d|d  tg dttdd
ddd}t	|j
t| t	|j
 | tg dttdd
ddd}t	|j
 | d S )N)ArP   BrQ   CrR   DrS   z1/1/2011   h)rC   rB   )classtime)r   r(   r   r   rW   )r   r+         rV   first)r   r+   rX      last)r   r   r   r   ZABCDnamerE   r^   )r   r(   r(   r   )r   r	   locZ	set_indexr    tmassert_frame_equalrZ   r$   assert_series_equalrW   r\   r   r   rM   r>   sizecount)r%   expr&   r   r   r   *test_agg_dict_parameter_cast_result_dtypesg   s2    

rg   c                  C   s^   dd t dD } td}t|| d}|dd t}|dd  }t|| d S )Nc                 S   s   g | ]}t d |d dqS )i  r(   )r7   r8   r-   r   r   r   r0      r1   z0test_agg_cast_results_dtypes.<locals>.<listcomp>   Zaaabbbbbbccd)XYri   rj   )	rangerM   r   r    r$   r>   re   ra   rc   )uvr%   resultexpectedr   r   r   test_agg_cast_results_dtypes   s    rp   c                  C   s   t g dg dg dd} t dg dig dd}d|j_| ddg  }t|| t g dg dd	g dd}d|j_| ddd
g  }t|| d S )N)r(   r   r+   rX   r   )r(   r   r   rX   r   r   r   cr   )r(   g      @rX   r   )r(   r   rX   r   rD   r   )r   rr   rr   )r   rE   r^   r    meanra   rb   )r%   ro   rn   r   r   r   test_aggregate_float64_no_int64   s    rt   c            	      C   s  t g dg dtjddd tdd} | ddg}|d	  }|d	  }|d
  }|d
  }|d
 	ddg}t
j||gdd}ddg|_tj||dd |	ddg}t
j||||gdd}td	d
gddgg|_tj||dd |d
d	g 	ddg}t
j||||gdd}td
d	gddgg|_tj||dd |	ddd}t
j||gdd}tj||dd |	ddgddgd}t
j||||gdd}td	d
gddgg|_d}tjt|d( |d
d	g 	ddd W d    n1 s0    Y  d S )Nfoobarrv   rw   rv   rw   rv   rv   r   r   r   r   r   r   r   r   r   rT   r*   rP   rQ   rR   rS   rP   rQ   rR   rS   r<   rs   r(   )ZaxisT)Z
check_likerR   rS   z&Column\(s\) \['r', 'r2'\] do not existr   )rr2)r   r   r   r   r   aranger    rs   r<   r$   pdconcatcolumnsra   rb   r   Zfrom_productr!   r"   KeyError)	r%   r&   Zc_meanZc_sumZd_meanZd_sumrn   ro   msgr   r   r   test_aggregate_api_consistency   sB    	
r   c                  C   s  t g dtdtdd} d}tjt|d6 | dddd	gid
ddgid W d    n1 sj0    Y  d}tjt|d, | dddg dd	i W d    n1 s0    Y  d}tjt|d& | djddi W d    n1 s0    Y  d S )N)r(   r(   r(   r   r   r   rP   rQ   rR   nested renamer is not supportedr   rP   rv   r<   r   rw   re   r   )rQ   rR   z!Column\(s\) \['ma'\] do not existrQ   rR   ma)	r   rk   r!   r"   r   r    r$   r   rQ   )r%   r   r   r   r   "test_agg_dict_renaming_deprecation   s    
":r   c                  C   s   t g dg dtjddd tdd} | ddg}d	}tjt	|d
& |d 
dddgi W d    n1 s~0    Y  tjt	|d
$ |d 
ddd W d    n1 s0    Y  d S )Nru   rx   r   rT   r*   ry   rP   rQ   r   r   rS   rR   r<   stdrz   )r   r   r   r   r   r}   r    r!   r"   r   r$   r%   gr   r   r   r   test_agg_compat   s    	4r   c                  C   st  t g dg dtjddd tdd} | ddg}d	}tjt	|d
0 |
dddgidddgid W d    n1 s0    Y  tjt	|d
0 |dddgidddgid W d    n1 s0    Y  tjt	|d
( |d tjtjd W d    n1 s0    Y  tjt	|d
( |d tjtjd W d    n1 sf0    Y  d S )Nru   rx   r   rT   r*   ry   rP   rQ   r   r   rR   rs   r<   rS   )r1r|   rar   rbrz   )result1result2)rS   r   )r   r   r   r   r   r}   r    r!   r"   r   r=   r$   r<   rs   r   r   r   r   test_agg_nested_dicts  s"    	>>8r   c                  C   sd   t tjdjddd} dd }tjtdd  | d		| W d    n1 sV0    Y  d S )
Nr   
   )   r   )rd   c                 S   s    t d t |   tdd S )Nz(----------------------------------------test)r   Z	to_stringr#   )r%   r   r   r   raiseException!  s    z=test_agg_item_by_item_raise_typeerror.<locals>.raiseExceptionr   r   r   )
r   r   r   r   Zintegersr!   r"   r#   r    r$   )r%   r   r   r   r   %test_agg_item_by_item_raise_typeerror  s    r   c                  C   sX   t tjdtjdtdddd} | dd dd g}|d	}| }t	|| d S )
Nr   dtypez
2020-01-01)rC   rD   c                 S   s   | j S r   )yearr/   r   r   r   <lambda>.  r1   z*test_series_agg_multikey.<locals>.<lambda>c                 S   s   | j S r   )rL   r   r   r   r   r   .  r1   r<   )
r   r   r}   Zfloat64r	   r    r$   r<   ra   rc   )tsr&   rn   ro   r   r   r   test_series_agg_multikey*  s    
r   c               	   C   s   t g dg dg dtjddtjddtjddd} dd }| d	d
g|}| d	d
gdd }t|| d S )N)rv   rv   rv   rv   rw   rw   rw   rw   rv   rv   rv   )r   r   r   r   r   r   r   r   r   r   r   )dullr   shinyr   r   r   r   r   r   r   r   r      )rP   rQ   rR   rS   EFc                 S   s   t | jjdksJ dS )Nr   rv   )r>   valuesbaser   r   r   r   bade  s    z.test_series_agg_multi_pure_python.<locals>.badrP   rQ   c                 S   s   dS )Nrv   r   r   r   r   r   r   j  r1   z3test_series_agg_multi_pure_python.<locals>.<lambda>)	r   r   r   r   r   r    r$   ra   rb   )r?   r   rn   ro   r   r   r   !test_series_agg_multi_pure_python5  s    /r   c                  C   s   dd } t g dg dtdddtdddtdddtdddgd	}|d
}|| g}|jjd |_|| }t|| d S )Nc                 S   s   t j|  ddS )Nr(   )q)r   Z
percentileZdropna)r   r   r   r   P1q  s    z test_agg_consistency.<locals>.P1)r(   r   r+   rX   )r            i  r   r   r   )Zcol1Zcol2r5   r5   r   )	r   r7   r5   r    r$   r   Zlevelsra   rb   )r   r%   r   ro   rn   r   r   r   test_agg_consistencyn  s     

r   c               	   C   s   t ddgddgdtj} G dd d}ttjdd	 d
d	 tt| g}| dd}|D ]l}|tu sv|tju rztnd }d}t	j
||d  | d|}W d    n1 s0    Y  t	|| q`d S )Nr(   r   r+   rX   )rv   rw   c                   @   s   e Zd Zdd ZdS )z$test_agg_callables.<locals>.fn_classc                 S   s   t |S r   r<   )selfr/   r   r   r   __call__  s    z-test_agg_callables.<locals>.fn_class.__call__N)__name__
__module____qualname__r   r   r   r   r   fn_class  s   r   c                 S   s   t | S r   r   r   r   r   r   r     r1   z$test_agg_callables.<locals>.<lambda>c                 S   s   |   S r   r   r   r   r   r   r     r1   rv   r<   zusing DataFrameGroupBy.sumr   )r   Zastyper   Zint64r<   r   r    r$   FutureWarningra   assert_produces_warningrb   )r%   r   Zequiv_callablesro   Zecallwarnr   rn   r   r   r   test_agg_callables  s     	.r   c                  C   s   t dtg dgdtg dgdtg dggddgd} | d}tg d	gtg dgg}tddgdd
}dg}t |||d}|jdd}t|| |jddd}t|| d S )Nr(   )r   r      )(   2   <   r   )r   r   r   categoryZ	arraydata)r   )r   F   Z   r]   rE   r   FZnumeric_onlyr<   )	r   r   arrayr    r   r<   ra   rb   r$   )r%   gbZexpected_dataZexpected_indexZexpected_columnro   Zaltrn   r   r   r   test_agg_over_numpy_arrays  s     
 r   	as_periodTFc                 C   s(  t dddd}| r"|d d}tg d|d}|d}|d	 d
d }t|d d d d	d}d|j_t	
|| |d	 dd }tddgd	d}d|j_t	
|| |d	 dd }ttjddtjddgd	d}d|j_| rttjdtjdgd	d}d|j_t	
|| d S )Nr)   rX   UTC)rC   tzrS   )r   r   r(   r(   r   r   r   r   c                 S   s
   | j d S Nr   ilocr   r   r   r   r     r1   z6test_agg_tzaware_non_datetime_result.<locals>.<lambda>r   r]   c                 S   s   | j d jS r   )r   r   r   r   r   r   r     r1   i  c                 S   s   | j d | j d  S )Nr   r   r   r   r   r   r     r1   r(   )days)r	   Ztz_localizeZ	to_periodr   r    r$   r   rE   r^   ra   rc   r~   Z	TimedeltaoffsetsZDay)r   Zdtir%   r   rn   ro   r   r   r   $test_agg_tzaware_non_datetime_result  s(    
  r   c                     s  t jddd td fddtdD d} | d	d
 djd }| d	d
 dd jd }| d	d
  jd }| ksJ | ksJ | ksJ dd tddD }td	d
gd |d} | d}| d jd   |dd jd ksJ  |	dd jd ks J  |
 d jd ks<J d}tjt|d.  |dd jd ksnJ W d    n1 s0    Y  | d jd   | d jd ksJ d}tjt|d.  |dd jd ksJ W d    n1 s 0    Y  d S )Nz2016-01-01 12:00:00
US/Pacificr   r(   c                    s   g | ]} t j|d  qS ))minutes)r7   	timedelta)r.   nnr   r   r   r0     r1   z0test_agg_timezone_round_trip.<locals>.<listcomp>r   r   r   r   r   r   c                 S   s
   t | S r   )r   r   r   r   r   r   r     r1   z.test_agg_timezone_round_trip.<locals>.<lambda>c                 S   s$   g | ]}t jd |ddddqS )z	2016-01-0dz	 12:00:00r   r   )r~   	Timestamp)r.   ir   r   r   r0     s   r   r   rP   rQ   rP   rQ   z7DataFrameGroupBy.apply operated on the grouping columnsr   c                 S   s
   | j d S r   r   r   r   r   r   r     r1   )r   r(   c                 S   s
   | j d S )Nr   r   r   r   r   r   r     r1   )r~   r   r   rk   r    r$   r   r   ZnthheadrZ   ra   r   DeprecationWarningapplyr\   )r%   r   r   Zresult3datesr&   r   r   r   r   test_agg_timezone_round_trip  s2     
>r   c                  C   s   t ddgddgddggtd} | d } tg d	tjd}t dg d
i|td}d|j_| djdd}t	
|| | djdd}|g  }t	
|| d S )Nr(   r   r+   rX   r   rY   r   l    )l            l           l           )l           l           l           )rE   r   r   Fr   T)r   objectr   r   Zuint64rE   r^   r    r<   ra   rb   )r%   rE   ro   rn   r   Z	expected2r   r   r   test_sum_uint64_overflow  s     

r   zstructure, expectedrR   )r(   r(   r(   )r+   rX   rX   ))r(   r(   )r+   rX   c                 C   s   t | S r   tupler   r   r   r   r   $  r1   r   c                 C   s   t | S r   rM   r   r   r   r   r   (  r1   c                 C   sL   t g dg dg dd}|ddg| }ddg|j_t|| d S )Nr(   r(   r(   r+   r+   r+   r(   r(   r(   rX   rX   rX   r(   r(   r(   r+   rX   rX   r   rP   rQ   )r   r    r=   rE   namesra   rb   Z	structurero   r%   rn   r   r   r   test_agg_structs_dataframe  s    r   r(   r+   r_   c                 C   s   t | S r   r   r   r   r   r   r   <  r1   c                 C   s   t | S r   r   r   r   r   r   r   =  r1   c                 C   sH   t g dg dg dd}|dd | }d|j_t|| d S )Nr   r   r   r   rP   rR   )r   r    r=   rE   r^   ra   rc   r   r   r   r   test_agg_structs_series7  s    r   c                 C   s   g d}t tjg d|dg dd}d}tjt|d( |jd| d	jt	j
}W d    n1 sh0    Y  tg d
tjg d|dddd}| r||dk }t|| d S )Nrq   )r   r   r   )
categories)r(   r   r+   r   zusing SeriesGroupBy.sumr   rP   )observed)r+   r+   r   )r   r^   rQ   r_   r   )r   r~   ZCategoricalra   r   r   r    rQ   r$   r   Znansumr   ZCategoricalIndexrc   )r   r   r%   r   rn   ro   r   r   r   test_agg_category_nansumK  s    6r   c                  C   s   t dd tdD dd tdD d} | jdddd}|dd	d
 i}t dd tdD dd tdD d}t|| d S )Nc                 S   s   g | ]}t |qS r   strr-   r   r   r   r0   _  r1   z+test_agg_list_like_func.<locals>.<listcomp>r+   r   rP   F)Zas_indexsortrQ   c                 S   s   t | S r   r   r   r   r   r   r   a  r1   z)test_agg_list_like_func.<locals>.<lambda>c                 S   s   g | ]}t |qS r   r   r-   r   r   r   r0   c  r1   c                 S   s   g | ]}t |gqS r   r   r-   r   r   r   r0   c  r1   )r   rk   r    r$   ra   rb   )r%   r&   rn   ro   r   r   r   test_agg_list_like_func]  s    *$r   c                  C   sx   t ddgtjdddtjdddgd} | ddd	d
 i}t tjdddgtdgdddgd}t|| d S )Nr(   z
2018-01-01r   r   z
2018-01-02)tagr5   r   r5   c                 S   s
   |  dS )Nr(   )r   )er   r   r   r   s  r1   z/test_agg_lambda_with_timezone.<locals>.<lambda>r]   r   )r   r~   r   r    r$   r   ra   rb   )r%   rn   ro   r   r   r   test_agg_lambda_with_timezoneh  s    	r   err_clsc           	         s   ddl m}m}m} | d d }tg dg d||d}t||d |d g} fdd	}|d
 |d |}tj	||dd d S )Nr   )DecimalArray	make_data
to_decimalr   )r   r   r   r(   r(   )r   r(   r   r(   r(   )id1Zid2decimalsr+   c                    s   t | dkr | jd S r   )r>   r   r   r   r   r   
weird_func  s    z1test_groupby_agg_err_catching.<locals>.weird_funcr   r   F)Zcheck_names)
Z$pandas.tests.extension.decimal.arrayr   r   r   r   r   r    r$   ra   rc   )	r   r   r   r   r?   r%   ro   r   rn   r   r   r   test_groupby_agg_err_catching|  s    r   )=__doc__r8   r7   	functoolsr   numpyr   r!   Zpandas.errorsr   Zpandasr~   r   r   r   r   r   r	   r
   Zpandas._testingZ_testingra   Zpandas.io.formats.printingr   r'   r@   rO   rg   rp   rt   r   r   r   r   r   r   r   r   r   r   markZparametrizer   r   r   r   rM   r   r   r   r   r   NotImplementedErrorRuntimeErrorr   
IndexErrorOSError
ValueErrorArithmeticErrorAttributeErrorr   r   r   r   r   <module>   s   $	%&/9
&

 $	
