o
    d/#                     @   s  d Z ddl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lmZ dd	 Zejd
dgddd Zejdd Zejdd Zejdd Zejdd Zejdd Zejdd Zejdd ZG dd dejZG dd  d ej Z!G d!d" d"ej"Z#G d#d$ d$ej$Z%G d%d& d&ej&Z'G d'd( d(ej(Z)G d)d* d*ej*Z+G d+d, d,ej,Z-G d-d. d.ej.Z/G d/d0 d0ej0Z1G d1d2 d2ej2Z3G d3d4 d4ej4Z5G d5d6 d6ej6Z7G d7d8 d8ej8Z9G d9d: d:ej: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is_string_dtype)ArrowStringArrayStringDtype)basec                 C   s&   | j jdkrtd dd }|| S )Npyarrowz-only applicable for pyarrow chunked array n/ac                 S   sZ   dd l }| j}t|d }|g |d | j||d  j}|jdks'J t| |S )Nr      )r	   _datalenZchunked_arraychunksZ
num_chunkstype)arrpaZarrow_arraysplit r   V/app/.heroku/python/lib/python3.10/site-packages/pandas/tests/extension/test_string.py_split_array#   s   z!split_array.<locals>._split_array)dtypestoragepytestskip)r   r   r   r   r   split_array   s   
r   TF)paramsc                 C   s   | j S N)param)requestr   r   r   chunked1      r   c                 C   s
   t | dS )Nr   r   )Zstring_storager   r   r   r   6   s   
r   c                 C   sj   t jjttjdd}|d |d kr&t jjttjdd}|d |d ks|  |}|r3t|S |S )Nd   )sizer      )	nprandomchoiceliststringascii_lettersconstruct_array_type_from_sequencer   )r   r   stringsr   r   r   r   data;   s   r-   c                 C   s$   |   tjdg}|rt|S |S )zLength 2 array with [NA, Valid]Ar*   r+   pdNAr   r   r   r   r   r   r   data_missingE   s   r3   c                 C   s"   |   g d}|rt|S |S )N)BCr.   )r*   r+   r   r2   r   r   r   data_for_sortingL   s   r6   c                 C   s&   |   dtjdg}|rt|S |S )Nr4   r.   r/   r2   r   r   r   data_missing_for_sortingR   s   r7   c                   C   s   t jS r   )r0   r1   r   r   r   r   na_valueX   r   r8   c              
   C   s2   |   ddtjtjddddg}|rt|S |S )Nr4   r.   r5   r/   r2   r   r   r   data_for_grouping]   s   r9   c                       s$   e Zd Z fddZdd Z  ZS )	TestDtypec                    s&   |d|j  dksJ t | d S )Nzstring[])r   supertest_eq_with_strselfr   	__class__r   r   r=   f   s   zTestDtype.test_eq_with_strc                 C   s   t |sJ d S r   r   r>   r   r   r   test_is_not_string_typej   s   z!TestDtype.test_is_not_string_type)__name__
__module____qualname__r=   rB   __classcell__r   r   r@   r   r:   e   s    r:   c                          e Zd Z fddZ  ZS )TestInterfacec                    (   |j jdkrtjdd t | d S Nr	   /2D support not implemented for ArrowStringArrayreason)r   r   r   r   r<   	test_viewr?   r-   r   r@   r   r   rN   q      zTestInterface.test_view)rC   rD   rE   rN   rF   r   r   r@   r   rH   p       rH   c                   @      e Zd Zdd Zdd ZdS )TestConstructorsc                 C   s   d S r   r   )r?   r-   r   r   r   test_from_dtypex   s   z TestConstructors.test_from_dtypec                 C   sH   t jddd tjdgtddd}t|jtsJ |jjdks"J d S )Nr	   z1.0.0)Z
minversionEr    r   )r   Zimportorskipr0   Seriesr   
isinstancer   r   )r?   resultr   r   r   test_constructor_from_list|   s   z+TestConstructors.test_constructor_from_listN)rC   rD   rE   rT   rZ   r   r   r   r   rS   w   s    rS   c                       rG   )TestReshapingc                    rI   rJ   )r   r   r   r   r<   test_transposerO   r@   r   r   r\      rP   zTestReshaping.test_transpose)rC   rD   rE   r\   rF   r   r   r@   r   r[      rQ   r[   c                   @      e Zd ZdS )TestGetitemNrC   rD   rE   r   r   r   r   r^          r^   c                       rG   )TestSetitemc                    rI   rJ   )r   r   r   r   r<   test_setitem_preserves_viewsrO   r@   r   r   rb      rP   z(TestSetitem.test_setitem_preserves_views)rC   rD   rE   rb   rF   r   r   r@   r   ra      rQ   ra   c                   @   r]   )	TestIndexNr_   r   r   r   r   rc      r`   rc   c                       s,   e Zd Zdd Zdd Z fddZ  ZS )TestMissingc                 C   s"   |  }|dg }| || d S )Nr#   )Zdropnaassert_extension_array_equal)r?   r3   rY   expectedr   r   r   test_dropna_array   s   
zTestMissing.test_dropna_arrayc                 C   s   ||    }|d }||}||usJ | || tt|jjdk |jdd}W d    n1 s7w   Y  ||usBJ | || d S )Nr   r	   Zbackfill)method)isnaZfillnare   tmmaybe_produces_warningr   r   r   )r?   r-   ZvalidrY   r   r   r   test_fillna_no_op_returns_copy   s   
z*TestMissing.test_fillna_no_op_returns_copyc                    sT   t jt|d uo|jjdkdd t || W d    d S 1 s#w   Y  d S )Nr	   F)Zcheck_stacklevel)rj   rk   r   r   r   r<   test_fillna_series_method)r?   r3   Zfillna_methodr@   r   r   rm      s   "z%TestMissing.test_fillna_series_method)rC   rD   rE   rg   rl   rm   rF   r   r   r@   r   rd      s    rd   c                   @   s&   e Zd Zejdddgdd ZdS )TestNoReduceskipnaTFc                 C   sZ   |}|dv rd S t |}tt t|||d W d    d S 1 s&w   Y  d S )N)minmax)ro   )r0   rW   r   raises	TypeErrorgetattr)r?   r-   Zall_numeric_reductionsro   op_nameserr   r   r   test_reduce_series_numeric   s   
"z'TestNoReduce.test_reduce_series_numericN)rC   rD   rE   r   markparametrizerw   r   r   r   r   rn      s    rn   c                   @   s   e Zd Zdd ZdS )TestMethodsc                 C   s   |d d   }t||   }tj||jd}|jdd }tjdt	| gt	| |j
dd}t|jdd	d
krC|d}n|d}| || d S )N
   rV   T)	normalizer#   Z
proportionindexnamer    r	   zdouble[pyarrow]ZFloat64)uniquer$   arrayri   r0   rW   r   Zvalue_countsZ
sort_indexr   r~   rt   astypeassert_series_equal)r?   r-   valuesrv   rY   rf   r   r   r    test_value_counts_with_normalize   s   
z,TestMethods.test_value_counts_with_normalizeN)rC   rD   rE   r   r   r   r   r   rz      s    rz   c                   @   r]   )TestCastingNr_   r   r   r   r   r      r`   r   c                   @   rR   )TestComparisonOpsc           	      C   sZ   d|j  d}t|||}|jjdkrdnd}t|t|||}| || d S )N__r	   zboolean[pyarrow]boolean)rC   rt   r   r   r   objectr   )	r?   rv   r-   opotherru   rY   r   rf   r   r   r   _compare_other   s
   z TestComparisonOps._compare_otherc                 C   s   t |}| |||d d S )Nabc)r0   rW   r   )r?   r-   Zcomparison_oprv   r   r   r   test_compare_scalar   s   
