U
    ¾mœdlb  ã                   @   s‚   d dl 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	m
Z
mZmZmZmZ d dlmZ d dlmZmZ G dd„ dƒZdS )é    N)Úis_integer_dtype)ÚCategoricalÚCategoricalIndexÚ	DataFrameÚ
RangeIndexÚSeriesÚget_dummies)ÚSparseArrayÚSparseDtypec                
   @   s(  e Zd Zejdd„ ƒZejddejedgd�dd„ ƒZ	ejd	d
gd�dd„ ƒZ
dd„ Zdd„ Zdd„ Zdd„ Zdd„ Zdd„ Zdd„ Zdd„ Zdd„ Zdd „ Zd!d"„ Zd#d$„ Zd%d&„ Zd'd(„ Zd)d*„ Zd+d,„ Zd-d.„ Zd/d0„ Zd1d2„ Zej d3d4e d5d6giƒie d7d8giƒfd4e d9d5giƒie d:d8giƒfe d9d6giƒd5d;œe d7d8giƒfe d9d6giƒd5d<œe d=d8giƒfg¡d>d?„ ƒZ!d@dA„ Z"dBdC„ Z#dDdE„ Z$dFdG„ Z%dHdI„ Z&dJdK„ Z'dLdM„ Z(dNdO„ Z)ej dPd8dQg¡dRdS„ ƒZ*ej d
d8dQg¡dTdU„ ƒZ+dVdW„ Z,dXdY„ Z-ej dZd[g¡d\d]„ ƒZ.d^d_„ Z/d`da„ Z0dS )bÚTestGetDummiesc                 C   s"   t dddgdddgdddgdœƒS )NÚaÚbÚcé   é   é   )ÚAÚBÚC)r   )Úself© r   ú^/home/sam/Atlas/atlas_env/lib/python3.8/site-packages/pandas/tests/reshape/test_get_dummies.pyÚdf   s    zTestGetDummies.dfÚuint8Úi8N)Úparamsc                 C   s   t  |j¡S ©N)ÚnpÚdtypeÚparam©r   Úrequestr   r   r   r      s    zTestGetDummies.dtypeZdenseÚsparsec                 C   s
   |j dkS )Nr"   )r   r    r   r   r   r"   "   s    zTestGetDummies.sparsec                 C   s   |d krt jS |S r   )r   r   )r   r   r   r   r   Úeffective_dtype(   s    zTestGetDummies.effective_dtypec              	   C   s.   d}t jt|d�� t|dd� W 5 Q R X d S )Nz1dtype=object is not a valid dtype for get_dummies©ÚmatchÚobject©r   )ÚpytestÚraisesÚ
ValueErrorr   )r   r   Úmsgr   r   r   Ú'test_get_dummies_raises_on_dtype_object-   s    z6TestGetDummies.test_get_dummies_raises_on_dtype_objectc                 C   s¸   t dƒ}t|ƒ}t|t dƒƒ}tdddgdddgdddgdœ|  |¡d�}|r\|jtdd�}t|||d	�}t ||¡ t|||d	�}t ||¡ t dƒ|_	t|||d	�}t ||¡ d S )
NÚabcÚABCr   r   ©r   r   r   r'   ç        ©Ú
fill_value©r"   r   )
Úlistr   r   r#   Úapplyr	   r   ÚtmÚassert_frame_equalÚindex)r   r"   r   Ús_listÚs_seriesÚs_series_indexÚexpectedÚresultr   r   r   Útest_get_dummies_basic2   s     þ
z%TestGetDummies.test_get_dummies_basicc                 C   sÂ  t dƒ}t|ƒ}tdddddgdddddgdddddgd	œƒ}tdddgdddgdddgd	œ|  |¡t dƒd
�}|r¢t|ƒr‚d}n|tkr�d}nd}|jt|d�}t|||d�}t	 
||¡ t|||d�}t	 
||¡ t||j||d�}|�r
d|  |¡j› d|› d�}	n|  |¡j}	t|	didd�}|j ¡ }dd„ |jD ƒ|_t	 ||¡ t|dg||d�}dddœ}
d|
 |	d¡ |
|	< t|
dd� ¡ }|j ¡ }dd„ |jD ƒ|_| ¡ }t	 ||¡ d S )Nr-   r   r   r   r   r   r   r   r/   )r   ÚcolumnsFr0   r1   r3   )r?   r"   r   zSparse[z, ú]é   Úcount©Únamec                 S   s   g | ]}t |ƒ‘qS r   ©Ústr©Ú.0Úir   r   r   Ú
<listcomp>k   s     z?TestGetDummies.test_get_dummies_basic_types.<locals>.<listcomp>r   )Úint64r&   c                 S   s   g | ]}t |ƒ‘qS r   rE   rG   r   r   r   rJ   u   s     )r4   r   r   r#   r   Úboolr5   r	   r   r6   r7   r?   rD   ZdtypesZvalue_countsr8   Zassert_series_equalÚgetÚ
sort_index)r   r"   r   r9   r:   Zs_dfr<   r2   r=   Z
dtype_nameZexpected_countsr   r   r   Útest_get_dummies_basic_typesG   sJ    (ÿý


z+TestGetDummies.test_get_dummies_basic_typesc                 C   s    t jg}t|ƒ}t|dgd�}t||d�}t||d�}t||d�}|jsLt‚|jsVt‚|js`t‚|j ¡ dgkstt‚|j ¡ dgksˆt‚|j ¡ dgksœt‚d S )Nr   ©r8   ©r"   r   )r   Únanr   r   ÚemptyÚAssertionErrorr8   Útolist)r   r"   Zjust_na_listZjust_na_seriesZjust_na_series_indexZres_listZ
res_seriesZres_series_indexr   r   r   Útest_get_dummies_just_nay   s    


