o
    d                     @   s   d dl 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lmZ ejdd Zejdd ZG d	d
 d
ZG dd dZG dd dZdS )    N)	DataFrameIndex
MultiIndexSeries	Timestamp
date_rangeto_datetime)BaseIndexer)get_groupbyc                   C   s    t g dg dtg ddS )z0Frame for testing times argument in EWM groupby.)
abcr   r   r   r   r   r   r   )
r   r   r      r   r      r   r      )

2020-01-01r   r   
2020-01-02
2020-01-10z
2020-01-22
2020-01-03
2020-01-23r   
2020-01-04ABC)r   r    r   r   T/app/.heroku/python/lib/python3.10/site-packages/pandas/tests/window/test_groupby.pytimes_frame   s   r   c                   C   .   t dgd dgd  dgd  tddS 	Nr      r      r      (   r   r   r   nparanger   r   r   r   
roll_frame+      .r(   c                   @   sD  e Zd Zdd Zdd Zdd Zejdg dd	d
 Z	ejdddgdd Z
ejdg ddd Zejdddgddggdd Zejdddgdd Zejdddgdd Zejddddgdd ggfdddgddggfgd!d" Zd#d$ Zd%d& Zejd'dd(gd)d*ggd+d, Zd-d. Zd/d0 Zejd1g d2d3d4 Zd5d6 Zd7d8 Zd9d: Zd;d< Zejd=d>d?gd@dA ZdBdC ZdDdE ZdFdG ZdHdI ZejdJdKdLddMfdNi fgdOdP ZejdQdRdSidSfdRdTidUfgdVdW Z ejdXd(d*gdYdZ Z!d[d\ Z"d]d^ Z#d_d` Z$dadb Z%dcdd Z&ejdedLdfdgg dhfdLdLdgg difdjdjdgg dkfdjdjdlg dmfgdndo Z'ejdpe()dqdrgdsdtggdudv Z*dwdx Z+ejdydzgd{g d|igdzdUge,d}e,d~e,d}e,d~gg ddggdd Z-dd Z.dd Z/dd Z0dd Z1dd Z2dTS )TestRollingc                 C   sF   d}t jt|d |jddd W d    d S 1 sw   Y  d S )Nz4groupby\(\) got an unexpected keyword argument 'foo'matchr   r   )foo)pytestraises	TypeErrorgroupby)selfr(   msgr   r   r   !test_groupby_unsupported_argument1   s   "z-TestRolling.test_groupby_unsupported_argumentc                 C   s   | d}t|dd}|jdd }|d j}t|| |dj }t|| |jd }t|| |j |jd }t|| d S )Nr   byc                 S      |  d S )Nr   rollingmeanxr   r   r   <lambda>:       z*TestRolling.test_getitem.<locals>.<lambda>r   )	r1   r
   r   applyr9   r:   tmassert_series_equalr   )r2   r(   g	g_mutatedexpectedresultr   r   r   test_getitem6   s   
zTestRolling.test_getitemc                 C   sd   | d}|jddd}t|dd}|jdd }|j }t|| |j }t|| d S )Nr   r   r   min_periodsr5   c                 S   s   | j ddd S )Nr   r   rG   )r9   countr;   r   r   r   r=   M       z3TestRolling.test_getitem_multiple.<locals>.<lambda>)r1   r9   r
   r   r?   rI   r@   rA   )r2   r(   rB   rrC   rD   rE   r   r   r   test_getitem_multipleH   s   


z!TestRolling.test_getitem_multiplefsumr:   minmaxrI   ZkurtZskewc                    sr   |j ddd}|jdd}t|  }| fdd}|jddd	}t|d td
g}||_t	
|| d S )Nr   F
group_keys   windowc                    s   t | d  S NrT   getattrr9   r;   rM   r   r   r=   f   rJ   z*TestRolling.test_rolling.<locals>.<lambda>r   Zaxisr#   r1   r9   rY   r?   dropr   from_arraysrangeindexr@   assert_frame_equalr2   rM   r(   rB   rK   rE   rD   expected_indexr   rZ   r   test_rollingU   s   zTestRolling.test_rollingstdvarc                    sv   |j ddd}|jdd}t| dd}| fdd	}|jddd
}t|d tdg}||_t	
|| d S )Nr   FrR   rT   rU   r   Zddofc                    s   t | d ddS )NrT   r   rg   rX   r;   rZ   r   r   r=   t       z/TestRolling.test_rolling_ddof.<locals>.<lambda>r[   r#   r\   rb   r   rZ   r   test_rolling_ddofn   s   zTestRolling.test_rolling_ddofinterpolationZlinearlowerZhigherZmidpointZnearestc                    st   |j ddd}|jdd}|jd d}| fdd	}|jdd
d}t|d tdg}||_t	
|| d S )Nr   FrR   rT   rU   皙?rj   c                    s   |  djd dS )NrT   rm   rn   )r9   quantiler;   rn   r   r   r=          z3TestRolling.test_rolling_quantile.<locals>.<lambda>r   r[   r#   )r1   r9   ro   r?   r]   r   r^   r_   r`   r@   ra   )r2   rj   r(   rB   rK   rE   rD   rc   r   rn   r   test_rolling_quantile|   s   
z!TestRolling.test_rolling_quantilezf, expected_valcorrr   cov      ?c              
   C   s   t tddgd dgd  g dd dddg}t tdg dd	d}t|jd
