o
    Eh[F                     @   s^  d Z ddlZddlZddlmZ ddlZddlmZ ddl	m
Z ddlmZ ddlmZ dd Zejd	d
 Zejdejgddd Zejdd Zejdejgddd Zejdejgddd Zejdejgddd Zejdejgddd Zejdd Zejdejgddd Zejdejgddd ZG dd dejZd d! ZdS )"aj  
This file contains a minimal set of tests for compliance with the extension
array interface test suite, and should contain no other tests.
The test suite for the full functionality of the array is located in
`pandas/tests/arrays/`.

The tests in this file are inherited from the BaseExtensionTests, and only
minimal tweaks should be applied to get the tests passing (by overwriting a
parent method).

Additional tests should either be added to one of the BaseExtensionTests
classes (if they are relevant for the extension interface for all dtypes), or
be added to the array-specific tests in `pandas/tests/arrays/`.

    N)PerformanceWarningSparseDtype)SparseArray)basec                 C   sh   t jd}t | r|jdd}n|jdddtd}|d |d kr+|d  d7  < | |dd d< |S )N   d   )size   )r	   dtyper      )npZrandomZdefault_rngisnanZuniformZintegersint)
fill_valuerngdata r   Blib/python3.10/site-packages/pandas/tests/extension/test_sparse.py	make_data   s   
r   c                   C   s   t  S Nr   r   r   r   r   r   *   s   r   )paramsc                 C   s   t t| j| jd}|S )z*Length-100 PeriodArray for semantics test.r   )r   r   param)requestresr   r   r   r   /   s   r   c                   C   s   t tdd S )Nr   r   )r   r   Zonesr   r   r   r   data_for_twos6      r   c                 C   s   t tjdg| jdS )zLength 2 array with [NA, Valid]r
   r   r   r   nanr   r   r   r   r   data_missing;   s   r!   c                 #   s     fdd}|V  dS )z1Return different versions of data for count timesc                 3   s*    t | D ]}tt j jdV  qd S )Nr   )ranger   r   r   )count_r    r   r   genE   s   zdata_repeated.<locals>.genNr   )r   r%   r   r    r   data_repeatedA   s   
r&   c                 C   s   t g d| jdS )N)r   r   r
   r   )r   r   r    r   r   r   data_for_sortingL   r   r'   c                 C   s   t dtjdg| jdS )Nr   r
   r   r   r    r   r   r   data_missing_for_sortingQ   s   r(   c                   C   s   dd S )Nc                 S   s   t | o	t |S r   )pdisna)leftrightr   r   r   <lambda>X       zna_cmp.<locals>.<lambda>r   r   r   r   r   na_cmpV   s   r/   c              	   C   s"   t ddtjtjddddg| jdS )Nr
   r   r   r   r   r    r   r   r   data_for_grouping[   s   "r0   c                 C   s$   t ddtjdddddddg
| jdS )Nr      r   r   r   r   r    r   r   r   data_for_compare`   s   $r4   c                       sh  e Zd ZdedefddZejdddg fdd	Z	ejdddg fd
dZ
dd Zdd Zejdejdddgejjddgddgdgejdddg fddZ fddZ fddZ fd d!Z fd"d#Z fd$d%Z fd&d'Zd(d) Z fd*d+Zd,d- Z fd.d/Z fd0d1Zejjd2d3 fd4d5Zd6d7 Zd8Z d9d: Z!d;d< Z"ejjd=d3 fd>d?Z#d@dA Z$ fdBdCZ%dDdE Z&ejdFdGdHg fdIdJZ'ejj(ejdKej)ej*ej+g fdLdMZ,ejj( fdNdOZ-ejdPdQdR dSe.dTe/j0gfdUdR dVe.dTe/j0gfedSe.dWdXgdXdYfedVe.dWe/j0gfgdZd[ Z1ejd\dSdVgd]d^ Z2ejje3d_d` fdadbZ4dSZ5dSZ6dSZ7dSZ8dcdd Z9 fdedfZ: fdgdhZ; fdidjZ<dkej*dle.fdmdnZ=dle.fdodpZ>dle.fdqdrZ?dle.fdsdtZ@ejjdud3 fdvdwZAejjdxd3ejdyddg fdzd{ZB  ZCS )|TestSparseArrayop_namereturnc                 C   s   dS )NTr   )selfobjr6   r   r   r   _supports_reductionf   s   z#TestSparseArray._supports_reductionskipnaTFc                    j   |dv rt jjdd}|j| n|dv r+|jjdkr+|s+t jjdd}|j| t ||| d S )NZprodZmedianvarZstdZsemZskewZkurt,This should be viable but is not implementedreasonsummaxminZmeanfzgetting a non-nan float)	pytestmarkxfailnode
add_markerr   kindsupertest_reduce_series_numericr8   r   Zall_numeric_reductionsr;   r   rH   	__class__r   r   rN   i      	z*TestSparseArray.test_reduce_series_numericc                    r<   )Nr=   r?   r@   rB   rF   z$ExtensionArray NA mask are different)	rG   rH   rI   rJ   rK   r   rL   rM   test_reduce_framerO   rP   r   r   rS      rR   z!TestSparseArray.test_reduce_framec                 C   s"   |j ttdkrtd d S d S )Nr   zCan't store nan in int array.)r   r   r   rG   skipr8   r   r   r   r   _check_unsupported   s   z"TestSparseArray._check_unsupportedc                 C   sx   t d|d d i}t dg di}t dg did}|||g}t |}t dd |D }t|| d S )NAr   )r
   r   r   )abccategoryc                 S   s   g | ]	}| d d qS )c                 S   s   t | tS r   )r   asarrayastypeobject)sr   r   r   r-      s    zETestSparseArray.test_concat_mixed_dtypes.<locals>.<listcomp>.<lambda>)Zapply).0xr   r   r   