z'TestGetDummies.test_get_dummies_just_nac           
   	   C   s$  ddt jg}t|||d�}tdddgdddgdœ|  |¡d�}|rP|jtdd	�}t ||¡ t|d
||d�}tt jdddgddddgddddgi|  |¡d�}|j	ddt jgdd�}|j
|_
|rÎ|jtdd	�}t ||¡ tt jgd
||d�}ttddgd�t jg|  |¡d�}	t |j|	j¡ d S )Nr   r   r3   r   r   )r   r   r'   r0   r1   T©Údummy_nar"   r   ©ZaxisrP   ©r?   r   )r   rR   r   r   r#   r5   r	   r6   r7   Úreindexr?   r   Zassert_numpy_array_equalÚvalues)
r   r"   r   ÚsÚresÚexpÚres_naÚexp_naÚres_just_naÚexp_just_nar   r   r   Útest_get_dummies_include_naŠ   s4     ÿ"þ  ÿz*TestGetDummies.test_get_dummies_include_nac                 C   sj   d}t  d¡}|||g}t|d|d�}tddddgd|› �dddgiƒ}|rZ|jtd	d
�}t ||¡ d S )NÚezLATIN SMALL LETTER E WITH ACUTEÚletter©Úprefixr"   Zletter_eTFZletter_r   r1   )ÚunicodedataÚlookupr   r   r5   r	   r6   r7   )r   r"   re   Úeacuter]   r^   r_   r   r   r   Útest_get_dummies_unicode§   s    