dd||}tjg|gd  d }t |dgtj	dgd dgd  dgd dgd  t
tddd gg ddd}t|| d S )N
   r      r   )r   r   r   rT   rv   )valueidx1idx2rx   ry   )rw   ry   r   levelrT   rw      )rx   rx   ry   namescolumnsr`   )r   r_   	set_indexrY   r1   r9   r&   nanr   r^   listr@   ra   )r2   rM   Zexpected_valdfotherrE   expected_datarD   r   r   r   /test_rolling_corr_cov_other_same_size_as_groups   s,   &
z;TestRolling.test_rolling_corr_cov_other_same_size_as_groupsc                    sV    d}|jdd}t| } fdd}||}tj|d< t|| d S )Nr   rT   rU   c                    s   t | d S rW   rX   r;   rM   r(   r   r   func   s   zITestRolling.test_rolling_corr_cov_other_diff_size_as_groups.<locals>.func)r1   r9   rY   r?   r&   r   r@   ra   r2   rM   r(   rB   rK   rE   r   rD   r   r   r   /test_rolling_corr_cov_other_diff_size_as_groups   s   


z;TestRolling.test_rolling_corr_cov_other_diff_size_as_groupsc                    sN   | d}|jdd}t|j dd} fdd}||}t|| d S )Nr   rT   rU   Tpairwisec                    s   t | jd ddS )NrT   Tr   )rY   r   r9   r;   rZ   r   r   r      s   z8TestRolling.test_rolling_corr_cov_pairwise.<locals>.func)r1   r9   rY   r   r?   r@   rA   r   r   rZ   r   test_rolling_corr_cov_pairwise   s   

z*TestRolling.test_rolling_corr_cov_pairwisezfunc, expected_values      ?      @c                 C   s   t g dg dg dd}|dd}t|| }t dtjg |d  d	tjg  dtjg |d
  d	tjg  dtjg dg ddd}t	|| d S )N)g1g2r   r   )r   r   r   r   )r   r   r|   rT   r   r   r   r   r   rT   r   r   r   )r   r   ))r   r   r   )r   r   r   )r   r   r   )r   r   r   )r   r   r   )r   r   r   )r   r   r   )r   r   r   )r   NNr}   r`   )
