o
    ҷh6                     @   s  d dl Z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 ejeg deg dfeg deg dfedd	d
edd	d
fe
dd	d
e
dd	d
fe
dd	dde
dd	ddfgdd ddd ZG dd dZG dd dZG dd dZdd ZG dd dZdd  Zd!d" Zd#d$ Zejjd%g d&e
d'd(d)d*gd+d,gd-d.d/ Zejjd%g d&e
d'd(d)d*gd+d,gd-d0d1 Zd2d3 Zejd4g d5d6d7 Z dS )8    N)IndexIntervalIntervalIndex	Timedelta	Timestamp
date_rangetimedelta_range)IntervalArray)r         )         )              ?       @)r   r         @0 daysr   periods1 day2017010120170102
US/Easternr   tzc                 C   s   t | d jS )Nr   )strdtype)x r   ]/var/www/html/venv/lib/python3.10/site-packages/pandas/tests/arrays/interval/test_interval.py<lambda>   s    r!   )paramsidsc                 C   s   | j S )zC
    Fixture for building an IntervalArray from various dtypes
    )param)requestr   r   r    left_right_dtypes   s   r&   c                   @   sh   e Zd Zejddededfededfedddedddfgejd	ee	gd
d Z
dS )TestAttributeszleft, rightr   r   r   r   z
2018-01-01z
2018-01-02r   )r   constructorc                 C   sH   ||f||ft jg}t |dkddg}|j||dj}t|| d S )NbothFclosed)npnanarrayfrom_tuplesis_emptytmassert_numpy_array_equal)selfr)   leftrightr,   tuplesexpectedresultr   r   r    test_is_empty'   s   zTestAttributes.test_is_emptyN)__name__
__module____qualname__pytestmarkparametrizer   r   r	   r   r:   r   r   r   r    r'   &   s    

r'   c                   @   sh   e Zd Zejdg ddd Zejdedddd	ej	g d
dd	gdd Z
dd Zdd ZdS )TestMethods
new_closed)r5   r6   r*   neitherc                 C   s>   t jtd|d}||}t jtd|d}t|| d S )N
   r+   )r	   from_breaksrange
set_closedr2   assert_extension_array_equal)r4   r,   rB   r/   r9   r8   r   r   r    test_set_closed=   s   
zTestMethods.test_set_closedotherr   r   r6   r+   r   r
   r   r   c                 C   s   t tjg ddd}tg d}d}tjt|d |j	|| W d    n1 s.w   Y  |j
||d}|t
||}t|| d S )NrK   r5   r+   )TFTz+'value.closed' is 'right', expected 'left'.match)rJ   )pdSeriesr	   rE   r-   r/   r>   raises
ValueError_wherewhereastypeobjectr2   assert_series_equal)r4   rJ   sermaskrM   resr8   r   r   r    test_where_raisesE   s   	zTestMethods.test_where_raisesc                 C   s   t g d}| }t tjtjfdg}t|| d}tj	t
|d |jdtjd W d    d S 1 s9w   Y  d S )N)r   r
   r   )r   r   =can only insert Interval objects and NA into an IntervalArrayrL   r   
fill_value)r	   rE   shiftr0   r-   r.   r2   assert_interval_array_equalr>   rP   	TypeErrorrN   NaTr4   ar9   r8   msgr   r   r    
test_shiftX   s   "zTestMethods.test_shiftc                 C   s   t tddd}|d}|jg ddd}t|| |d}|jg d	dd}t|| d
}tjt	|d |jdt
ddd W d    d S 1 sQw   Y  d S )N2000r   r   r
   )rg   r   T)