<listcomp>   s    z<TestSparseArray.test_concat_mixed_dtypes.<locals>.<listcomp>)r)   	DataFramer]   concattmassert_frame_equal)r8   r   Zdf1Zdf2Zdf3Zdfsresultexpectedr   r   r   test_concat_mixed_dtypes   s   

z(TestSparseArray.test_concat_mixed_dtypesz9ignore:The previous implementation of stack is deprecatedcolumnsrW   B)rW   rX   )rW   rY   Zouterinner)namesfuture_stackc                    s   t  ||| d S r   )rM   
test_stack)r8   r   rj   rn   rP   r   r   ro      s   zTestSparseArray.test_stackc                       |  | t || d S r   )rV   rM   test_concat_columnsr8   r   na_valuerP   r   r   rq         
z#TestSparseArray.test_concat_columnsc                    rp   r   )rV   rM   'test_concat_extension_arrays_copy_falserr   rP   r   r   ru      rt   z7TestSparseArray.test_concat_extension_arrays_copy_falsec                    rp   r   )rV   rM   
test_alignrr   rP   r   r   rv      rt   zTestSparseArray.test_alignc                    rp   r   )rV   rM   test_align_framerr   rP   r   r   rw      rt   z TestSparseArray.test_align_framec                    rp   r   )rV   rM   test_align_series_framerr   rP   r   r   rx      rt   z'TestSparseArray.test_align_series_framec                    rp   r   )rV   rM   
test_mergerr   rP   r   r   ry      rt   zTestSparseArray.test_mergec                 C   s   t j|dd tt|D d}t|jjr*t|dr't|j	d s)J n|d|j	d ks6J |d|j	d ksBJ d S )Nc                 S   s   g | ]}d | qS )r   r   )r`   ir   r   r   rb      r.   z,TestSparseArray.test_get.<locals>.<listcomp>)indexr3   r   r
   )
r)   Seriesr"   lenr   r   valuesr   getiloc)r8   r   serr   r   r   test_get   s
    &zTestSparseArray.test_getc                    rp   r   )rV   rM   test_reindexrr   rP   r   r   r      rt   zTestSparseArray.test_reindexc                 C   s   t |}ttt|jj}t ddg|d}| }t|| |	d}ttt|jj}t ddgd|d}t