r   r1   r9   rY   r&   r   r   from_tuplesr@   ra   )r2   r   Zexpected_valuesr   ZrolrE   rD   r   r   r   test_rolling_corr_cov_unordered   s$   
z+TestRolling.test_rolling_corr_cov_unorderedc                    sx   |j ddd}|jdd}|jdd  d}| fd	d}|jdd
d}t|d tdg}||_t	|| d S )Nr   FrR   rT   rU   c                 S      |   S NrO   r;   r   r   r   r=          z0TestRolling.test_rolling_apply.<locals>.<lambda>rawc                    s   |  djdd  dS )NrT   c                 S   r   r   r   yr   r   r   r=      r   zBTestRolling.test_rolling_apply.<locals>.<lambda>.<locals>.<lambda>r   )r9   r?   r;   r   r   r   r=          r   r[   r#   )
r1   r9   r?   r]   r   r^   r_   r`   r@   ra   )r2   r   r(   rB   rK   rE   rD   rc   r   r   r   test_rolling_apply   s   zTestRolling.test_rolling_applyc                 C   s   t dgd dgd  dgd d}|d}tg d}dd g|_t tjd	d	gd
 dg|d}|jd
d }t	
|| |  |jd
d }t	
|| d S )Nr-   r   barr   r|   r$   r   ))r   r   )r   rT   )r   rv   )r-   r   )r-   r   )r-   r          @r   r   r   rU   )r   r1   r   r   r~   r&   r   r9   rO   r@   ra   )r2   r   rB   mirD   rE   r   r   r   test_rolling_apply_mutability   s   $

z)TestRolling.test_rolling_apply_mutabilityzexpected_value,raw_valueT        Fc                 C   sl   dd }t g dg dd}|djdj||d}t|gd	 tjd
dd gddd}t	|| d S )Nc                 S   s   t t| tjS r   )int
isinstancer&   Zndarrayr;   r   r   r   isnumpyarray  s   z6TestRolling.test_groupby_rolling.<locals>.isnumpyarrayr   r   r   r   r   r   )idrw   r   r   r   r   )r   r   r   r   r   r   r}   rw   r`   name)
r   r1   rw   r9   r?   r   r   r   r@   rA   )r2   Zexpected_valueZ	raw_valuer   r   rE   rD   r   r   r   test_groupby_rolling  s   z TestRolling.test_groupby_rollingc                 C   s  t tdd}||jddd }t tjgd tdd}t	
|| t tdd}||jddd }t tjgd	 td
d}t	
|| tdgd dgd  tdd}|djddd }ttjdddtjtjddddtjgtjddd gddgd}t	|| tdgd dgd  tdd}|djddd }ttjdddtjtjdddtjg
tjddd gddgd}t	|| d S )Nr   r|   Tr   )centerrV   rv   )r   r   r   r   r   rT   r   )rv   rT   r   rT   )r   r   r   r   r   r      r   r   r      r"   	   )r   r   r   r   r   r   r   r   r   rT   r   rv   r   r|   r   r   r   r"   r   r   )r   ru   r}   )r`   r   ru   )
r   r   r   r   r   r   r   r   r   r   )r   r_   r1   r9   r:   r&   r   r   r   r@   rA   r   ra   )r2   ZseriesrE   rD   r   r   r   r   "test_groupby_rolling_center_center  sH   

" "z.TestRolling.test_groupby_rolling_center_centerc                 C   s   t tdddgd dgd  tddd	}|d
jdddddj }tg ddtj	dt
dfdt
dfdt
dfdt
dfdt
dfdt
dfdt
dfdt
dfdt
dfdt
dff
d
dgdd}t|| d S )Nr   r   group_1r|   group_2rT   ru   )Dategbrw   )datar   r   Tr   )onr   rH   )
r         ?r   g      @      @g      @      @      @r   r   rw   r   r   r   z
2020-01-05z
2020-01-06z
2020-01-07z
2020-01-08z
2020-01-09r}   r   r`   )r   r   r_   r1   r9   rw   r:   r   r   r   r   r@   rA   r2   r   rE   rD   r   r   r   test_groupby_rolling_center_on^  s<   










z*TestRolling.test_groupby_rolling_center_onrH   )rv   rT   r   c                 C   s   t dgd dgd  tdd}d}|dj|d|d	 }| dd
g }g d}dd |D }td|d }tjg| }|||d|   | }	|||d|   | }
t dgd dgd  |	|
 d}t	
|| d S )Nr   ru   r   r    )groupr   rv   r   T)r   rH   r   )
r   r   r   r   r         @g      @r   r          @c                 S   s   g | ]}|d  qS )g      $@r   .0r<   r   r   r   
<listcomp>  rp   zGTestRolling.test_groupby_rolling_center_min_periods.<locals>.<listcomp>r   r   )r   r_   r1   r9   r:   Zreset_indexrQ   r&   r   r@   ra   )r2   rH   r   window_sizerE   Z
grp_A_meanZ
grp_B_meanZnum_nansZnansZgrp_A_expectedZgrp_B_expectedrD   r   r   r   'test_groupby_rolling_center_min_periods  s$   "z3TestRolling.test_groupby_rolling_center_min_periodsc                 C   s   t g dg dg dd}|ddg d }t tjtjdtjgdgtjd	dd gd
d}t	|| |dd d }t
tjtjdtjgtjd	dd gd
dd}t|| d S )N)r   r   r   r   )r   r   r   r   )ru   r       r    r   r   r   r   r   )r   r   r   r   r   r}   r   r   )r   r1   r9   rQ   r&   r   r   r   r@   ra   r   rA   r   r   r   r   test_groupby_subselect_rolling  s*   z*TestRolling.test_groupby_subselect_rollingc                 C   s   G dd dt }tdg dd idgd dgd  d	gd  d
}||jj|dddd }||jjddd }t|| d S )Nc                   @   s    e Zd Z					dddZdS )zFTestRolling.test_groupby_rolling_custom_indexer.<locals>.SimpleIndexerr   Nc                 S   sH   |d u r| j nd}tj|tjdd }| | j  }|||dk < ||fS )Nr   dtyper   )r   r&   r'   int64copy)r2   Z
num_valuesrH   r   closedstependstartr   r   r   get_window_bounds  s
   zXTestRolling.test_groupby_rolling_custom_indexer.<locals>.SimpleIndexer.get_window_bounds)r   NNNN)__name__
