a
    GÿGbÊ8  ã                   @   s`   d dl Zd dlZd dlZd dlmZmZmZ d dlm	Z
 g d¢ZG dd„ dƒZG dd„ dƒZdS )é    N)Ú	DataFrameÚ
MultiIndexÚSeries)ÚsumZprodÚminÚmaxZmedianÚmeanZskewZmadÚstdÚvarZsemc                	   @   s|  e Zd Zdd„ Zej 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ej de¡ej dddg¡ej dddg¡ej dddg¡dd„ ƒƒƒƒZej de¡ej dddg¡ej dddg¡ej dddg¡ej dddg¡d d!„ ƒƒƒƒƒZej d"d#d$g¡d%d&„ ƒZd'd(„ Zd)d*„ Zd+d,„ Zd-d.„ Zd/d0„ Zej d1d2d1g¡d3d4„ ƒZd5d6„ Zd7d8„ Zd9S ):ÚTestMultiLevelc                 C   s  |}t  t¡ |jdd}W d   ƒ n1 s00    Y  |j|jdd}|jdd tj¡}t  	||¡ |d j|jdd}|d jdd tj¡}t j
||dd t  t¡  |jjddd}W d   ƒ n1 sÖ0    Y  |j|jdd}|jdd tj¡j}t  	||¡ d S )	NÚmonth©Úlevelé   ÚAF©Zcheck_names©Úaxisr   )Úcolumnsr   )ÚtmÚassert_produces_warningÚFutureWarningr   ÚreindexÚindexÚgroupbyÚ	transformÚnpÚassert_frame_equalÚassert_series_equalÚT)ÚselfÚ/multiindex_year_month_day_dataframe_random_dataÚymdÚ
month_sumsÚresultÚexpected© r&   ú;lib/python3.9/site-packages/pandas/tests/test_multilevel.pyÚtest_reindex_level   s    *.z!TestMultiLevel.test_reindex_levelÚopname)ÚsubÚaddÚmulZdivc           	      C   sØ   |}t t|ƒ}t t¡ |jdd}W d   ƒ n1 s:0    Y  |||dd}|jdd tj¡}|||ƒ}t 	||¡ t t
|ƒ}||d |d dd}|d jdd tj¡}||d |ƒ}d|_t ||¡ d S )Nr   r   r   )Úgetattrr   r   r   r   r   r   r   r   r   r   Únamer   )	r    r)   r!   r"   Úopr#   r$   Zbroadcastedr%   r&   r&   r'   Útest_binops_level3   s    
*

z TestMultiLevel.test_binops_levelc                 C   s0   |}|j ddg }|jddg }t ||¡ d S )Nr   é   ©ÚfooÚone©Úbarr4   )ÚilocÚlocr   r   )r    Ú multiindex_dataframe_random_dataÚframer%   Z	reindexedr&   r&   r'   Útest_reindexJ   s    zTestMultiLevel.test_reindexc                 C   sŒ   |}|j d d d… }| |¡}|j |u s,J ‚|j| }|j  |¡sFJ ‚|j}|j|d}|j|u sfJ ‚|jd d …|f }|j |¡sˆJ ‚d S )Né
   ©r   )r   r   r8   Zequalsr   r   )r    r!   r"   Z	new_indexÚchunkZymdTr&   r&   r'   Útest_reindex_preserve_levelsQ   s    

z+TestMultiLevel.test_reindex_preserve_levelsc           	      C   s^   |}|d }|j  d¡}| |¡}| dd„ ¡}| dd„ ¡}| |j ¡}tj||dd d S )Nr   r   c                 S   s   | d S ©Né   r&   ©Úxr&   r&   r'   Ú<lambda>l   ó    z7TestMultiLevel.test_groupby_transform.<locals>.<lambda>c                 S   s   | d S r@   r&   rB   r&   r&   r'   rD   m   rE   Fr   )r   Úget_level_valuesr   Zapplyr   r   r   r   )	r    r9   r:   ÚsZgrouperÚgroupedZappliedr%   r$   r&   r&   r'   Útest_groupby_transformd   s    
z%TestMultiLevel.test_groupby_transformc                 C   sV   t dgdgdggdgdgdggg d¢d}ttj d¡gg d¢|d	}|jd
d d S )Nr3   r6   Úbazr   )r4   ÚtwoÚthree©ÚlevelsÚcodesÚnamesé   ©ÚaÚbÚcÚd)r   r   rL   r   )r   r   r   ÚrandomZrandr   )r    ÚmidxÚdfr&   r&   r'   Útest_groupby_cornerq   s    ýz"TestMultiLevel.test_groupby_cornerc                 C   sr   t  g d¢¡}tg d¢g d¢g|d}|jdd|j dd„ ¡ }|jdd	d
}| ¡ }|jddgk ¡ snJ ‚d S )N))Úf1Ús1)r[   Ús2)Úf2r\   )r^   r]   )Úf3r\   )r_   r]   )r   rA   r1   rQ   é   é   )é   é   é	   r<   é   é   r=   r   )r   c                 S   s   | d dv S )Nr   )r^   r_   r&   )Úur&   r&   r'   rD   ˆ   rE   z:TestMultiLevel.test_groupby_level_no_obs.<locals>.<lambda>r   r   r^   r_   )	r   Úfrom_tuplesr   r8   r   Úmapr   r   Úall)r    rX   rY   Zdf1rH   r$   r&   r&   r'   Útest_groupby_level_no_obs{   s    ÿ
z(TestMultiLevel.test_groupby_level_no_obsc                 C   sJ   |}|d d… j }|d |d< t|jtƒs.J ‚|d |d k ¡ sFJ ‚d S )Nr`   )éÐ  r   rb   )rl   r   r<   )r   Ú
isinstancer   r   rj   )r    r!   r"   rY   r&   r&   r'   Ú.test_setitem_with_expansion_multiindex_columnsŽ   s
    z=TestMultiLevel.test_setitem_with_expansion_multiindex_columnsc                 C   s¸   t g d¢t g d¢¡d}t g d¢t g d¢¡d}|| }|j |j¡}| |¡| |¡ }t ||¡ |d d d… |d d d…  }|j |j¡}| |¡| |¡ }t ||¡ d S )N)r   rA   r1   ))r   r   )r   rA   ©ÚBr1   )Údatar   )rQ   r`   ra   ))ÚZr   )rr   rA   ro   éÿÿÿÿ)r   r   rh   r   Úunionr   r   r   )r    rC   ÚyÚresZ	exp_indexZexpr&   r&   r'   Útest_alignment˜   s    ÿÿzTestMultiLevel.test_alignmentr/   r   r   r   ÚskipnaTFÚsortc           
         s‚   |}|j ||d}| ‡ ‡fdd„¡}t t¡" t|ˆ ƒ|ˆd}	W d   ƒ n1 sX0    Y  |rr|	j|d}	t ||	¡ d S )N)r   ry   c                    s   t | ˆ ƒˆdS )N)rx   ©r-   rB   ©r/   rx   r&   r'   rD   ¸   rE   z:TestMultiLevel.test_series_group_min_max.<locals>.<lambda>)r   rx   r   )r   Úaggr   r   r   r-   Ú
sort_indexr   )
r    r/   r   rx   ry   Zseries_with_multilevel_indexZserrH   ÚleftsideÚ	rightsider&   r{   r'   Útest_series_group_min_max¬   s    0z(TestMultiLevel.test_series_group_min_maxr   c                    s$  |}t j|jdddgf< t j|jdddgf< |jj| }ˆ dkrF|}n|j}|j|ˆ |d}	g ‰‡ ‡‡‡fdd„}
|	 |
¡}t 	t
¡$ t|ˆƒ|ˆ ˆd}W d   ƒ n1 s°0    Y  |rÚ|j|ˆ d	}|j|ˆ d	}| ˆ ¡j|  |¡}t | ˆ ¡|¡ t | ˆ ¡|¡ t ||¡ d S )
Nr   rA   rb   r   )r   r   ry   c                    s   ˆ  | ¡ t| ˆƒˆˆ dS )N)rx   r   )Úappendr-   rB   ©r   r/   Úpiecesrx   r&   r'   ÚaggfØ   s    
z1TestMultiLevel.test_frame_group_ops.<locals>.aggf)r   r   rx   )r   r   )r   Únanr7   r   rP   r   r   r|   r   r   r   r-   r}   Z	_get_axisrN   ÚrenameÚassert_index_equalr   )r    r/   r   r   rx   ry   r9   r:   Z
level_namerH   r„   r~   r   Zlevel_indexr&   r‚   r'   Útest_frame_group_ops¿   s(    	
2z#TestMultiLevel.test_frame_group_opsÚmethr
   r	   c                    s  t  t d¡ d¡t t d¡d¡g¡}ttj t	|ƒd¡|d}d‰ ‡ ‡fdd„}t
 t¡& t|d ˆƒdˆ d}W d   ƒ n1 sŠ0    Y  |d jdd	 |¡}t
 ||¡ t
 t¡" t|ˆƒdˆ d}W d   ƒ n1 sè0    Y  |jdd	 |¡}t
 ||¡ d S )
Nr`   r<   ©r   rQ   c                    s   t | ˆƒˆ dS )N)Úddofrz   rB   ©r‹   r‰   r&   r'   rD   ó   rE   z7TestMultiLevel.test_std_var_pass_ddof.<locals>.<lambda>r   )r   r‹   r   )r   Zfrom_arraysr   ZarangeÚrepeatZtiler   rW   ÚrandnÚlenr   r   r   r-   r   r|   r   r   )r    r‰   r   rY   Zaltr$   r%   r&   rŒ   r'   Útest_std_var_pass_ddofë   s     ÿ40z%TestMultiLevel.test_std_var_pass_ddofc                 C   sn   |}t  ||¡}t  t¡  |jddgd}W d   ƒ n1 s@0    Y  |jddgd ¡ }t  ||¡ d S )NZyearr   r   )r   Zget_objr   r   r   r   Zassert_equal)r    r!   Zframe_or_seriesr"   r$   r%   r&   r&   r'   Útest_agg_multiple_levelsÿ   s    .z'TestMultiLevel.test_agg_multiple_levelsc                 C   sš   |}|j ddgd ¡ }|j d¡}|j d¡}|  ||g¡ ¡ }tj||dd |jj|jjd d… ksnJ ‚|j |jjd d… d ¡ }t ||¡ d S )Nr   r   r   Fr   rA   )r   r   r   rF   r   r   rP   )r    r!   r"   r$   Zk1Zk2r%   Úresult2r&   r&   r'   Útest_groupby_multilevel
  s    z&TestMultiLevel.test_groupby_multilevelc                 C   s@   t  g d¢¡}ttj dd¡||d}| d¡|d< | ¡ }d S )N)r2   )r3   rK   r5   )r6   rK   rQ   )r   r   r   )ZTotalsÚ )r   rh   r   r   rW   rŽ   r   Z_consolidate)r    r   rY   r&   r&   r'   Útest_multilevel_consolidate  s    ÿz*TestMultiLevel.test_multilevel_consolidatec                 C   sê  t g d¢ddggg d¢g d¢gd}ttj d¡|d}ttj dd	¡|d}|d
 }|jd
 }|d d… }|j d¡|_t	 