ÿz'TestGetDummies.test_get_dummies_unicodec                 C   s¤   |ddg }t ||d�}tdddgdddgdddgdddgdœtd�}|r”ttdddgdd�tdddgdd�tdddgdd�tdddgdd�dœƒ}t ||¡ d S )	Nr   r   rQ   r   r   ©ÚA_aÚA_bÚB_bÚB_cr'   rL   )r   r   rL   r	   r6   r7   ©r   r   r"   r=   r<   r   r   r   Útest_dataframe_dummies_all_obj´   s    $þüÿ	z-TestGetDummies.test_dataframe_dummies_all_objc                 C   sb   |ddg }|  dddœ¡}t|ƒ}tdddgdddgdddgdddgdœtd	�}t ||¡ d S )
Nr   r   r&   Ústring©r   r   r   r   rm   r'   )Úastyper   r   rL   r6   r7   ©r   r   r=   r<   r   r   r   Ú#test_dataframe_dummies_string_dtypeÇ   s    üù	z2TestGetDummies.test_dataframe_dummies_string_dtypec              	   C   s    t |||d�}|r"t}t|dƒ}n
tj}|}tdddg|dddg|d�|dddg|d�|dddg|d�|dddg|d�dœƒ}|dd	d
ddg }t ||¡ d S )Nr3   r   r   r   r   r'   )r   rn   ro   rp   rq   r   rn   ro   rp   rq   )r   r	   r
   r   Úarrayr   r6   r7   ©r   r   r"   r   r=   ZarrÚtypr<   r   r   r   Ú"test_dataframe_dummies_mix_default×   s     ûÿ	z1TestGetDummies.test_dataframe_dummies_mix_defaultc                    s¨   ddg}t |||d�}tdddgdddgdddgdddgdddgd	œƒ}|d
g |d
g< ddddg}|d
g|  }|rztnt‰ ||  ‡ fdd„¡||< t ||¡ d S )NÚfrom_AÚfrom_Brg   r   r   r   TF©r   Úfrom_A_aÚfrom_A_bÚfrom_B_bÚfrom_B_cr   r€   r�   r‚   rƒ   c                    s   ˆ | ƒS r   r   )Úx©r{   r   r   Ú<lambda>ü   ó    zCTestGetDummies.test_dataframe_dummies_prefix_list.<locals>.<lambda>)r   r   r	   r   r5   r6   r7   )r   r   r"   Úprefixesr=   r<   Úcolsr   r…   r   Ú"test_dataframe_dummies_prefix_listë   s     ûÿ	z1TestGetDummies.test_dataframe_dummies_prefix_listc              
   C   sÚ   t |d|d�}ddddg}tdddddgd	ddddgd
ddddggdg| d�}| dtji¡}|rÊtjtdd	d
gdd�tdddgddd�tdddgddd�tdddgddd�tdddgddd�gdd�}t 	||¡ d S )NÚbadrg   Zbad_aZbad_bZbad_cr   TFr   r   r   ©r?   rC   zSparse[bool])rD   r   rY   )
r   r   rv   r   rK   ÚpdÚconcatr   r6   r7   )r   r   r"   r=   Zbad_columnsr<   r   r   r   Ú!test_dataframe_dummies_prefix_strÿ   s*    ýúûøz0TestGetDummies.test_dataframe_dummies_prefix_strc                 C   sª   t |dgdg|d�}tdddgdddgdd	dgd	dd	gd
œƒ}|j}||dd …   t¡||dd … < |dg |dg< |ršddg}||  tdd	ƒ¡||< t ||¡ d S )Nr}   r   )rh   r?   r"   r   r   r   r   r   r   )r   r   r€   r�   r   r€   r�   rL   )r   r   r?   rv   rL   r
   r6   r7   ©r   r   r"   r=   r<   r‰   r   r   r   Útest_dataframe_dummies_subset  s    üÿ"z,TestGetDummies.test_dataframe_dummies_subsetc                 C   sê   t |d|d�}tdddgdddgdddgdddgdddgdœƒ}|d	g |d	g< |d	d
dddg }|rŠd
dddg}||  tddƒ¡||< t ||¡ t |ddg|d�}|jdddœd�}t ||¡ t |dddœ|d�}t ||¡ d S )Nz..©Ú
prefix_sepr"   r   r   r   TF)r   úA..aúA..búB..búB..cr   r”   r•   r–   r—   rL   r   Ú__ZB__bZB__c)r–   r—   rŒ   ru   )r   r   rv   r
   r6   r7   Úrenamer�   r   r   r   Ú!test_dataframe_dummies_prefix_sep.  s(    ûÿ	z0TestGetDummies.test_dataframe_dummies_prefix_sepc              	   C   s8   t  d¡}tjt|d�� t|dg|d� W 5 Q R X d S )NzPLength of 'prefix' (1) did not match the length of the columns being encoded (2)r$   ztoo fewrg   ©ÚreÚescaper(   r)   r*   r   ©r   r   r"   r+   r   r   r   Ú(test_dataframe_dummies_prefix_bad_lengthH  s
    ÿz7TestGetDummies.test_dataframe_dummies_prefix_bad_lengthc              	   C   s8   t  d¡}tjt|d�� t|dg|d� W 5 Q R X d S )NzTLength of 'prefix_sep' (1) did not match the length of the columns being encoded (2)r$   r‹   r’   r›   rž   r   r   r   Ú,test_dataframe_dummies_prefix_sep_bad_lengthP  s
    ÿz;TestGetDummies.test_dataframe_dummies_prefix_sep_bad_lengthc                 C   s¶   dddœ}t dddgdddgddd	gd
