a
    oÝEbi:  ã                   @   s   d dl Zd dlZd dlmZmZmZ d dlmZ d dl	m
  mZ d dlmZmZmZmZ d dlmZ d dlmZ ejjeƒ oˆeƒ pˆeƒ dd�Ze d¡ej d	¡G d
d„ dƒƒƒZe d¡G dd„ dƒƒZe d¡dd„ ƒZe d¡dd„ ƒZe d¡ejj ej d	¡G dd„ dƒƒƒƒZ!dS )é    N)Úis_ci_environmentÚis_platform_macÚis_platform_windows)ÚNumbaUtilError)Ú	DataFrameÚSeriesÚoption_contextÚto_datetime)ÚNUMBA_FUNC_CACHEzbOn Azure CI, Windows can fail with 'Windows fatal exception: stack overflow' and MacOS can timeout)ÚreasonÚnumbazignore:
c                   @   sÈ   e Zd Zej dddg¡dd„ ƒZej dee 	d¡ƒe
edƒdd	�g¡d
d„ ƒZej dee 	d¡ƒe
edƒdd	�g¡dd„ ƒZej dddg¡dd„ ƒZej dddddœgdi gg¡dd„ ƒZdS )Ú
TestEngineÚjitTFc                 C   s„   dd„ }|rdd l }| |¡}|||dœ}d}	ttdƒƒ}
|
jd|d�j||	d	|d
d�}|
jd|d�j|d|	d
d�}t ||¡ d S )Nc                 W   s$   d}|D ]}||7 }qt  | ¡| S )Nr   ©ÚnpÚmean)ÚxÚargsZarg_sumÚarg© r   ú^/home/ja/django-apps/lartica_env/lib/python3.9/site-packages/pandas/tests/window/test_numba.pyÚf$   s    
z0TestEngine.test_numba_vs_cython_apply.<locals>.fr   ©ÚnogilÚparallelÚnopython©é   é
   r   )Úcenterr   T)r   ÚengineÚengine_kwargsÚrawÚcython)r    r   r"   )r   r   r   ÚrangeÚrollingÚapplyÚtmÚassert_series_equal)Úselfr   r   r   r   r   r   r   r!   r   ÚsÚresultÚexpectedr   r   r   Útest_numba_vs_cython_apply"   s    

ÿÿz%TestEngine.test_numba_vs_cython_applyÚdataé   Úfoo)Únamec                 C   s„   |\}}|||dœ}|  d¡}	t|	|ƒf d|dœ|¤Ž}
t|	|ƒf ddi|¤Ž}|dvrtttd|› �ƒd	ftv stJ ‚t |
|¡ d S )
Nr   r   r   ©r    r!   r    r#   ©r   ÚsumÚvarÚstdÚmaxÚminÚnanÚRolling_apply_single)r%   Úgetattrr   r
   r'   Úassert_equal)r)   r.   r   r   r   Ú$arithmetic_numba_supported_operatorsÚmethodÚkwargsr!   Úrollr+   r,   r   r   r   Ú$test_numba_vs_cython_rolling_methods;   s     

ÿÿþýz/TestEngine.test_numba_vs_cython_rolling_methodsc                 C   s�   |\}}|||dœ}t t d¡ƒ}| ¡ }	t|	|ƒf d|dœ|¤Ž}
t|	|ƒf ddi|¤Ž}|dvr€ttd|› �ƒd	ftv s€J ‚t |
|¡ d S )
Nr   r/   r   r2   r    r#   r3   r9   ZExpanding_apply_single)r   r   ÚeyeÚ	expandingr;   r
   r'   r<   )r)   r.   r   r   r   r=   r>   r?   r!   Úexpandr+   r,   r   r   r   Ú&test_numba_vs_cython_expanding_methodsU   s"    
ÿÿþýz1TestEngine.test_numba_vs_cython_expanding_methodsc                 C   sì   dd„ }dd„ }|r0dd l }| |¡}| |¡}|||dœ}ttdƒƒ d¡}	|	j|d	|d
d�}
|	j|dd
d�}t |
|¡ |dftv sŒJ ‚|	j|d	|d
d�}
|	j|dd
d�}t |
|¡ |	j|d	|d
d�}
|	j|dd
d�}t |
|¡ d S )Nc                 S   s   t  | ¡d S )Né   r   ©r   r   r   r   Úfunc_1s   s    z+TestEngine.test_cache_apply.<locals>.func_1c                 S   s   t  | ¡d S )Nr/   )r   r6   rG   r   r   r   Úfunc_2v   s    z+TestEngine.test_cache_apply.<locals>.func_2r   r   r   r   r   T)r    r!   r"   r#   ©r    r"   r:   )	r   r   r   r$   r%   r&   r'   r(   r
   )r)   r   r   r   r   rH   rI   r   r!   r@   r+   r,   r   r   r   Útest_cache_applyp   s0    