| | d S )NTFr   r   r   r   )r   r   boolr)   r*   r   r   re   Zassert_sp_array_equalfillnaassert_equal)r8   r!   ZsarrZexpected_dtyperh   rg   r   r   r   	test_isna   s   
zTestSparseArray.test_isnac                    sH   t tf}tj|dd t | W d    d S 1 sw   Y  d S NF)Zcheck_stacklevel)r   FutureWarningre   assert_produces_warningrM   test_fillna_limit_backfill)r8   r!   ZwarnsrP   r   r   r      s   "z*TestSparseArray.test_fillna_limit_backfillc                    s0   t |jr|tjjdd t | d S )Nz'returns array with different fill valuer@   )	r   r   r   applymarkerrG   rH   rI   rM   test_fillna_no_op_returns_copy)r8   r   r   rP   r   r   r      s
   z.TestSparseArray.test_fillna_no_op_returns_copyZUnsupportedr@   c                    s   t    d S r   )rM   test_fillna_seriesr8   r!   rP   r   r   r      s   z"TestSparseArray.test_fillna_seriesc                 C   st   |d }t |ddgd|}t |jrt|j|}n|j}t |j||g|dddgd}t	|| d S )Nr
   r   )rW   rk   r   )
r)   rc   r   r*   r   r   r   _from_sequencere   rf   )r8   r!   r   rg   r   rh   r   r   r   test_fillna_frame  s   z!TestSparseArray.test_fillna_framezSparse[bool]c                 C   s   | ddg}tjd|idd}|jd }||}t|jdr7|r-|jj|jju s,J n
|jj|jjus7J |j	j
 | u sCJ d S )Nr
   rW   Fcopy)r   r   Zblocks)taker)   rc   r   r   hasattrZ_mgrr~   r   rW   _valuesto_dense)r8   r!   using_copy_on_writearrZdf
filled_valrg   r   r   r   test_fillna_copy_frame  s   

z&TestSparseArray.test_fillna_copy_framec                 C   sn   | ddg}tj|dd}|d }||}|r"|j|ju s!J n|j|jus*J |j | u s5J d S )Nr
   Fr   r   )r   r)   r|   r   r   r   )r8   r!   r   r   r   r   rg   r   r   r   test_fillna_copy_series&  s   
z'TestSparseArray.test_fillna_copy_serieszNot Applicablec                       t  | d S r   )rM   test_fillna_length_mismatchr   rP   r   r   r   4     z+TestSparseArray.test_fillna_length_mismatchc                 C   s   |d |d ks
J t |}|d d \}}t|j||||g|jd}tg d}||}tdd}	t|j||||g|	d}
t	
||
 |j||||g|jd}tg d}|||}t|j||||g|jd}
t	
||
 d S )	Nr   r
   r   r   )TTFFfloatg        )TFTT)typer)   r|   r   r   r   arraywherer   re   Zassert_series_equal)r8   r   rs   clsrX   rY   r   Zcondrg   Z	new_dtyperh   otherr   r   r   test_where_series8  s    

z!TestSparseArray.test_where_seriesc                    sB   t jtdd t || W d    d S 1 sw   Y  d S r   )re   r   r   rM   test_searchsorted)r8   r'   	as_seriesrP   r   r   r   N  s   "z!TestSparseArray.test_searchsortedc                 C   s6   | d}|jd |jd< |jd |jd ksJ d S )Nr   r
   )shiftZ_sparse_values)r8   r   rg   r   r   r   test_shift_0_periodsR  s   
z$TestSparseArray.test_shift_0_periodsmethodZargmaxZargminc                    s   |  | t ||| d S r   )rV   rM   test_argmin_argmax_all_na)r8   r   r   rs   rP   r   r   r   Y  s   
z)TestSparseArray.test_argmin_argmax_all_naboxc                    s    |  | t |||| d S r   )rV   rM   test_equals)r8   r   rs   r   r   rP   r   r   r   _  s   
zTestSparseArray.test_equalsc                    r   r   )rM   &test_equals_same_data_different_objectrU   rP   r   r   r   e  r   z6TestSparseArray.test_equals_same_data_different_objectzfunc, na_action, expectedc                 C      | S r   r   ra   r   r   r   r-   l      zTestSparseArray.<lambda>Ng      ?c                 C   r   r   r   r   r   r   r   r-   m  r   ignorez1.0r   r   c                 C   s,   t dtjg}|j||d}t|| d S )Nr
   	na_action)r   r   r   mapre   Zassert_extension_array_equal)r8   funcr   rh   r   rg   r   r   r   test_mapi  s   zTestSparseArray.test_mapr   c                 C   sJ   d}t jt|d |jdd |d W d    d S 1 sw   Y  d S )Nz-fill value in the sparse values not supported)matchc                 S   s   t jS r   )r   r   r   r   r   r   r-   }  s    z1TestSparseArray.test_map_raises.<locals>.<lambda>r   )rG   raises
