U
    f/e,                     @   sh   d dl Zd dlZd dlm  mZ d dlZd dlm	Z	m
Z
mZmZmZmZ d dlmZ G dd dZdS )    N)CategoricalIndex	DataFrameIndexSeries
date_rangeoffsetsc                   @   s   e Z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 Zdd Zdd Zdd Zdd Zdd Zejdd Zdd Zd S )!TestDataFrameShiftc              	   C   s  | d}t|j|j |d  d}t|d | | d}t|j|j |d  d}t|d | | d}t|| |j dt d}t|t|kst	|j ddd}t|| |jd }|td }tj|
||
|dd | d	}	t }
|
 d	}| d
}t|j|
j t|j|
j t|jd d df  j|
jd d
df j |
 d	d}|
 d	t }t|| t|
| d
d d}tjt|d |
j dd W 5 Q R X ttjdd}tjttj|jdgd|jd d dd
f gdd	d}|j d	d	d}t|| ttjdd}tjttj|jdgd|jd d dd
f gdd	d}|j d	dd}t|| d S )N   Ar   freqBF)Zcheck_names   zdoes not match PeriodIndex freqmatchD
   indexcolumnsT)Zignore_indexaxisr   r   )shifttmZassert_index_equalr   assert_series_equalassert_frame_equalr   BDaylenAssertionErrorxsmakePeriodFrameZassert_numpy_array_equalilocZdropnavaluespytestraises
ValueErrorr   nprandomZrandpdconcatnan)selfdatetime_frameZ	int_frameZshiftedFrameZshiftedSeries	unshiftedZshiftedFrame2dZ	shifted_dZint_shiftedpsshiftedshifted2shifted3msgdfexpectedresult r9   I/tmp/pip-unpacked-wheel-tiezk1ph/pandas/tests/frame/methods/test_shift.py
test_shift   sl    



  


 ((zTestDataFrameShift.test_shiftc                 C   sZ   t ddgddgd}|d}t tjtjtjgddggtdddgd}t|| d S )	NTF)highlowr   Zdtyper<   r=   r   )r   r   r(   arrayr,   objectr   r   )r-   r6   rsxpr9   r9   r:   test_shift_bool]   s    
z"TestDataFrameShift.test_shift_boolc                 C   sf   t dddgdd}t dddgdd}t||d	}|d
}t|d
|d
d	}t|| d S )Nabccategoryr>   r
   r   C)ZoneZtwor   )r   r   r   r   r   )r-   s1s2r6   rB   rC   r9   r9   r:   test_shift_categoricalf   s    
z)TestDataFrameShift.test_shift_categoricalc                 C   s   t dddddgtddddd	}t d
ddddgtddddd	}|jdd
d}t|| t d
d
dddgtddddd	}|jdd
d}t|| d S )Nr            r	   z1/1/2000H)Zperiodsr   )r   r   
fill_value)r   r   r   r   r   )r-   r6   expr8   r9   r9   r:   test_shift_fill_valueo   s        z(TestDataFrameShift.test_shift_fill_valuec                 C   s&   t dg i}|d}t|| d S )NZfoor   )r   r   r   r   )r-   r6   rB   r9   r9   r:   test_shift_empty   s    
z#TestDataFrameShift.test_shift_emptyc                 C   s   t tddgd dddddgg}tjdd}g }|D ]`}t| |d}tdD ].}|jd d |f |d |jd d |f< qTtd|_	|
| q8|d   }t|ttdddd	 t|d |d  t|d |d  d S )
Nr	   r   rM      r?   r      Zint64r>   )listranger(   r)   Zrandnr   copyr#   r   r   appendZisnasumr   r   r   r   )r-   Zcolumn_listsdatar2   r   r6   sZnullsr9   r9   r:   test_shift_duplicate_columns   s    ",
z/TestDataFrameShift.test_shift_duplicate_columnsc                 C   s$  t tjjddd}t tjjddd}tj||gdd}|sRt|jjdksRt	|j
ddd}|jddd	ddgdd}tj|jd d d df< |j|_t|| tj||gdd}|st|jjdkst	|j
d
dd}|jdddddgdd}tj|jd d d
d f< |j|_t|| d S )Ni  )r	   rN   )size)r	   rM   r   r   rM   r   r   rN   rO   )r   r(   r)   randintr*   r+   r   Z_mgrblocksr    r   Ztaker,   r#   r   r   r   )r-   Zusing_array_managerZdf1df2Zdf3r8   r7   r9   r9   r:    test_shift_axis1_multiple_blocks   s$    z3TestDataFrameShift.test_shift_axis1_multiple_blocksz)ignore:tshift is deprecated:FutureWarningc              	   C   sf  t  }|d}|d}t || |jdd}t || |jt d}t || d}tjt|d |jdd W 5 Q R X |d}|d}t || |j|j	j
d}t || t|jtt|j	|jd}|d}|d}	|	j	d |	_	t ||	 |d}t || |jd	d
dgd d f }
d}tjt|d |
  W 5 Q R X d S )Nr   r   r   r   .Given freq M does not match PeriodIndex freq Br   Mr?   r   r	      6Freq was not set in the index hence cannot be inferred)r   r"   tshiftr   r   r   r%   r&   r'   r   r   r   r$   r   r(   asarrayr   