__module____qualname__r   r   r   r   r   SimpleIndexer  s    r   r   )r   r   r   r   r   r   r   rv   r   r   r   )r   rG   rV   rH   )r	   r   r1   r`   r9   rO   r@   ra   )r2   r   r   rE   rD   r   r   r   #test_groupby_rolling_custom_indexer  s   *
z/TestRolling.test_groupby_rolling_custom_indexerc                 C   s   t tdtddddg tdgd d}|djdd	d
dd  }ttjddtjddgt	j
dtdfgd dtdfgd  dd	gddd}t|| d S )Nr|   r   r   r   
2019-01-01column1Zcolumn2r   dater   1Dr   leftr   r   r   r   r   r   r   r}   r   r   r_   r   r1   r9   rO   r   r&   r   r   r   r@   rA   r   r   r   r   'test_groupby_rolling_subset_with_closed  s(   
		z3TestRolling.test_groupby_rolling_subset_with_closedc                 C   s   t tdtddddg tdgd d}|ddd	g jd
d	ddd  }ttjddtjddgt	j
dtdfgd dtdfgd  dd	gddd}t|| d S )Nr|   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r}   r   r   r   r   r   r   .test_groupby_subset_rolling_subset_with_closed  s2   

	z:TestRolling.test_groupby_subset_rolling_subset_with_closedr   rQ   rP   c                 C   sv   t g dtjg dddgddd}t||d| }t tjtjd	gtjg d
g dddd}t	|| d S )N)r   r   r   ))r   r<   )r   r   )r   z12r}   r   r   r   r   ))r   r   r<   )r   r   r   )r   r   r   )r   r   r   )
r   r   r   rY   r1   r9   r&   r   r@   rA   )r2   r   ZdsrE   rD   r   r   r   "test_groupby_rolling_index_changed  s    z.TestRolling.test_groupby_rolling_index_changedc                 C   s   t dg i}|djdd }|jdd}tjtg ddtg ddgdd gd|_t	
|| t g g d	}|dd
gjdd }|jdd
gd}tjtg ddtg ddtg ddgg dd|_t	
|| d S )Ns1r   rU   r   float64r   r   r}   )r   s2r   )r   r   N)r   r1   r9   rO   r]   r   Zfrom_productr   r`   r@   ra   )r2   rD   rE   r   r   r    test_groupby_rolling_empty_frame  s$   


z,TestRolling.test_groupby_rolling_empty_framec                 C   s0  t ddtddddgddtddddgdd	tddd
dgddtddddgdd	tddddggg ddd}|d}| |d< |jddd}|dd }t ddtdddddgddtdddddgddtdddddgdd	tddd
ddgdd	tdddddggg ddddg}t|| d S )Nr   r     r   r   r   r   Zr   r   Hr|   Er    )r`   r   	eventTimer   r`   r   count_to_dateZ10dr  r   c                 S   s
   | j d S )Nr   )shaper   r   r   r   r=   D  s   
 z?TestRolling.test_groupby_rolling_string_index.<locals>.<lambda>r   r   r   )r`   r   r  r  )	r   r   r   r1   Zcumcountr9   r?   r@   ra   )r2   r   groupsZrolling_groupsrE   rD   r   r   r   !test_groupby_rolling_string_index4  s6   	

	
z-TestRolling.test_groupby_rolling_string_indexc                 C   s|   t ddgddgdjdddd }t tddgddggdd	gtjd
dgdd gdd}|jdd}t	
|| d S )Nr   r   )r-   r   r-   F)sortr   r   r   )r   r   r   r}   r   r   )r   r1   r9   rP   r&   arrayr   r   r]   r@   ra   )r2   rE   rD   r   r   r   test_groupby_rolling_no_sortQ  s   z(TestRolling.test_groupby_rolling_no_sortc                 C   s   t tdtddddg tdddd}|djd	d
ddd  }ttjddtjddgdt	j
dtdfdtdfdtdfdtdfdtdfdtdfgdd
gdd}t|| d S )Nr|   r   r   r   Z20190101)r   periodsr   r   Z3dr   r   r   r   r   z
2018-12-27z
2018-12-29z
2018-12-31z
2018-12-28z
2018-12-30r   r}   r   )r   r_   r   r1   r9   rI   r   r&   r   r   r   r   r@   rA   r   r   r   r   $test_groupby_rolling_count_closed_onb  s:   

	





z0TestRolling.test_groupby_rolling_count_closed_on)r   kwargsr9   r   r   	expandingc                 C   s   t ddgddgddgddgddggddgd}t|d|di | }t tjgd tjdtjddgd	tjg d
dd gdd}|jdd}t	
|| d S )Nr   r   r   r   r   r   rv   g&?r   )r   r   )r   r   )r   r   r   rT   r}   r   r   )r   rY   r1   Zsemr&   r   r   r   r]   r@   ra   )r2   r   r  r   rE   rD   r   r   r   test_groupby_rolling_sem  s   &z$TestRolling.test_groupby_rolling_sem)rollingskeyr   r   Nr`   c                 C   s   t tddtjgg dg dd}|dkr|d}tjt| dd	 |d
j	di | W d    d S 1 s<w   Y  d S )Nz2020-06-01 12:00z2020-06-01 14:00r   r   r   r`   r   z values must not have NaTr+   r   60min)r  )
r   r   r&   r   r   r.   r/   
ValueErrorr1   r9   )r2   r  r  r   r   r   r   "test_groupby_rolling_nans_in_index  s   
"z.TestRolling.test_groupby_rolling_nans_in_indexrS   c                 C   s|   g dg dg}t j|dd}tg d|d}|jddg|dd	 }tg d
t jg dg ddd}t|| d S )Nval1r  val2rx   ry   r}   r   r   rx   ry   rR   r   )r   r   r   )r  r  r  r  r  )r  r  r  r  )rx   ry   rx   ry   )	r   r^   r   r1   r9   r:   r   r@   rA   )r2   rS   arraysr`   srE   rD   r   r   r   test_groupby_rolling_group_keys  s   z+TestRolling.test_groupby_rolling_group_keysc                 C   s   g dg dg}t j|dd}tg dtdd|d}|dd	gd
 }tdg dit jg dg ddd}t	|| d S )Nr  r  r}   )r   r   r   r   r$   r   rx   r   r   r   )r   r   r   )r  r   r  r  r   )r  r   r  r  )rx   r   rx   ry   )