œƒ}t|||d�}t dddgdddgdddgdddgdddgdœƒ}ddddg}||  t¡||< |r¦||  tddƒ¡||< t ||¡ d S )Nr}   r~   ru   r   r   r   r   r   r   )r   r   r   rg   r   r   r€   r�   r‚   rƒ   rL   )r   r   rv   rL   r
   r6   r7   )r   r"   rˆ   r   r=   r<   r?   r   r   r   Ú"test_dataframe_dummies_prefix_dictX  s     
"ûÿ
z1TestGetDummies.test_dataframe_dummies_prefix_dictc                 C   s  t jt jt jg|jdd d …f< t|d||d�jdd�}|rJt}t|dƒ}n
t j}|}tdddt jg|ddddg|d�|ddddg|d�|ddddg|d�|ddddg|d�|ddddg|d�|ddddg|d�d	œƒjdd�}t	 
||¡ t|d
||d�}|dddddg }t	 
||¡ d S )Nr   TrW   r   rY   r   r   r'   )r   rn   ro   ÚA_nanrp   rq   ÚB_nanFr   rn   ro   rp   rq   )r   rR   Úlocr   rN   r	   r
   ry   r   r6   r7   rz   r   r   r   Útest_dataframe_dummies_with_nan  s2    ÿùÿ
öz-TestGetDummies.test_dataframe_dummies_with_nac                 C   sÐ   t dddgƒ|d< t|||d�jdd�}|r<t}t|dƒ}n
tj}|}tddd	g|dddg|d
�|dddg|d
�|dddg|d
�|dddg|d
�|dddg|d
�|dddg|d
�dœƒjdd�}t 	||¡ d S )Nr„   ÚyÚcatr3   r   rY   r   r   r   r'   )r   rn   ro   rp   rq   Zcat_xÚcat_y)
r   r   rN   r	   r
   r   ry   r   r6   r7   rz   r   r   r   Ú'test_dataframe_dummies_with_categorical�  s(    ùÿ
