o
    թZhR>                     @   s  d dl Zd dlZd dlmZmZmZmZmZ d dl	m
Z d dlmZmZ d dlmZmZmZ d dlmZ dd Zdd	 Zd
d Zdd Zejddedddddddddejg
i fdejdddddddddejg
i fdejdddddddddejg
i fdejddddddddd ejg
d!d"ifd#ej dddd$d%d&ddd'ejg
d!d"ifd(ej!ddddddddd)ejg
i fgd*d+ Z"d,d- Z#ejd.d/dddd0dd1ddejejg
fd2dddd3d4d5ddejejg
fgd6d7 Z$ejd8d9g d:gd;g d<ggd=d> Z%d?d@ Z&dAdB Z'dCdD Z(ejdEd"g dFfdGg dHfgejdId(g fdJd'gfgdKdL Z)ejdMeeegejdNg dOejdPd"d"gd d"gdQd"dRgd d"gdQd"gdS ejd"dRejge*e+dTdS dQgdUdV Z,ejdWd"d"d gd"gfd"dRd d"gd"dRgfdRd"d gd"gfdRdRd d"gdRdRgfdXdYe+dYe*e+dXdYdYgdX  fdYdXe+dXdXgdX fd d e-g e-g fd"d e-g e-g fd d"d gd gfg	dZd[ Z.ejd\eg d]g d^d_dReg d`ej/g d^e+dagdbdgdcddej0defeejd"dRejge*e+dTdf d"gdg dRgdh  e+dfdidYeg djej/d"gdg dRgdh  e+dfgdbdgdcddej0defgdkdl Z1ejdmg dnejdNg dodpdq Z2drds Z3dtdu Z4dS )v    N)	DataFrame
MultiIndexSeriesconcat
date_range)BaseIndexerFixedForwardWindowIndexer)ExpandingIndexerFixedWindowIndexerVariableOffsetWindowIndexer)BusinessDayc                  C   s\   G dd dt } |  }tjtdd ttd| W d    d S 1 s'w   Y  d S )Nc                   @      e Zd Zdd ZdS )z8test_bad_get_window_bounds_signature.<locals>.BadIndexerc                 S   s   d S N )selfr   r   \/var/www/html/lang_env/lib/python3.10/site-packages/pandas/tests/window/test_base_indexer.pyget_window_bounds   s   zJtest_bad_get_window_bounds_signature.<locals>.BadIndexer.get_window_boundsN__name__
__module____qualname__r   r   r   r   r   
BadIndexer       r   zBadIndexer does not implementmatch   )r   pytestraises
ValueErrorr   rangerolling)r   indexerr   r   r   $test_bad_get_window_bounds_signature   s
   "r"   c                  C   s<   t td} t }| | }|   }t|| d S )N
   )r   r   r	   r    meanZ	expandingtmassert_series_equal)sr!   resultexpectedr   r   r   test_expanding_indexer#   s
   r*   c                  C   sb   g d} t dtdi}G dd dt}|d| d}|| }t dg di}t|| d S )	NTFTFTvaluesr   c                   @   r   )z3test_indexer_constructor_arg.<locals>.CustomIndexerc           	      S   sl   t j|t jd}t j|t jd}t|D ]}| j| r&d||< |d ||< q|||< || j ||< q||fS Ndtyper      )npemptyint64r   use_expandingwindow_size	r   
num_valuesmin_periodscenterclosedstepstartendir   r   r   r   1   s   
zEtest_indexer_constructor_arg.<locals>.CustomIndexer.get_window_boundsNr   r   r   r   r   CustomIndexer0   r   r?   r0   r5   r4   )              ?      @rC         $@r   r   r   r    sumr%   assert_frame_equal)r4   dfr?   r!   r(   r)   r   r   r   test_indexer_constructor_arg+   s   rI   c                  C   sb   t dtdi} G dd dt}|dd}| j|ddddd	 }t dg d
i}t|| d S )Nr,   r   c                   @   r   )z8test_indexer_accepts_rolling_args.<locals>.CustomIndexerc           	      S   s   t j|t jd}t j|t jd}t|D ](}|r1|dkr1|dkr1|dkr1|dkr1d||< |||< q|||< || j ||< q||fS )Nr.   r0   both   r   )r1   r2   r3   r   r5   r6   r   r   r   r   G   s   
zJtest_indexer_accepts_rolling_args.<locals>.CustomIndexer.get_window_boundsNr   r   r   r   r   r?   F   r   r?   r0   r5   TrJ   )r9   r8   r:   r;   )rA   rB   rD   rC         @rE   )rH   r?   r!   r(   r)   r   r   r   !test_indexer_accepts_rolling_argsC   s   

rN   zfunc,np_func,expected,np_kwargscountrC          @minrA   rB   rM         @      @       @max      Y@      "@stdgK@gDmK@g!?J@g'e?Zddofr0   varg[{@@g[{@g[{Ħ@g      ?mediang      !@c                    st  t d}d|d< tdd}d}tjt|d | |j|dd	}	t|	|  W d    n1 s1w   Y  d
}tjt|d | |j|dd}	t|	|  W d    n1 sYw   Y  | |j|d|d}	t|	| }
| |d d | }t	|
| | |	
 fdd}t	|
| |dkrdnd }| |j||d}t|| }| |
 fdd}t	|| d S )NrD   rV   r      rL   z.Forward-looking windows can't have center=Truer   T)windowr9   zAForward-looking windows don't support setting the closed argumentright)r\   r:   rK   r\   r8   r;   c                        | fi S r   r   xnp_func	np_kwargsr   r   <lambda>       z-test_rolling_forward_window.<locals>.<lambda>rO   r   r\   r8   c                    r_   r   r   r`   rb   r   r   re      rf   )r1   aranger   r   r   r   r    getattrr%   assert_equalapply)frame_or_seriesfuncrc   r)   rd   r;   r,   r!   r   r    r(   Z	expected2r8   Zrolling3Zresult3Z	expected3r   rb   r   test_rolling_forward_windowa   s0   
9
rn   c                 C   sr   t d}d|d< tdd}| |j|d|d}| }| ddd	d
ddddt jt jg
d d | }t|| d S )NrD   rV   r   rL   r[   r^   rA   go@g\QJ@g@gv-@gA@)r1   rh   r   r    Zskewnanr%   rj   )rl   r;   r,   r!   r    r(   r)   r   r   r   test_rolling_forward_skewness   s*   