r   r^   r   r_   r1   r9   r:   r   r@   ra   r2   r  r`   r   rE   rD   r   r   r   1test_groupby_rolling_index_level_and_column_label  s   
z=TestRolling.test_groupby_rolling_index_level_and_column_labelc                 C   sX   t tdddgd d}|dd }tjg ddd gd	}t	|j
| d S )
Nr   r   r   rT   r   r   r   )r   r   )r   rT   )r   r|   r   r   )r   rv   )r   r   r}   )r   r&   r'   r1   r9   r:   r   r   r@   assert_index_equalr`   r2   r   rE   rc   r   r   r   )test_groupby_rolling_resulting_multiindex  s   z5TestRolling.test_groupby_rolling_resulting_multiindexc                 C   sf   t tdddgd g dd d}|dd	gd }tjg d
g dd}t	|j
| d S )Ng      (@r   r   r|   r   r   r   rT   r   r   r   r   ))r   r   r   )r   r   rT   )r   r   r"   )r   r   r   )r   r   r|   )r   r   ru   )r   r   r   )r   r   rv   )r   r   r   )r   rT   r   )r   rT   r   )r   rT   r   )r   r   Nr}   )r   r&   r'   r1   r9   rO   r   r   r@   r#  r`   r$  r   r   r   *test_groupby_rolling_resulting_multiindex2  s   &z6TestRolling.test_groupby_rolling_resulting_multiindex2c                 C   st   t tdddgd g dd d}|jddd	}|d
d }tjg dg dd}t	j
|j|dd d S )Nr   r   r   rT   r&  r   r   T)appendr   r   ))r   r   r   r   )r   rT   r   )r   r|   r   )r   r   r   )r   r   rT   )r   rv   r   )r   r   rT   )r   Nr   r}   equiv)exact)r   r&   r'   r   r1   r9   r:   r   r   r@   r#  r`   r$  r   r   r   *test_groupby_rolling_resulting_multiindex3  s   &
z6TestRolling.test_groupby_rolling_resulting_multiindex3c                 C   sJ   |j ddd}|dd j}|jdd}|dd j}t|| d S )	Nr   FrR   c                 S   r7   rW   r9   rO   r;   r   r   r   r=     r>   zUTestRolling.test_groupby_rolling_object_doesnt_affect_groupby_apply.<locals>.<lambda>rT   rU   c                 S   r7   rW   r,  r;   r   r   r   r=     r>   )r1   r?   r`   r9   r@   r#  )r2   r(   rB   rD   _rE   r   r   r   7test_groupby_rolling_object_doesnt_affect_groupby_apply
  s
   zCTestRolling.test_groupby_rolling_object_doesnt_affect_groupby_apply)rV   rH   r   rD   r   r   )Nr   r   r   Nr   r   r   )NNr   r   NNr   r   rT   )NNNNNNNNright)NNNr   NNNr   c                 C   sx   t g d}|g dj|||dd}t tj|ddttddgg d	gg d
g dgdd}t|| d S )N)r   r   r   rT   rv   r|   r   r"   )r   r   r   r   r   r   r   r   )rV   rH   r   r   r   r   r   r   )r   r   r   r   rT   rv   r|   r   )r   r   r   r   r   r   r   r   )r   r   rT   r|   r   r   rv   r   )Zlevelscodesr   )	r   r1   r9   rf   r&   r  r   r@   ra   )r2   rV   rH   r   rD   r   rE   Zexpected_resultr   r   r   test_groupby_rolling_var  s   
z$TestRolling.test_groupby_rolling_varr   )r    )r   r   r   r   c                 C   s|   t ddggd ddggd  ddggd  |d}|d	}|jjd
d}|d}| }d	|jvs5J t|j| d S )Nr   r   r    r   r!   r   r"   r   r   T)deeprT   )	r   r1   objr   r9   rO   r   r@   ra   )r2   r   r   rB   Zoriginal_objrK   rE   r   r   r   test_by_column_not_in_values+  s   2

z(TestRolling.test_by_column_not_in_valuesc                 C   s   g dg dg}t j|dd}tdg di|d}|jdd	d d
 }ttjdtjdgt j	g dg dddd}t
|| d S )N)Falconr6  Parrotr7  )CaptiveWildr8  r9  )AnimalTyper}   z	Max Speed)g     `x@g     u@g      >@g      4@r   r   rz   r   g      @g      I@))r6  r6  r8  )r6  r6  r9  )r7  r7  r8  )r7  r7  r9  )r:  r:  r;  r   )r   r^   r   r1   r9   rO   r   r&   r   r   r@   rA   r!  r   r   r   test_groupby_level8  s   	zTestRolling.test_groupby_levelzby, expected_datar   num)      Y@     b@r?        i@
2018-01-01
2018-01-02)r>  r@  r?       @o@)r   r=  c                    s   g dg dg dg dg}t |g dd t d  d<  dg  fdd	|D } j|d
djddd }dg di}|| t | jd}t	|| d S )N)r   rA  r>  )r   rB  r@  )r   rA  r?  )r   rB  rC  )r   r   r=  r   r   c                    s   g | ]}t  |qS r   )rY   )r   attrr  r   r   r   n  rh   z3TestRolling.test_as_index_false.<locals>.<listcomp>F)Zas_indexr   r   r   r   )r   r   r   r   r   )
r   r   r   r1   r9   r:   updater`   r@   ra   )r2   r6   r   r   Zgp_byrE   rD   r   r  r   test_as_index_falseP  s$   
zTestRolling.test_as_index_falsec           
      C   s   t |j}d}t |d|}d|d< d}t ||}t j|d< d|d< t||dd}|dd jd	dd
	 }t
|dttdgd dg |ddtdgd dg |ddgd}	t||	 d S )N  r   r   g52DE)r`   adl2r`   rI  ru   r   i  )r   r   r   )r&   r   typerepeatr   r   r   r1   r9   r:   r   r   r^   r   r@   rA   )
r2   Zany_int_numpy_dtypetypsizeidxvalZarrr   rE   rD   r   r   r   test_nan_and_zero_endpoints{  s4   

z'TestRolling.test_nan_and_zero_endpointsc                    sj   g d}d t  fdd|D }tjtdd |djdd	d
 W d    d S 1 s.w   Y  d S )N)r   r   r   r   rG  c                    s(   g | ]}t d |   |d ddqS )r   r   *   )tr<   r   r   r   secr   r   r     s   ( zBTestRolling.test_groupby_rolling_non_monotonic.<locals>.<listcomp>z.* must be monotonicr+   r   rR  Z3s)r   rV   )r   r.   r/   r  r1   r9   )r2   Zshuffledr   r   rT  r   "test_groupby_rolling_non_monotonic  s   "z.TestRolling.test_groupby_rolling_non_monotonicc                 C   s   g dg dg dg dg dg dg dg dg d	g d
g
}t |g dd}t|d |d< |d}|dddd }|djdddd  }t	|| d S )N)Davidz1/1/2015d   )rW  z1/5/2015  )rW  z	5/30/20152   )rW  z	7/25/2015rZ  )Ryanz1/4/2014rX  )r[  z	1/19/2015rY  )r[  z	3/31/2016rZ  )Joez7/1/2015rX  )r\  z9/9/2015rY  )r\  z
10/15/2015rZ  )r   r   amount)r   r   r   r   c                 S      |  dd  S )N180Dr]  r,  r;   r   r   r   r=     rJ   z4TestRolling.test_groupby_monotonic.<locals>.<lambda>r_  r  r]  )
r   r   Zsort_valuesr   r1   r?   r9   rO   r@   rA   )r2   r   r   rD   rE   r   r   r   test_groupby_monotonic  s(   

z"TestRolling.test_groupby_monotonicc                 C   s   t dddd}tdgd dgd  dgd	  t||ftd
d}|dddd }|djdddj	
 }t|| d S )Nz2016-01-01 09:30:00r    r  )r   r  freqr   r   r!   r   r"   r#   r   r   r   c                 S   r^  )N4sr   r8   r;   r   r   r   r=     rJ   zJTestRolling.test_datelike_on_monotonic_within_each_group.<locals>.<lambda>rb  r  )r   r   r&   Zconcatenater'   r   r1   r?   r9   r   r:   r@   rA   )r2   datesr   rD   rE   r   r   r   ,test_datelike_on_monotonic_within_each_group  s   	z8TestRolling.test_datelike_on_monotonic_within_each_groupc                 C   sz   t dgd dgd  dd dD d tdd}tjtd	d
 |djddd W d    d S 1 s6w   Y  d S )Nr   r   r   c                 S   s   g | ]}t |d d qS )r   rS  )r   yearr   r   r   r     r   zPTestRolling.test_datelike_on_not_monotonic_within_each_group.<locals>.<listcomp>)i  i  r   r|   r   z&Each group within B must be monotonic.r+   r   Z365Dr   r  )r   r_   r.   r/   r  r1   r9   )r2   r   r   r   r   0test_datelike_on_not_monotonic_within_each_group  s   "z<TestRolling.test_datelike_on_not_monotonic_within_each_group)3r   r   r   r4   rF   rL   r.   markparametrizerd   ri   rq   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r	  r  r  r  r  r  r"  r%  r'  r+  r.  r1  r   r   r5  r<  r   rF  rP  rV  r`  rd  rf  r   r   r   r   r*   0   s    