allow_fillrg   )r   r
   rg   r[   rL   r   ra   nsr\   )r	   rE   r   r^   taker2   r_   r>   rP   r`   r-   timedelta64rb   r   r   r    test_shift_datetimed   s   

"zTestMethods.test_shift_datetimeN)r;   r<   r=   r>   r?   r@   rI   r   r	   rE   rZ   re   rl   r   r   r   r    rA   <   s    

rA   c                   @      e Zd Zdd Zdd ZdS )TestSetitemc           	      C   s$  |\}}|j dd}|j dd}t||}|jjjdvr;d}tjt|d t	j
|d< W d    n1 s6w   Y  |jjjdv rcd}tjt|d tj|d< W d    d S 1 s\w   Y  d S tj|d< t|jgt|d	d   }t|jgt|d	d   }t||}t|| d S )
NTdeep)mMz='value' should be an interval type, got <.*NaTType'> instead.rL   r   )iuz4Cannot set float NaN to integer-backed IntervalArrayr   )copyr	   from_arraysr   subtypekindr>   rP   r`   rN   ra   r-   r.   r   	_na_valuelistr2   rH   )	r4   r&   r5   r6   r9   rd   expected_leftexpected_rightr8   r   r   r    test_set_nau   s,   

zTestSetitem.test_set_nac                 C   s  t td}| }|d}d}tjt|d |d |d< W d    n1 s*w   Y  tjt|d |d d |d d< W d    n1 sKw   Y  tjt|d |d d |d d< W d    n1 slw   Y  tjt|d |d d d |d d < W d    n1 sw   Y  tjt|d t|d d d |d d < W d    n1 sw   Y  tjt|d |d d d 	t
|d d < W d    n1 sw   Y  tjt|d |d d d 	d|d d < W d    n1 sw   Y  g |d d< t|| d S )	Nr   r*   z*'value.closed' is 'both', expected 'right'rL   r   r   rg   category)r	   rE   rF   ru   rG   r>   rP   rQ   rz   rT   rU   r2   r_   )r4   arrorigrJ   rd   r   r   r    test_setitem_mismatched_closed   s6   
z*TestSetitem.test_setitem_mismatched_closedN)r;   r<   r=   r}   r   r   r   r   r    rn   t   s    rn   c                  C   s*   t ddg} t| }d}||ksJ d S )Nr(   )r   r
   zI<IntervalArray>
[(0, 1], (1, 2]]
Length: 2, dtype: interval[int64, right])r	   r0   repr)r   r9   r8   r   r   r    	test_repr   s
   r   c                   @   rm   )TestReductionsc              	   C   s,  |\}}|j dd}|j dd}t||}d}dD ]<}tjt|d |j|d W d    n1 s4w   Y  tjt|d |j|d W d    n1 sQw   Y  qd}tjt|d |jdd W d    n1 sqw   Y  tjt|d |jdd W d    d S 1 sw   Y  d S )	NTro   z2`axis` must be fewer than the number of dimensions)r   rL   )axisz'>=' not supported betweenfoo)	ru   r	   rv   r>   rP   rQ   minmaxr`   )r4   r&   r5   r6   r   rd   r   r   r   r    test_min_max_invalid_axis   s(   "z(TestReductions.test_min_max_invalid_axisc                 C   s  |\}}|j dd}|j dd}t||}|jsJ t|js"J |d }|d }tt|}tj	d
| ||}|dtj}	||}||	}	dD ].}
|j|
d}||ks^J t|t|kshJ |j|
d}||kstJ t|t|ks~J qP|	jdd}t|sJ |	jdd}t|sJ |	jdd}||ksJ t|t|ksJ |	jdd}||ksJ t|t|ksJ d S )	NTro   r   rg   r
   )TF)skipnaF)ru   r	   rv   is_monotonic_increasingr   r-   arangelenrandomdefault_rngshufflerj   insertr.   r   typer   isnan)r4   r&   index_or_series_or_arrayr5   r6   r   MINMAXindexerarr_nar   rY   r   r   r    test_min_max   s>   

zTestReductions.test_min_maxN)r;   r<   r=   r   r   r   r   r   r    r      s    r   c                  C   s   t d} ddlm} ||  d}||  d}||  d}|jdks'J ||ks-J ||ks3J t|t|ks=J t|t|ksGJ d S )Npyarrowr   ArrowIntervalTyper5   r6   )r>   importorskip(pandas.core.arrays.arrow.extension_typesr   int64r,   hash)par   p1p2p3r   r   r    test_arrow_extension_type   s   
r   c                  C   s  t d} ddlm} tjddddj}| |}t|j|s!J |jj	|j	ks*J |jj
|  ks4J |jd| jg dd	d
sFJ |jd| jg dd	d
sXJ | dd tddD }|j|smJ | j||jd
}||s|J t jtdd | j|dd
 W d    n1 sw   Y  t jtdd | j|||  dd
 W d    d S 1 sw   Y  d S )Nr   r   r   r   r   )freqr5   rK   r   r   r6   )r
   r   r   r   c                 S   s   g | ]	}||d  dqS )r   r5   r6   r   ).0rs   r   r   r    
<listcomp>  s    z$test_arrow_array.<locals>.<listcomp>z&Not supported to convert IntervalArrayrL   float64)r>   r   r   r   rN   interval_ranger/   
isinstancer   r,   rw   r   storagefieldequalsrF   rP   r`   r   )r   r   	intervalsr9   r8   r   r   r    test_arrow_array
  s&   