öz6TestGetDummies.test_dataframe_dummies_with_categoricalzget_dummies_kwargs,expectedÚdataõ   Ã¤r   u   Ã¤_aTr„   u   x_Ã¤)rª   rh   )rª   r“   u   xÃ¤ac                 C   s   t f |Ž}t ||¡ d S r   )r   r6   r7   )r   Zget_dummies_kwargsr<   r=   r   r   r   Útest_dataframe_dummies_unicode¥  s    
z-TestGetDummies.test_dataframe_dummies_unicodec                 C   sª   t dƒ}t|ƒ}t|t dƒƒ}tdddgdddgdœtd�}t|d|d�}|r\|jtdd	�}t ||¡ t|d|d�}t ||¡ t dƒ|_	t|d|d�}t ||¡ d S )
Nr-   r.   r   r   )r   r   r'   T©Ú
drop_firstr"   r1   )
r4   r   r   rL   r   r5   r	   r6   r7   r8   ©r   r"   r9   r:   r;   r<   r=   r   r   r   Ú!test_get_dummies_basic_drop_firstÀ  s    
z0TestGetDummies.test_get_dummies_basic_drop_firstc                 C   sŒ   t dƒ}t|ƒ}t|t dƒƒ}ttdƒd�}t|d|d�}t ||¡ t|d|d�}t ||¡ tt dƒd�}t|d|d�}t ||¡ d S )NZaaar.   r   rP   Tr­   )r4   r   r   r   r   r6   r7   r¯   r   r   r   Ú+test_get_dummies_basic_drop_first_one_levelÕ  s    z:TestGetDummies.test_get_dummies_basic_drop_first_one_levelc           	      C   sà   ddt jg}t|d|d�}tddddgitd�}|rB|jtdd�}t ||¡ t|dd|d	�}tddddgt jdddgitd�j	dt jgdd
�}|r¢|jtdd�}t ||¡ tt jgdd|d	�}tt
dƒd�}t ||¡ d S )Nr   r   Tr­   r   r   r'   r1   ©rX   r®   r"   rY   rP   )r   rR   r   r   rL   r5   r	   r6   r7   r[   r   )	r   r"   Zs_NAr^   r_   r`   ra   rb   rc   r   r   r   Ú$test_get_dummies_basic_drop_first_NAç  s,    " ÿ   ÿz3TestGetDummies.test_get_dummies_basic_drop_first_NAc                 C   sZ   |ddg }t |d|d�}tdddgdddgdœtd�}|rJ|jtdd	�}t ||¡ d S )
Nr   r   Tr­   r   r   )ro   rq   r'   r1   )r   r   rL   r5   r	   r6   r7   rr   r   r   r   Ú!test_dataframe_dummies_drop_firstÿ  s    z0TestGetDummies.test_dataframe_dummies_drop_firstc                 C   s¤   t dddgƒ|d< t|d|d�}tdddgd	dd	gd	d	dgd	ddgd
œƒ}dddg}||  t¡||< |ddddg }|r”|D ]}t|| ƒ||< q~t ||¡ d S )Nr„   r¦   r§   Tr­   r   r   r   r   )r   ro   rq   r¨   ro   rq   r¨   r   )r   r   r   rv   rL   r	   r6   r7   )r   r   r"   r   r=   r<   r‰   Úcolr   r   r   Ú2test_dataframe_dummies_drop_first_with_categorical  s    $ÿ
zATestGetDummies.test_dataframe_dummies_drop_first_with_categoricalc              	   C   sö   t jt jt jg|jdd d …f< t|dd|d�jdd�}tdddt jgddddgddddgddddgddddgdœƒ}d	d
ddg}||  t¡||< |jdd�}|r¼|D ]}t|| ƒ||< q¦t	 
||¡ t|dd|d�}|dd	dg }t	 
||¡ d S )Nr   Tr²   r   rY   r   r   )r   ro   r¢   rq   r£   ro   r¢   rq   r£   Fr   )r   rR   r¤   r   rN   r   rv   rL   r	   r6   r7   )r   r   r"   r=   r<   r‰   rµ   r   r   r   Ú)test_dataframe_dummies_drop_first_with_na  s6       ÿþ



ûÿ	z8TestGetDummies.test_dataframe_dummies_drop_first_with_nac                 C   s˜   t dddgƒ}t|ƒ}tddgddgddggddgtd�}t ||¡ t tdddgƒƒ}t|ƒ}tddgddgddggtddgƒtd�}t ||¡ d S )Nr   r   r   rZ   r   r   )r   r   r   rL   r6   r7   r   )r   rª   r=   r<   r   r   r   Útest_get_dummies_int_int0  s    $ 
 ÿz'TestGetDummies.test_get_dummies_int_intc              	   C   s¸   t dddgtdddgƒdddgdddgdœƒ}dd	d