&
 
A#

 



	


 r*   c                   @   s   e Zd Zejdd Zejdg ddd Zejdddgd	d
 Z	ejdg ddd Z
ejdddgdd Zdd ZdS )TestExpandingc                 C   r   r   r%   )r2   r   r   r   frame  r)   zTestExpanding.framerM   rN   c                    sn   |j ddd}| }t|  }| fdd}|jddd}t|d tdg}||_t	
|| d S )	Nr   FrR   c                    s   t |    S r   rY   r  r;   rZ   r   r   r=     s    z.TestExpanding.test_expanding.<locals>.<lambda>r   r[   r#   r1   r  rY   r?   r]   r   r^   r_   r`   r@   ra   r2   rM   rj  rB   rK   rE   rD   rc   r   rZ   r   test_expanding  s   zTestExpanding.test_expandingre   rf   c                    sr   |j ddd}| }t| dd}| fdd}|jddd	}t|d td
g}||_t	
|| d S )Nr   FrR   r   rg   c                    s   t |   ddS )Nr   rg   rk  r;   rZ   r   r   r=     rp   z3TestExpanding.test_expanding_ddof.<locals>.<lambda>r   r[   r#   rl  rm  r   rZ   r   test_expanding_ddof  s   z!TestExpanding.test_expanding_ddofrj   rk   c                    sp   |j ddd}| }|jd d}| fdd}|jddd	}t|d td
g}||_t	
|| d S )Nr   FrR   rm   rn   c                    s   |   jd dS )Nrm   rn   )r  ro   r;   rn   r   r   r=     rJ   z7TestExpanding.test_expanding_quantile.<locals>.<lambda>r   r[   r#   )r1   r  ro   r?   r]   r   r^   r_   r`   r@   ra   )r2   rj   rj  rB   rK   rE   rD   rc   r   rn   r   test_expanding_quantile
  s   
z%TestExpanding.test_expanding_quantilerr   rs   c           
         s    d}| }t| } fdd}||}ttddttdd }tj|j|df< tj|d< t	