||¡ t	 
||¡ tjtdd |d  W d   ƒ n1 s¾0    Y  |jd
 }| d
¡}|d d… }|j d¡|_t	 ||¡ t	 ||¡ t g d¢ddggg d¢g d¢gd}ttj d¡|d}ttj dd	¡|d}|d }|jd }|d d… }|j d¡|_t	 
||¡ t	 
||¡ |jd }| d¡}|d d… }|j d¡|_t	 ||¡ t	 ||¡ d S )N)©r3   r6   r   )r3   rJ   r   )r3   Úquxr   r   r   )r   r   r   r   rA   rA   )r   r   r   r   r   r   )rN   rO   ra   rŠ   rQ   r–   rA   z^\(\('foo', 'bar', 0\), 2\)$)Úmatch)r–   rA   )©r3   r6   )r3   rJ   )r3   r—   r™   )r   r   r   rW   rŽ   r   r8   r   Z	droplevelr   r   ÚpytestZraisesÚKeyErrorZxsr   )r    r   Zseriesr:   r$   r’   r%   r&   r&   r'   Útest_level_with_tuples#  sL    þ
&

þ


z%TestMultiLevel.test_level_with_tuplesc                 C   sº   |}|j ddgdd}|jg d¢ }t ||¡ |jj ddgddd}t ||j¡ |jddg }t ||¡ |d jddg }t ||d ¡ |jjd d …ddgf }t ||j¡ d S )	Nr3   r—   r   r   )r   r   rA   rb   rc   rd   r   r   r   )r   r7   r   r   r   r8   r   )r    r9   r:   r$   r%   r&   r&   r'   Ú$test_reindex_level_partial_selectionS  s    z3TestMultiLevel.test_reindex_level_partial_selectionrV   rQ   c                 C   st   g d¢}t dgdgdg|gdœƒ}||jdk  |¡}| ¡ j}tdgdgdggg g g gg d¢d}t ||¡ d S )N)rS   rT   rU   r   rA   r1   rR   rM   )r   rS   r   Úfirstr   r   r   r‡   )r    rV   Z
group_keysrY   Úgr$   r%   r&   r&   r'   Ú+test_empty_frame_groupby_dtypes_consistencyf  s    
ÿz:TestMultiLevel.test_empty_frame_groupby_dtypes_consistencyc                 C   sR   g d¢}g d¢}t j|ddgd}t||d}| |j¡ ¡ }t|ƒdksNJ ‚d S )N))Ú600809Z20061231©r¡   Z20070331)r¡   Z20070630r¢   )Údemor£   r£   r£   ZSTK_IDZRPT_Date)rP   rŠ   r1   )r   rh   r   r   r   rž   r   )r    Zidx_tpZdtÚidxrG   r$   r&   r&   r'   Útest_duplicate_groupby_issuest  s    z,TestMultiLevel.test_duplicate_groupby_issuesc                 C   sP   ddgg}ddt jfdg}t|t |¡d}|jjj}|jjj}t 	||¡ d S )NrC   r   rS   rT   )rS   rU   g        r=   )
r   r…   r   r   rh   ZdtypesrS   rT   r   r   )r    rq   r   rY   r%   r$   r&   r&   r'   Útest_subsets_multiindex_dtypeƒ  s    