ÿÿÿzTestEngine.test_cache_applyzwindow,window_kwargsr%   é   r   )ÚwindowÚmin_periodsrC   c                 C   s°   dd„ }|||dœ}t dg d¢iƒ}	t|	|ƒf d|i|¤Žj|dd|d	d
�}
t dg d¢iƒ}t |
|¡ t|	|ƒf d|i|¤Žj|dd|dd
�}
t dg d¢iƒ}t |
|¡ d S )Nc                 S   s   t  | ¡| S ©N©r   r4   )Úvaluesr   r   r   r   Úadd£   s    z,TestEngine.test_dont_cache_args.<locals>.add)r   r   r   Úvalue)r   r   r   r>   Tr   )é   )r"   r    r!   r   )ç      ð?rU   rU   r   )ç       @rV   rV   )r   r;   r&   r'   Úassert_frame_equal)r)   rM   Zwindow_kwargsr   r   r   r>   rR   r!   Údfr+   r,   r   r   r   Útest_dont_cache_args—   s    
ÿ
ÿzTestEngine.test_dont_cache_argsN)Ú__name__Ú
__module__Ú__qualname__ÚpytestÚmarkÚparametrizer-   r   r   rB   r   r$   rA   rE   rK   rY   r   r   r   r   r      s&   
ÿ
ÿ

&þþr   c                   @   sÊ   e Zd Zejjddd„ dd„ gddgd�ej dd	d
g¡dd„ ƒƒZejjddd„ dd„ gddgd�ej dd	d
g¡dd„ ƒƒZej dddg¡ej dd	d
g¡dd„ ƒƒZej dddg¡dd„ ƒZ	dS )ÚTestEWMÚgrouperc                 C   s   | S rO   r   rG   r   r   r   Ú<lambda>¸   ó    zTestEWM.<lambda>c                 C   s
   |   d¡S ©NÚA©ÚgroupbyrG   r   r   r   rb   ¸   rc   ÚNonerg   )Úidsr>   r   r4   c                 C   sd   t g d¢tdƒdœƒ}tjtdd��, t||ƒjdd�|ƒdd	� W d   ƒ n1 sV0    Y  d S )
N©ÚaÚbrk   rl   rF   ©re   ÚBzengine must be either©ÚmatchrU   ©Úcomr0   ©r    ©r   r$   r]   ÚraisesÚ
ValueErrorr;   Úewm©r)   ra   r>   rX   r   r   r   Útest_invalid_engine·   s    zTestEWM.test_invalid_enginec                 C   s   | S rO   r   rG   r   r   r   rb   Á   rc   c                 C   s
   |   d¡S rd   rf   rG   r   r   r   rb   Á   rc   c                 C   sj   t g d¢tdƒdœƒ}tjtdd��2 t||ƒjdd�|ƒdd	d
id� W d   ƒ n1 s\0    Y  d S )Nrj   rF   rm   zcython engine does notro   rU   rq   r#   r   Tr2   rt   rx   r   r   r   Útest_invalid_engine_kwargsÀ   s
    ÿz"TestEWM.test_invalid_engine_kwargsc                 C   sÄ   |dkrdd„ }t }ndd„ }d }|dkr.d}tg d¢tdƒd	œƒ}	||	ƒjd
||d�}
|||dœ}tj|dd��2 t|
|ƒd|d�}t|
|ƒdd�}W d   ƒ n1 sª0    Y  t ||¡ d S )Nrh   c                 S   s   | S rO   r   rG   r   r   r   rb   Ñ   rc   z.TestEWM.test_cython_vs_numba.<locals>.<lambda>c                 S   s
   |   d¡S rd   rf   rG   r   r   r   rb   Ô   rc   r4   Trj   rF   rm   rU   )rr   ÚadjustÚ	ignore_nar   Únuisancero   r   r2   r#   rs   )ÚFutureWarningr   r$   rw   r'   Úassert_produces_warningr;   rW   )r)   ra   r>   r   r   r   r|   r{   ÚwarnrX   rw   r!   r+   r,   r   r   r   Útest_cython_vs_numbaË   s    .zTestEWM.test_cython_vs_numbac                 C   sÂ   |dkrdd„ }t }ndd„ }d }d}tg d¢ƒ}tg d¢g d¢d	œƒ}	||	ƒj|d