_with_freqr#   )r-   r.   r1   r2   r/   r3   r4   r5   inferred_tsr7   no_freqr9   r9   r:   test_tshift   s>    






zTestDataFrameShift.test_tshiftc              	   C   s"   t t |  W 5 Q R X d S )N)r   assert_produces_warningFutureWarningrj   )r-   r.   r9   r9   r:   test_tshift_deprecated   s    z)TestDataFrameShift.test_tshift_deprecatedc                 C   sh   t  }|jddd}|jddd}t || |jdd}t || |jt d}t || d S )Nr   inferr   r   r   )r   r"   r   r   r   r   )r-   r1   r2   r/   r3   r4   r9   r9   r:   'test_period_index_frame_shift_with_freq   s    z:TestDataFrameShift.test_period_index_frame_shift_with_freqc                 C   s   |j ddd}|j ddd}t|| |j |jjd}t|| t|jtt	|j|j
d}|j ddd}|j ddd}|jd |_t|| |j ddd}t|| d S )Nr   rs   r   r   r?   )r   r   r   r   r   r   r$   r   r(   rk   r   rl   )r-   r.   r2   r/   r3   rm   r7   r9   r9   r:   #test_datetime_frame_shift_with_freq   s     z6TestDataFrameShift.test_datetime_frame_shift_with_freqc              	   C   s6   t  }d}tjt|d |jdd W 5 Q R X d S )Nrf   r   rg   r   )r   r"   r%   r&   r'   r   )r-   r1   r5   r9   r9   r:   -test_period_index_frame_shift_with_freq_error  s    z@TestDataFrameShift.test_period_index_frame_shift_with_freq_errorc              	   C   sF   |j dddgd d f }d}tjt|d |jdd W 5 Q R X d S )Nr   r	   rh   ri   r   rs   r   )r#   r%   r&   r'   r   )r-   r.   rn   r5   r9   r9   r:   )test_datetime_frame_shift_with_freq_error  s    z<TestDataFrameShift.test_datetime_frame_shift_with_freq_errorc              	   C   s   t tdtdg}| }tt |jddd}W 5 Q R X t td|d g }t|| t	||d}|
  tt |jdddd}W 5 Q R X t	tdtdg|d d}t|| d S )	Nz
2020-01-01z
2020-01-02r   r   rQ   )r
   r   )r   rR   r
   )r   r*   Z	TimestampZto_framer   rp   rq   r   r   r   Z_consolidate_inplace)r-   Zserr6   r8   r7   rd   r9   r9   r:   )test_shift_dt64values_int_fill_deprecated  s    "z<TestDataFrameShift.test_shift_dt64values_int_fill_deprecatedc                 C   s   t dddg}tddgddgdd	gd
|d d |d}|jdd}ttjtjgddgddgd
|d d |d}t|| |jddd}ttjtjgtjtjgddgd
|d d |d}t|| d S )NrE   rF   rG   r   rN   rM   rO   r	   rW   )rE   rF   rG   r   r   r   )r   r   r   r(   r,   r   r   )r-   cir6   r8   r7   r9   r9   r:   $test_shift_axis1_categorical_columns4  s(     
  
 
z7TestDataFrameShift.test_shift_axis1_categorical_columnsN)__name__
__module____qualname__r;   rD   rL   rT   rU   r_   re   r%   markfilterwarningsro   rr   rt   ru   rv   rw   tdZ&skip_array_manager_not_yet_implementedrx   rz   r9   r9   r9   r:   r      s"   J		

0
r   )Znumpyr(   r%   Zpandas.util._test_decoratorsutilZ_test_decoratorsr   Zpandasr*   r   r   r   r   r   r   Zpandas._testingZ_testingr   r   r9   r9   r9   r:   <module>   s    