z,TestMultiLevel.test_subsets_multiindex_dtypeN)Ú__name__Ú
__module__Ú__qualname__r(   rš   ZmarkZparametrizer0   r;   r?   rI   rZ   rk   rn   rw   ÚAGG_FUNCTIONSr€   rˆ   r   r‘   r“   r•   rœ   r   r    r¥   r¦   r&   r&   r&   r'   r      s@   


'
0
r   c                   @   s   e Zd ZdZdd„ ZdS )Ú
TestSortedz+everything you wanted to test about sortingc                 C   s¸   t g d¢g d¢gg d¢g d¢gƒ}tdtt|ƒƒi|dd}|jjdu sLJ ‚| ¡ }|jjd	u sdJ ‚tdg d
¢it  g d¢¡dd}|jt	j
dd…dd…f d d …f }t ||¡ d S )N)r   rp   ÚC)rU   rT   rS   )r   r   rA   r   r   rA   )r   rA   r   r   r   rA   ÚcolZint64)r   ZdtypeFT)r   rQ   r`   rA   ))rp   rS   )rp   rU   )r¬   rS   )r¬   rT   rp   r¬   rS   rU   )r   r   Úranger   r   Zis_monotonicr}   rh   r8   ÚpdZ
IndexSlicer   r   )r    r¤   rY   Úsortedr%   r$   r&   r&   r'   Útest_sort_non_lexsorted  s     ÿ
ÿû$z"TestSorted.test_sort_non_lexsortedN)r§   r¨   r©   Ú__doc__r±   r&   r&   r&   r'   r«     s   r«   )Znumpyr   rš   Zpandasr¯   r   r   r   Zpandas._testingZ_testingr   rª   r   r«   r&   r&   r&   r'   Ú<module>   s     t