o
    i=                     @   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eeg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ej"d+d,d- Z#ejdeegd.d/ Z$ejd0d1dddd2dd3ddejejg
fd4dddd5d6d7ddejejg
fgd8d9 Z%ejd:d;g d<gd=g d>ggd?d@ Z&dAdB Z'ejdCd#g dDfdEg dFfgejdGd)g fdHd(gfgdIdJ Z(ejdKeeegejdLg dMejdNd#d#gd d#gdOd#dPgd d#gdOd#gdQ ejd#dPejge)e*dRdQ dOgdSdT Z+ejdUd#d#d gd#gfd#dPd d#gd#dPgfdPd#d gd#gfdPdPd d#gdPdPgfdVdWe*dWe)e*dVdWdWgdV  fdWdVe*dVdVgdV fd d e,g e,g fd#d e,g e,g fd d#d gd gfg	dXdY Z-ejdZeg d[g d\d]dPeg d^ej.g d\e*d_gd`dgdadbej/dcfeejd#dPejge)e*dRdd d#gde dPgdf  e*dddgdWeg dhej.d#gde dPgdf  e*ddgd`dgdadbej/dcfgdidj Z0ejdkg dlejdLg dmdndo Z1dpdq Z2drds Z3dS )t    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/edux/Edux_v2/venv/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    mean	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@   r1   r6   r5   )              ?      @rD         $@r   r   r   r    sumr&   assert_frame_equal)r5   dfr@   r!   r)   r*   r   r   r   test_indexer_constructor_arg+   s   rJ   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/   r1   both   r   )r2   r3   r4   r   r6   r7   r   r   r   r   G   s   
zJtest_indexer_accepts_rolling_args.<locals>.CustomIndexer.get_window_boundsNr   r   r   r   r   r@   F   r   r@   r1   r6   TrK   )r:   r9   r;   r<   )rB   rC   rE   rD         @rF   )rI   r@   r!   r)   r*   r   r   r   !test_indexer_accepts_rolling_argsC   s   

rO   constructorzfunc,np_func,expected,np_kwargscountrD          @minrB   rC   rN         @      @       @max      Y@      "@stdgK@gDmK@g!?J@g'e?ddofr1   varg[{@@g[{@g[{Ħ@g      ?mediang      !@z ignore:min_periods:FutureWarningc                    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 )NrE   rX   r      rM   z.Forward-looking windows can't have center=Truer   T)windowr:   zAForward-looking windows don't support setting the closed argumentright)r_   r;   rL   r_   r9   r<   c                        | fi S r   r   xnp_func	np_kwargsr   r   <lambda>       z-test_rolling_forward_window.<locals>.<lambda>rQ   r   r_   r9   c                    rb   r   r   rc   re   r   r   rh      ri   )r2   aranger   r   r   r   r    getattrr&   assert_equalapply)rP   funcrf   r*   rg   r<   r-   r!   r   r    r)   	expected2r9   rolling3result3	expected3r   re   r   test_rolling_forward_windowa   s0   
9
rt   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 )NrE   rX   r   rM   r^   ra   rB   go@g\QJ@g@gv-@gA@)r2   rk   r   r    skewnanr&   rm   )rP   r<   r-   r!   r    r)   r*   r   r   r   test_rolling_forward_skewness   s*   