dddg}t ddddddgddddddgddddddgg|d�}||dd …   |¡||dd … < t|ddg|d�}t ||¡ d S )Nr   r   r   r   g      ð?g       @)r   r   r   ÚDr   r¹   ÚA_1ÚA_2ZB_arp   r   rŒ   r   r   rZ   )r   r   rv   r   r6   r7   )r   r   rª   r?   r<   r=   r   r   r   Útest_get_dummies_int_df=  s    üÿ,þ"z&TestGetDummies.test_get_dummies_int_dfÚorderedFc                 C   s|   t tdƒtdƒ|d�}t||d�}tjdddgdddgg|  |¡d�}t|j|j|d�}t|||  |¡d�}t	 
||¡ d S )NZxyZxyz)Ú
categoriesr½   r'   r   r   rZ   )r   r4   r   r   ry   r#   r   r¾   r   r6   r7   )r   r   r½   r§   r=   rª   r‰   r<   r   r   r   Ú1test_dataframe_dummies_preserve_categorical_dtypeO  s    $  ÿz@TestGetDummies.test_dataframe_dummies_preserve_categorical_dtypec                 C   sL   t  ddgddgdœ¡}t|dg|d�}|jdgd	�}t |dg |¡ d S )
Nr   r   ZABZCD)ÚGDPÚNationrÁ   ©r?   r"   rÀ   rŒ   )r   Ú	from_dictr   r[   r6   r7   )r   r"   r   Zdf2r   r   r   Ú*test_get_dummies_dont_sparsify_all_columns]  s    z9TestGetDummies.test_get_dummies_dont_sparsify_all_columnsc                 C   s~   dddg|_ t|ƒjdd�}tdddddgdddddgdddddggddd	d	d
gd�jdd�}| dtji¡}t ||¡ d S )Nr   r   rY   TFr   r   rn   ro   ZA_crŒ   )	r?   r   rN   r   rv   r   rK   r6   r7   rw   r   r   r   Ú"test_get_dummies_duplicate_columnsf  s    ýúù	z1TestGetDummies.test_get_dummies_duplicate_columnsc                 C   s`   t dddgiƒ}t|dgdd�}tddƒ}t tddg|d�tddg|d�d	œƒ}t ||¡ d S )
Nr   r   r   TrÂ   rL   r   r'   )rº   r»   )r   r   r
   r	   r6   r7   )r   r   r=   r   r<   r   r   r   Útest_get_dummies_all_sparsex  s    
þÿz*TestGetDummies.test_get_dummies_all_sparser\   Úbazc              
   C   sp   t ddddddgddddddgd	d
dd	d
dgddddddgdœƒ}d}tjt|d�� t||d� W 5 Q R X d S )Nr   r   r   é   é   é   ÚoneÚtwor   r   r   r„   r¦   ÚzÚqÚwÚt)ÚbarZfoorÇ   Zzooz1Input must be a list-like for parameter `columns`r$   rŒ   )r   r(   r)   Ú	TypeErrorr   )r   r\   r   r+   r   r   r   Ú#test_get_dummies_with_string_values„  s    üÿ	z2TestGetDummies.test_get_dummies_with_string_valuesc                 C   sT   t tdƒƒ}t||d�}tddddgddddgddddgdœ|d�}t ||¡ d S )NÚabcar'   r   r   r/   )r   r4   r   r   r6   r7   )r   Úany_numeric_ea_and_arrow_dtypeZserr=   r<   r   r   r   Ú test_get_dummies_ea_dtype_series•  s    "þz/TestGetDummies.test_get_dummies_ea_dtype_seriesc                 C   sX   t dtdƒiƒ}t||d�}t ddddgddddgddddgdœ|d�}t ||¡ d S )Nr„   rÔ   r'   r   r   )Zx_aZx_bZx_c)r   r4   r   r6   r7   )r   rÕ   r   r=   r<   r   r   r   Ú#test_get_dummies_ea_dtype_dataframeŸ  s    "þz2TestGetDummies.test_get_dummies_ea_dtype_dataframe)1Ú__name__Ú
__module__Ú__qualname__r(   Zfixturer   r   Zfloat64rL   r   r"   r#   r,   r>   rO   rV   rd   rl   rs   rx   r|   rŠ   r�   r‘   rš   rŸ   r    r¡   r¥   r©   ÚmarkZparametrizer   r¬   r°   r±   r³   r´   r¶   r·   r¸   r¼   r¿   rÄ   rÅ   rÆ   rÓ   rÖ   r×   r   r   r   r   r      sx   


2þþþþóþ




r   )rœ   ri   Únumpyr   r(   Zpandas.core.dtypes.commonr   Zpandasr�   r   r   r   r   r   r   Zpandas._testingZ_testingr6   Zpandas.core.arrays.sparser	   r
   r   r   r   r   r   Ú<module>   s    