$$"r   c                  C   s  t d} ddlm} tg d}d |d< | |}t|j|s#J |jj	|j	ks,J |jj
|  ks6J | jg ddd}| jg d	dd}|jd
|sSJ |jd|s^J dddd d ddddg}| jj|tg dd}|j|sJ d S )Nr   r   r   r   r   r   r   r   )r   Nr   r   r   )r   Nr   r5   r6   r   r   r   r   r   )FTF)rX   )r>   r   r   r   r	   rE   r/   r   r   r,   rw   r   r   r   r   StructArrayfrom_pandasr-   )r   r   r   r9   r5   r6   valsr8   r   r   r    test_arrow_array_missing'  s$   

r   breaksr   2017r   D)r   r   floatzdatetime64[ns])r#   c           	      C   s   t d}ddlm} t| }d |d< td|i}||}t	|
dj|s+J | }t	|d jtjs:J t|| |||g}| }tj||gdd}t|| |j|jg |djdg|jd	}| }t||dd  d S )
Nr   r   r   r   rc   T)ignore_indexr   )schema)r>   r   r   r   r	   rE   rN   	DataFrametabler   r   r   	to_pandasr   IntervalDtyper2   assert_frame_equalconcat_tablesconcatchunked_arraycolumnr   )	r   r   r   r   dfr   r9   table2r8   r   r   r    test_arrow_table_roundtripD  s&   


r   c                 C   sz   t d}t| }d |d< td|i}||}| }|jj	d u s&J |
 }t|d jtjs5J t|| d S )Nr   r   rc   )r>   r   r	   rE   rN   r   r   replace_schema_metadatar   metadatar   r   r   r   r2   r   )r   r   r   r   r   r9   r   r   r    +test_arrow_table_roundtrip_without_metadatae  s   


r   c                  C   s   t d} | ddddddg}tjtddd}||}tj	tjg d	dd
dd}t
|| || |g}t
|| d S )Nr   r   r   r   r
   r   rC   r+   )r   r   r
   )r   )r>   r   r/   rN   r   r-   r   __from_arrow__r	   rE   r2   rH   r   )r   r   r   r9   r8   r   r   r    %test_from_arrow_from_raw_struct_array{  s   

r   timezone)UTCz
US/PacificGMTc                 C   sj   t dd| d}d|  d| d}tjddgdd	g||d
}tj|d d |dd  |d}t|| d S )N2022r   r   zinterval[datetime64[ns, z], ]z
2022-01-01z
2022-01-02z
2022-01-03)r,   r   rg   r   r+   )r   r   rv   r2   assert_index_equal)r   inclusive_endpoints_fixturedatesr   r9   r8   r   r   r    test_interval_index_subtype  s   r   )!numpyr-   r>   pandasrN   r   r   r   r   r   r   r   pandas._testing_testingr2   pandas.core.arraysr	   fixturer&   r'   rA   rn   r   r   r   r   r   r?   r@   r   r   r   r   r   r   r   r    <module>   sR    $	

85E