||d�}
|||dœ}tj|dd��* |
jd|d�}|
jdd�}W d   ƒ n1 s¨0    Y  t ||¡ d S )Nrh   c                 S   s   | S rO   r   rG   r   r   r   rb   è   rc   z4TestEWM.test_cython_vs_numba_times.<locals>.<lambda>c                 S   s
   |   d¡S rd   rf   rG   r   r   r   rb   ë   rc   z23 days)ú
2020-01-01r‚   z
2020-01-02z
2020-01-10z
2020-02-23z
2020-01-03)rk   rl   rk   rl   rl   rk   )r   r   rT   rT   r   r   rm   T)Úhalflifer{   r|   Útimesr   r}   ro   r   r2   r#   rs   )r~   r	   r   rw   r'   r   r   rW   )r)   ra   r   r   r   r|   r€   rƒ   r„   rX   rw   r!   r+   r,   r   r   r   Útest_cython_vs_numba_timesã   s$    ÿ
ÿ*z"TestEWM.test_cython_vs_numba_timesN)
rZ   r[   r\   r]   r^   r_   ry   rz   r�   r…   r   r   r   r   r`   µ   s   ÿÿr`   c                  C   sz   dd„ } t tdƒƒ}tddƒ�& | d¡j| d dd�}W d   ƒ n1 sJ0    Y  | d¡j| ddd�}t ||¡ d S )	Nc                 S   s   t  | ¡d S )Nr   r   rG   r   r   r   r   
  s    z!test_use_global_config.<locals>.fr   zcompute.use_numbaTr   rJ   r   )r   r$   r   r%   r&   r'   r(   )r   r*   r+   r,   r   r   r   Útest_use_global_config  s    4r†   c                   C   sZ   t jtdd��8 ttdƒƒ d¡jdd„ ddiddd	� W d   ƒ n1 sL0    Y  d S )
Nz"numba does not support kwargs withro   rT   c                 S   s   | S rO   r   rG   r   r   r   rb     rc   z.test_invalid_kwargs_nopython.<locals>.<lambda>rk   r   T)r?   r    r"   )r]   ru   r   r   r$   r%   r&   r   r   r   r   Útest_invalid_kwargs_nopython  s    ÿr‡   c                   @   s|   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e d¡e d¡e d¡g¡e	j
 dddg¡dd„ ƒƒZdS )ÚTestTableMethodc                 C   s\   dd„ }t jtdd��2 ttdƒƒjddd�j|dd	d
� W d   ƒ n1 sN0    Y  d S )Nc                 S   s   t j| dd�d S ©Nr   ©ÚaxisrT   rP   rG   r   r   r   r   "  s    z7TestTableMethod.test_table_series_valueerror.<locals>.fz1method='table' not applicable for Series objects.ro   rT   Útable)r>   r   TrJ   )r]   ru   rv   r   r$   r%   r&   )r)   r   r   r   r   Útest_table_series_valueerror!  s    ÿÿz,TestTableMethod.test_table_series_valueerrorc                 C   sâ   |\}}|||dœ}t t d¡ƒ}	|	jdd|dd�}
|dv rŒtjt|› d�d	��* t|
|ƒf |d
dœ|¤Ž W d   ƒ qÞ1 s€0    Y  nR|	jdd|dd�}t|
|ƒf |d
dœ|¤Ž}t||ƒf |d
dœ|¤Ž}t 	||¡ d S )Nr   rL   r   rŒ   r   ©r>   r‹   rN   ©r5   r6   ú not supportedro   r   ©r!   r    Úsingle)
r   r   rB   r%   r]   ru   ÚNotImplementedErrorr;   r'   rW   )r)   r‹   r   r   r   r=   r>   r?   r!   rX   Z
roll_tableZroll_singler+   r,   r   r   r   Ú!test_table_method_rolling_methods,  s.    
ÿÿ&
ÿÿ
ÿÿz1TestTableMethod.test_table_method_rolling_methodsc           
      C   sr   |||dœ}dd„ }t t d¡ƒ}|jdd|dd�j|d	|d