rp   zfunc,expectedcovg     @X@g     @Wcorrg?gh"?g1=c                 C   s   t ddd}|d }d|d< t j||gdd}tdd	}t|j|dd
}t||  jt	d dfdf }|j
dd}t|j
dd}|j|_t|| d S )Nr#   r0   rK   d   )r   r   )Zaxisr[   rL   rg   r   T)Zdrop)r1   rh   ZreshapeZconcatenater   r   r    ri   locsliceZreset_indexr   namer%   rj   )rm   r)   Zvalues1Zvalues2r,   r!   r    r(   r   r   r   test_rolling_forward_cov_corr   s   
rx   zclosed,expected_datar]   )
rA   rB   rP   rC   rS         (@rR   rS   rT   rW   left)
rA   rA   rB   rP         @rW   r{   rR   rS   rT   c                 C   s^   t ddd}ttd|d}td}t||d}|j|| d }t||d}t|| d S )N2020r#   Zperiods)indexr0   r~   offset)r:   )	r   r   r   r   r   r    rF   r%   rG   )r:   Zexpected_datar~   rH   r   r!   r(   r)   r   r   r   &test_non_fixed_variable_window_indexer   s   r   c                   C   sD   t jtdd tdtdd W d    d S 1 sw   Y  d S )Nzindex must be a DatetimeIndex.r   foor0   r   )r   r   r   r   r   r   r   r   r   (test_variableoffsetwindowindexer_not_dti  s   "r   c                  C   sL   t ddd} tjtdd t| dd W d    d S 1 sw   Y  d S )Nr|   r#   r}   z(offset must be a DateOffset-like object.r   r   r   )r   r   r   r   r   )idxr   r   r   +test_variableoffsetwindowindexer_not_offset  s   "r   c                 C   sX   t dg di}tdd}|j|d| d }t dg did d |  }t|| d S )Nb)NNN   rK   rL   r   r^   )rA   rA   rB   rB   )r   r   r    rO   r%   rG   )r;   rH   r!   r(   r)   r   r   r    test_fixed_forward_indexer_count  s
   
r   )	end_valuer,   )rA   r0   r0   r[   rK   rs   )rA   r0   r   r[   r0   )rm   argsZquantilec           
         sh   G  fdddt }g d}tdtdi}|d|d}t|||| }td|i}	t||	 d S )Nc                       s   e Zd Z fddZdS )z0test_indexer_quantile_sum.<locals>.CustomIndexerc           	         sr   t j|t jd}t j|t jd}t|D ] }| j| r)d||< t|  d||< q|||< || j ||< q||fS r-   )r1   r2   r3   r   r4   rU   r5   r6   r   r   r   r   ,  s   