rw   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#   r1   rL   d   )r   r   )axisr^   rM   rj   r   T)drop)r2   rk   reshapeconcatenater   r   r    rl   locslicereset_indexr   namer&   rm   )ro   r*   values1values2r-   r!   r    r)   r   r   r   test_rolling_forward_cov_corr   s   
r   zclosed,expected_datar`   )
rB   rC   rR   rD   rU         (@rT   rU   rV   rY   left)
rB   rB   rC   rR         @rY   r   rT   rU   rV   c                 C   s^   t ddd}ttd|d}td}t||d}|j|| d }t||d}t|| d S )N2020r#   )periods)indexr1   )r   offset)r;   )	r   r   r   r   r   r    rG   r&   rH   )r;   expected_datar   rI   r   r!   r)   r*   r   r   r   &test_non_fixed_variable_window_indexer   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   rL   rM   r   ra   )rB   rB   rC   rC   )r   r   r    rQ   r&   rH   )r<   rI   r!   r)   r*   r   r   r    test_fixed_forward_indexer_count  s
   
r   )	end_valuer-   )rB   r1   r1   r^   rL   rz   )rB   r1   r   r^   r1   )ro   args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.   )r2   r3   r4   r   r5   rW   r6   r7   r   r   r   r      s   
zBtest_indexer_quantile_sum.<locals>.CustomIndexer.get_window_boundsNr   r   r   r   r   r@     s    r@   r,   r-   r   r1   rA   )r   r   r   rl   r    r&   rH   )
r   r-   ro   r   r@   r5   rI   r!   r)   r*   r   r   r   test_indexer_quantile_sum  s   r   indexer_classr6   )r1   rL      df_data)ar   rL         c                 C   sX   t |}d}| |d}|j}t|D ]}|dd j|dd  |j|ks)J qd S )Nr^   rM   r   r   r1   rj   )r   r6   r   groupbyr    r$   )r   r6   r   rI   
num_trialsr!   original_window_sizer?   r   r   r   0test_indexers_are_reusable_after_groupby_rolling5  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 )NrM   )r8   r<   F)check_dtype)r   r   r&   assert_numpy_array_equalr2   arraylen)r6   r8   expected_startexpected_endr<   r!   r=   r>   r   r   r   !test_fixed_forward_indexer_boundsN  s   
 r   zdf, window_size, expected)r   r1   rL   )r1   rL   rL   )r   r   )r   g      ?rR   r^   r   )namesr   )r   r   r0      r      )r   r   c)@r   g      @r   r   g      @rT   r   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 )NrM   r   r   r1   rj   )r   r   r    r$   r&   r'   )rI   r6   r*   r!   r)   r   r   r   0test_rolling_groupby_with_fixed_forward_specificj  s   
6r   
group_keys)	)r1   )r1   rL   )rL   r1   )r1   r1   rL   )r1   rL   r1   )r1   r1   rL   rL   )r1   rL   r^   rL   r^   )r1   r1   rL   r1   r1   rL   r1   r1   rL   r1   r1   rL   )r1   rL   r^   r1   rL   r^   r1   rL   r^   r1   rL   r^   r1   rL   r^   )r1   rL   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   rM   r   r   r1   rj   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   r1   )	min_count)ilocrG   ).0r?   )gr6   r   r   
<listcomp>  s    zKtest_rolling_groupby_with_fixed_forward_many.<locals>.<listcomp>.<listcomp>)r   )assignr   r   )r   _rM   )r   r   r     s    
z@test_rolling_groupby_with_fixed_forward_many.<locals>.<listcomp>)r   r2   r   listrk   r   float64r4   r   r   r    rG   r   r   r   	set_indexr&   r'   )r   r6   rI   r!   r)   groupsmanualr   rM   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 Nr1   rL   r2   r   r   r8   r9   r:   r;   r<   r   r   r   r     s   zFtest_unequal_start_end_bounds.<locals>.CustomIndexer.get_window_boundsNr   r   r   r   r   r@     r   r@   r1   r=   r   Tpairwiser   r   r    r   r   r   r$   nextiterry   rx   r@   r!   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@   r1   zstart and endr   Tr   r   r   r   r   r   test_unequal_bounds_to_object  s    
"r   )4numpyr2   r   pandasr   r   r   r   r   pandas._testing_testingr&   pandas.api.indexersr   r   pandas.core.indexers.objectsr	   r
   r   pandas.tseries.offsetsr   r"   r+   rJ   rO   markparametrizer   rv   rS   rW   rZ   r\   r]   filterwarningsrt   rw   r   r   r   r   r   r   r   r   r   from_arraysr   r   r   r   r   r   r   r   r   <module>   s<   
 "
5(





	
("
 
4