z%TestComparisonOps.test_compare_scalarN)rC   rD   rE   r   r   r   r   r   r   r      s    r   c                   @   r]   )TestParsingNr_   r   r   r   r   r      r`   r   c                   @   r]   )TestPrintingNr_   r   r   r   r   r      r`   r   c                       sB   e Zd Zejdddgdd Zejd fddZ  Z	S )	TestGroupByas_indexTFc           	      C   s   t g d|d}|jd|dj }t j|dd\}}|r8t j|dd}t jg d|d	d
}| || d S t |g dd}| 	|| d S )N)r#   r#   r
   r
      r   r#      )r.   r4   r4   )r   T)sort)r   )g      @g      ?g      @r.   r}   )r4   r.   )
r0   Z	DataFramegroupbyr.   meanZ	factorizeIndexrW   r   Zassert_frame_equal)	r?   r   r9   ZdfrY   _Zuniquesr~   rf   r   r   r   test_groupby_extension_agg   s   z&TestGroupBy.test_groupby_extension_aggz4ignore:Falling back:pandas.errors.PerformanceWarningc                    s   t  || d S r   )r<   test_groupby_extension_apply)r?   r9   Zgroupby_apply_opr@   r   r   r      s   z(TestGroupBy.test_groupby_extension_apply)
rC   rD   rE   r   rx   ry   r   filterwarningsr   rF   r   r   r@   r   r      s
    

r   c                   @   s    e Zd Zejdddd ZdS )Test2DCompatT)Zautousec                 C   s   t |trtjdd d S d S )NrK   rL   )rX   r   r   r   rO   r   r   r   arrow_not_supported  s   
z Test2DCompat.arrow_not_supportedN)rC   rD   rE   r   fixturer   r   r   r   r   r     s    
r   c                 C   sr   | \}}}|  g d}tj|d< |rt|}d}tjt|d || W d    d S 1 s2w   Y  d S )N)r
   r   r#   zOsearchsorted requires array to be sorted, which is impossible with NAs present.)match)Ztaker0   r1   rW   r   rr   
ValueErrorZsearchsorted)r6   Z	as_seriesbcar   msgr   r   r    test_searchsorted_with_na_raises  s   


"r   )?__doc__r(   numpyr$   r   Zpandas.errorsr   Zpandasr0   Zpandas._testingZ_testingrj   Zpandas.api.typesr   Zpandas.core.arraysr   Zpandas.core.arrays.string_r   Zpandas.tests.extensionr   r   r   r   r   r-   r3   r6   r7   r8   r9   ZBaseDtypeTestsr:   ZBaseInterfaceTestsrH   ZBaseConstructorsTestsrS   ZBaseReshapingTestsr[   ZBaseGetitemTestsr^   ZBaseSetitemTestsra   ZBaseIndexTestsrc   ZBaseMissingTestsrd   ZBaseNoReduceTestsrn   ZBaseMethodsTestsrz   ZBaseCastingTestsr   ZBaseComparisonOpsTestsr   ZBaseParsingTestsr   ZBasePrintingTestsr   ZBaseGroupbyTestsr   ZDim2CompatTestsr   r   r   r   r   r   <module>   sZ    


	