ValueErrorr   )r8   r   r   msgr   r   r   test_map_raisesx  s   "zTestSparseArray.test_map_raiseszno sparse StringDtype)r   rA   c                    r   r   )rM   test_astype_string)r8   r   Znullable_string_dtyperP   r   r   r     s   z"TestSparseArray.test_astype_stringc                 C   s   |j dkrtd d S d S )Nr   z=Incorrected expected from Series.combine and tested elsewhere)r   rG   rT   rU   r   r   r   _skip_if_different_combine  s   
z*TestSparseArray._skip_if_different_combinec                    rp   r   )r   rM   test_arith_series_with_scalarr8   r   all_arithmetic_operatorsrP   r   r   r     rt   z-TestSparseArray.test_arith_series_with_scalarc                    rp   r   )r   rM   test_arith_series_with_arrayr   rP   r   r   r     rt   z,TestSparseArray.test_arith_series_with_arrayc                    sF   |j jdkrn|ddvrtjjdd}|| t || d S )Nr   r$   )mulZrmulfloordivZ	rfloordivpowmodZrmodz result dtype.fill_value mismatchr@   )	r   r   striprG   rH   rI   r   rM   test_arith_frame_with_scalar)r8   r   r   r   rH   rP   r   r   r     s   	
z,TestSparseArray.test_arith_frame_with_scalarr   r4   c           	      C   s   |}|||}t |tjrt |tjsJ t |jtsJ nt |ts%J |jjtjks.J t |tjrM||j	|j
j	}t|| t||tjd}nt|t|j	t|}t|| t||tjd}t |tjrwt|}t|| d S )Nr   )
isinstancer)   r|   r   r   r   Zsubtyper   Zbool_r   r   r   r\   allre   r   )	r8   r   r4   comparison_opr   oprg   r   rh   r   r   r   _compare_other  s2   

zTestSparseArray._compare_otherc                 C   sP   t |}| |||d | |||d | |||d | |||tj d S )Nr   r
   r2   )r)   r|   r   r   r   )r8   r4   r   r   r   r   r   test_scalar  s
   
zTestSparseArray.test_scalarc                 C   sZ   |j jdkr|jdv rtjjdd}|| tddd}t	
|}| |||| d S )Nr   )eqgeleWrong fill_valuer@      
   )r   r   __name__rG   rH   rI   r   r   Zlinspacer)   r|   r   )r8   r4   r   r   rH   r   r   r   r   r   
test_array  s   

zTestSparseArray.test_arrayc                 C   sl   |j jdkr|jdkrtjjdd}|| t|}|d }| 	|||| |d }| 	|||| d S )Nr   gtr   r@   r
   r   )
r   r   r   rG   rH   rI   r   r)   r|   r   )r8   r4   r   r   rH   r   r   r   r   r   test_sparse_array  s   

z!TestSparseArray.test_sparse_arrayzDifferent reprc                       t  || d S r   )rM   test_array_repr)r8   r   r	   rP   r   r   r     r   zTestSparseArray.test_array_reprzresult does not match expectedas_indexc                    r   r   )rM   test_groupby_extension_agg)r8   r   r0   rP   r   r   r     s   z*TestSparseArray.test_groupby_extension_agg)Dr   
__module____qualname__strr   r:   rG   rH   ZparametrizerN   rS   rV   ri   filterwarningsr)   Z
MultiIndexZfrom_tuplesro   rq   ru   rv   rw   rx   ry   r   r   r   r   r   rI   r   r   Z_combine_le_expected_dtyper   r   r   r   r   r   r   Zfails_arm_wheelsr   r|   rc   r   r   r   r   r   r   r   	TypeErrorr   Zseries_scalar_excZframe_scalar_excZ
divmod_excZseries_array_excr   r   r   r   r   r   r   r   r   r   __classcell__r   r   rP   r   r5   e   s    	
	

$r5   c                 C   s   |   tu sJ d S r   )Zconstruct_array_typer   r   r   r   r   test_array_type_with_arg  s   r   ) __doc__Znumpyr   rG   Zpandas.errorsr   Zpandasr)   r   Zpandas._testingZ_testingre   Zpandas.arraysr   Zpandas.tests.extensionr   r   Zfixturer   r   r   r   r!   r&   r'   r(   r/   r0   r4   ZExtensionTestsr5   r   r   r   r   r   <module>   sF    










   