zBtest_indexer_quantile_sum.<locals>.CustomIndexer.get_window_boundsNr   r   r   r   r   r?   +  s    r?   r+   r,   r   r0   r@   )r   r   r   ri   r    r%   rG   )
r   r,   rm   r   r?   r4   rH   r!   r(   r)   r   r   r   test_indexer_quantile_sum%  s   r   indexer_classr5   )r0   rK      df_data)ar   rK         c                 C   sX   t |}d}| |d}|j}t|D ]}|dd j|dd  |j|ks)J qd S )Nr[   rL   r   r   r0   rg   )r   r5   r   groupbyr    r$   )r   r5   r   rH   Z
num_trialsr!   Zoriginal_window_sizer>   r   r   r   0test_indexers_are_reusable_after_groupby_rollingA  s   
r   z5window_size, num_values, expected_start, expected_endr   r   c                 C   st   t | d}|j||d\}}tj|t|d d | dd tj|t|d d | dd t|t|ks8J d S )NrL   )r7   r;   F)Zcheck_dtype)r   r   r%   Zassert_numpy_array_equalr1   arraylen)r5   r7   Zexpected_startZexpected_endr;   r!   r<   r=   r   r   r   !test_fixed_forward_indexer_boundsZ  s   
 r   zdf, window_size, expected)r   r0   rK   )r0   rK   rK   )r   r   )r   g      ?rP   r[   r   )namesr   )r~   rw   r/      r      )r   r   c)@r   g      @r{   r{   g      @rR   ry   g      )@g      *@g      +@g      ,@g      -@g      .@g      /@g      0@g     0@g      1@c                 C   s6   t |d}| dd j|dd }t|| d S )NrL   r   r   r0   rg   )r   r   r    r$   r%   r&   )rH   r5   r)   r!   r(   r   r   r   0test_rolling_groupby_with_fixed_forward_specificv  s   
6r   
group_keys)	)r0   )r0   rK   )rK   r0   )r0   r0   rK   )r0   rK   r0   )r0   r0   rK   rK   )r0   rK   r[   rK   r[   )r0   r0   rK   r0   r0   rK   r0   r0   rK   r0   r0   rK   )r0   rK   r[   r0   rK   r[   r0   rK   r[   r0   rK   r[   r0   rK   r[   )r0   rK   r[   r   r         c                    s   t tt| tjt| tjdd tjt| tjdd}t d}|	dd j
|dd }dd	g|j_|	dg d }t fd
d|D }|dd	gd }t|| d S )Nr.      )r   r   r   rL   r   r   r0   rg   r   c                    s4   g | ]\}  j  fd dtt D dqS )c                    s*   g | ]} d  j ||  jddqS )r   r0   )Z	min_count)ZilocrF   ).0r>   )gr5   r   r   
<listcomp>  s    zKtest_rolling_groupby_with_fixed_forward_many.<locals>.<listcomp>.<listcomp>)r   )Zassignr   r   )r   _rL   )r   r   r     s    
z@test_rolling_groupby_with_fixed_forward_many.<locals>.<listcomp>)r   r1   r   listrh   r   float64r3   r   r   r    rF   r~   r   r   Z	set_indexr%   r&   )r   r5   rH   r!   r(   groupsZmanualr   rL   r   ,test_rolling_groupby_with_fixed_forward_many  s"   

r   c                  C   s  G dd dt } |  }td|}d}tjt|d |  W d    n1 s*w   Y  tjt|d tt| W d    n1 sGw   Y  tjt|d |j	dd W d    n1 sdw   Y  tjt|d |j
dd W d    d S 1 sw   Y  d S )Nc                   @   r   )z4test_unequal_start_end_bounds.<locals>.CustomIndexerc                 S   s   t dgt ddgfS Nr0   rK   r1   r   r   r7   r8   r9   r:   r;   r   r   r   r     s   zFtest_unequal_start_end_bounds.<locals>.CustomIndexer.get_window_boundsNr   r   r   r   r   r?     r   r?   r0   r<   r   Tpairwiser   r   r    r   r   r   r$   nextiterrr   rq   r?   r!   Zrollr   r   r   r   test_unequal_start_end_bounds  s    
"r   c                  C   s  G dd dt } |  }tddg|}d}tjt|d |  W d    n1 s,w   Y  tjt|d tt| W d    n1 sIw   Y  tjt|d |j	dd W d    n1 sfw   Y  tjt|d |j
dd W d    d S 1 sw   Y  d S )Nc                   @   r   )z4test_unequal_bounds_to_object.<locals>.CustomIndexerc                 S   s   t dgt dgfS r   r   r   r   r   r   r     s   zFtest_unequal_bounds_to_object.<locals>.CustomIndexer.get_window_boundsNr   r   r   r   r   r?     r   r?   r0   zstart and endr   Tr   r   r   r   r   r   test_unequal_bounds_to_object  s    
"r   )5numpyr1   r   Zpandasr   r   r   r   r   Zpandas._testingZ_testingr%   Zpandas.api.indexersr   r   Zpandas.core.indexers.objectsr	   r
   r   Zpandas.tseries.offsetsr   r"   r*   rI   rN   markZparametrizer   ro   rQ   rU   rX   rY   rZ   rn   rp   rx   r   r   r   r   r   r   r   r   r   r   Zfrom_arraysr   r   r   r   r   r   r   r   r   <module>   s:   
 "
5*




	
("
 
4