|| t|j d	d
} fdd}	||	}t	|| d S )Nr   c                    s   t |   S r   rk  r;   rM   rj  r   r   func_0#  s   z5TestExpanding.test_expanding_corr_cov.<locals>.func_0r    =   H   q   r   Tr   c                    s   t | j  ddS )NTr   )rY   r   r  r;   rZ   r   r   func_12  s   z5TestExpanding.test_expanding_corr_cov.<locals>.func_1)r1   r  rY   r?   r   r_   r&   r   Zilocr@   ra   r   rA   )
r2   rM   rj  rB   rK   rE   rr  rD   Znull_idxrv  r   rq  r   test_expanding_corr_cov  s   



z%TestExpanding.test_expanding_corr_covc                    st   |j ddd}| }|jdd  d}| fdd}|jddd	}t|d td
g}||_t	|| d S )Nr   FrR   c                 S   r   r   r   r;   r   r   r   r=   =  r   z4TestExpanding.test_expanding_apply.<locals>.<lambda>r   c                    s   |   jdd  dS )Nc                 S   r   r   r   r   r   r   r   r=   >  r   zFTestExpanding.test_expanding_apply.<locals>.<lambda>.<locals>.<lambda>r   )r  r?   r;   r   r   r   r=   >  rh   r   r[   r#   )
r1   r  r?   r]   r   r^   r_   r`   r@   ra   )r2   r   rj  rB   rK   rE   rD   rc   r   r   r   test_expanding_apply8  s   z"TestExpanding.test_expanding_applyN)r   r   r   r.   fixturerj  rg  rh  rn  ro  rp  rw  rx  r   r   r   r   ri    s     




ri  c                
   @   s   e Zd Zejddg dgdejdddggdejd	d
dgggdd Zejddejdddggdejd	ddgggdd Z	dd Z
dd Zdd ZdS )TestEWMzmethod, expected_datar:   )r   gUUUUUU?gm۶m?g""""""@re   g۞?g?gt ?rf   rt   gn۶m۶?g++?c                 C   sd   t dgd tdd}t|djdd| }t d|itjg ddd gd	d
}t|| d S )Nr   rT   r$   r   r   comr   )r   r   r   r   r}   r   )	r   r_   rY   r1   ewmr   r   r@   ra   r2   methodr   r   rE   rD   r   r   r   test_methodsH  s   
zTestEWM.test_methodsrr   r   r   rs   grڶ?g߿yq+?c                    s   t dgd tdd}t|djdd  }t d|itjg dg d	d
d}t|| |d	 fdd}t|| d S )Nr   rT   r$   r   r   r{  r   ))r   r   r   )r   r   r   )r   r   r   )r   r   r   )r   NNr}   r   c                    s   t | jdd  S )Nr   r{  )rY   r}  r;   r  r   r   r=   x  rp   z/TestEWM.test_pairwise_methods.<locals>.<lambda>)
r   r_   rY   r1   r}  r   r   r@   ra   r?   r~  r   r  r   test_pairwise_methodsb  s   zTestEWM.test_pairwise_methodsc                 C   s\   d}| d}|dj||d }tdg ditjg ddd gdd	}t|| d S )
N23 daysr   r   halflifetimesr   )
r   gCus=?g;pGR?gZkB?r   g>(?g8Kr?r   gҧU?gA9w?)
r   r   )r   r|   )r   r   )r   r   r  r   )r   r   )r   rv   )r   r"   r}   r   )	popr1   r}  r:   r   r   r   r@   ra   )r2   r   r  r  rE   rD   r   r   r   
test_times{  s   
zTestEWM.test_timesc                 C   sN   d}| d}|d}|j||d }|j||jd }t|| d S )Nr  r   r   r  )r  r1   r}  r:   valuesr@   ra   )r2   r   r  r  r   rE   rD   r   r   r   test_times_array  s   

zTestEWM.test_times_arrayc              	   C   s   t g dtddddtdd}|djdd	d
}| }t tdddddgd dtjg dttdgdd gdd}t	
|| |d  }tdgd tjg dtddddgdd	gddd}t	|| | }t	
|| d S )N)r   r   r   r   r   z2021-9-1rv   r  )r  ra  )r   	timestampr   r   Z1Hr  r  r   )r  r   r}   r   r   r   )r   r   r_   r1   r9   rI   r   r^   r   r@   ra   r   rA   )r2   r   grprE   Zexpected_dfZexpected_seriesr   r   r   "test_dont_mutate_obj_after_slicing  s>   	z*TestEWM.test_dont_mutate_obj_after_slicingN)r   r   r   r.   rg  rh  r&   r   r  r  r  r  r  r   r   r   r   rz  G  s"    

"
'	rz  )numpyr&   r.   Zpandasr   r   r   r   r   r   r   Zpandas._testingZ_testingr@   Zpandas.api.indexersr	   Zpandas.core.groupby.groupbyr
   ry  r   r(   r*   ri  rz  r   r   r   r   <module>   s(    $	

       >`