d�}|jdd|dd�j|d	|d
d�}	t ||	¡ d S )Nr   c                 S   s   t j| dd�d S r‰   rP   rG   r   r   r   r   G  s    z:TestTableMethod.test_table_method_rolling_apply.<locals>.frL   r   rŒ   r   rŽ   Tr   ©r"   r!   r    r’   )r   r   rB   r%   r&   r'   rW   ©
r)   r‹   r   r   r   r!   r   rX   r+   r,   r   r   r   Útest_table_method_rolling_applyD  s    ÿÿz/TestTableMethod.test_table_method_rolling_applyc                 C   st   dd„ }t g d¢g d¢g d¢g d¢gƒ}|jddd	d
�j|ddd�}t g d¢g d¢g d¢g d¢gƒ}t ||¡ d S )Nc                 S   sh   t  d| jd f¡}| d d …d d…f | d d …df  jdd�| d d …df  ¡  |d d …d d…f< |S )NrT   r   r   rŠ   )r   ÚonesÚshaper4   )r   Zarrr   r   r   Úweighted_meanT  s    PzNTestTableMethod.test_table_method_rolling_weighted_mean.<locals>.weighted_mean)rT   r   g333333ã?)r   rL   gš™™™™™Ù?)rL   rF   gš™™™™™É?)rF   r/   gffffffæ?r   rŒ   r   )r>   rN   Tr   )r"   r    )rU   rV   rU   )gÍÌÌÌÌÌü?rV   rU   )g©kí}ªª
@g©kí}ªª@rU   )g“‹1°Žãø?é   rU   )r   r%   r&   r'   rW   )r)   rš   rX   r+   r,   r   r   r   Ú'test_table_method_rolling_weighted_meanS  s     ÿüÿz7TestTableMethod.test_table_method_rolling_weighted_meanc           
      C   sj   |||dœ}dd„ }t t d¡ƒ}|jd|d�j|d|dd	�}|jd
|d�j|d|dd	�}	t ||	¡ d S )Nr   c                 S   s   t j| dd�d S r‰   rP   rG   r   r   r   r   j  s    z<TestTableMethod.test_table_method_expanding_apply.<locals>.frL   rŒ   ©r>   r‹   Tr   r•   r’   )r   r   rB   rC   r&   r'   rW   r–   r   r   r   Ú!test_table_method_expanding_applyg  s    ÿÿz1TestTableMethod.test_table_method_expanding_applyc                 C   sÚ   |\}}|||dœ}t t d¡ƒ}	|	jd|d�}
|dv rˆtjt|› d�d��* t|
|ƒf |dd	œ|¤Ž W d   ƒ qÖ1 s|0    Y  nN|	jd
|d�}t|
|ƒf |dd	œ|¤Ž}t||ƒf |dd	œ|¤Ž}t 	||¡ d S )Nr   rL   rŒ   r�   r�   r�   ro   r   r‘   r’   )
r   r   rB   rC   r]   ru   r“   r;   r'   rW   )r)   r‹   r   r   r   r=   r>   r?   r!   rX   Zexpand_tableZexpand_singler+   r,   r   r   r   Ú#test_table_method_expanding_methodsv  s.    
ÿÿ&
ÿÿ
ÿÿz3TestTableMethod.test_table_method_expanding_methodsr.   rL   )r   rL   )rL   r   r>   r   r4   c                 C   s`   |||dœ}t |ƒ}t|jdd|d�|ƒ|dd�}	t|jdd|d�|ƒ|dd�}
t |	|
¡ d S )Nr   rT   rŒ   )rr   r>   r‹   r   r‘   r’   )r   r;   rw   r'   rW   )r)   r.   r>   r‹   r   r   r   r!   rX   r+   r,   r   r   r   Útest_table_method_ewmŽ  s    ÿÿz%TestTableMethod.test_table_method_ewmN)rZ   r[   r\   r�   r”   r—   rœ   rž   rŸ   r]   r^   r_   r   rB   r˜   r    r   r   r   r   rˆ     s   $rˆ   )"Únumpyr   r]   Zpandas.compatr   r   r   Zpandas.errorsr   Zpandas.util._test_decoratorsÚutilZ_test_decoratorsÚtdZpandasr   r   r   r	   Zpandas._testingZ_testingr'   Zpandas.core.util.numba_r
   r^   ZskipifZ
pytestmarkZ
skip_if_noÚfilterwarningsr   r`   r†   r‡   Zslowrˆ   r   r   r   r   Ú<module>   s2   þ
 R


