U
    +ifNB                     @   s.  d dl Zd dlZd dlmZ d dlZd dlmZmZ d dl	m
Z d0d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ejdddgdd Zdd  Zejd!ddgd"d# Zd$d% Zd&d' Z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S )1    N)DatetimeTZDtype)CategoricalIndex
MultiIndexFc                 C   sP   t | t |kstt| |D ],\}}t|}t|}tj|||d qd S )Ncheck_dtype)lenAssertionErrorzipnpZasarraytmZassert_numpy_array_equal)actualexpectedr   Zactexp r   Z/home/mars/bis/venv/lib/python3.8/site-packages/pandas/tests/indexes/multi/test_get_set.pyassert_matching   s
    

r   c              	   C   s~   ddg| _ | ddkst| ddks.td}tjt|d | d W 5 Q R X tjtdd | d W 5 Q R X d S )N   r   z/Too many levels: Index has only 2 levels, not 3match   zLevel fourth not foundZfourth)names_get_level_numberr   pytestraises
IndexErrorKeyError)idxmsgr   r   r   test_get_level_number_integer   s    
r   c                  C   sh   t jdddgdddgtjdddd	gd
ddgd} ttdtdtddd}t	|| j
 d S )Nr   r      abc20200101UTCZperiodstzintstringdtr   int64Outcr&   )r'   r(   r)   r   from_productpd
date_rangeSeriesr
   dtyper   r   assert_series_equaldtypesZidx_multityper   r   r   r   test_get_dtypes#   s     r8   c               	   C   s^   t dddgdddgtjdddd	g} ttd
tdtddd}t	|| j
 d S )Nr   r   r   r    r!   r"   r#   r$   r%   r+   r,   r-   r.   )Zlevel_0Zlevel_1Zlevel_2r/   r7   r   r   r   test_get_dtypes_no_level_name3   s    r9   c                  C   sp   t jdddgdddgtjdddd	gd
d
d
gdj} tjtdtdtddgd
d
d
gd}t	
| | d S )Nr   r   r   r    r!   r"   r#   r$   r%   Ar*   r+   r,   r-   r.   )index)r   r0   r1   r2   r6   r3   r
   r4   r   r   r5   )resultr   r   r   r   %test_get_dtypes_duplicate_level_namesF   s    r=   c              	   C   sT   | }t jtdd |jd W 5 Q R X t jtdd |jd W 5 Q R X d S )NzToo many levelsr   r   znot a valid level number)r   r   r   r;   r   )Z multiindex_dataframe_random_dataframer   r   r   #test_get_level_number_out_of_boundsW   s
    r@   c              	   C   s  | j | jkstdd |D }| |}| j|ks6t|j|ksDtd}tjt|d |||  W 5 Q R X dd |D }|j|dd}|d kst|j|kst| j|d dd	}| j|kst|j|d |d
 gkst|j|d ddd}|d kst|j|d |d
 gkst| j|dd
gd	}| j|ks@t|j|ksPt|j|dd
gdd}|d ksrt|j|kstd S )Nc                 S   s   g | ]}|d  qS )ZSUFFIXr   .0namer   r   r   
<listcomp>c   s     z)test_set_name_methods.<locals>.<listcomp>z9Length of names must match number of levels in MultiIndexr   c                 S   s   g | ]}|d  qS )ZSUFFIX2r   rA   r   r   r   rD   j   s     Tinplacer   levelr   rH   rF   )rename	set_namesr   r   r   r   
ValueError)r   Zindex_namesZ	new_namesindr   Z
new_names2resr   r   r   test_set_name_methods`   s0    
rO   c              	   C   s   | j }dd |D }| j}|\}}dd |D }dd |D }||g}d}tjt|d || _ W 5 Q R X tjt|d || _W 5 Q R X d S )Nc                 S   s   g | ]}d d |D qS )c                 S   s   g | ]}|d  qS r    r   rB   Zlevr   r   r   rD      s     z=test_set_levels_codes_directly.<locals>.<listcomp>.<listcomp>r   rB   rH   r   r   r   rD      s     z2test_set_levels_codes_directly.<locals>.<listcomp>c                 S   s   g | ]}|d  d qS r   r   r   rB   xr   r   r   rD      s     c                 S   s   g | ]}|d  d  qS r   r   rT   r   r   r   rD      s     z[Cc]an't set attributer   )levelscodesr   r   AttributeError)r   rW   
new_levelsrX   major_codesminor_codes	new_codesr   r   r   r   test_set_levels_codes_directly   s    r^   c                 C   s  | j }dd |D }| |}t|j | t| j | |  }tt |j|dd}W 5 Q R X |d ksntt|j | | j|d dd}t|j |d |d g t| j | | j|d dd}t|j |d |d g t| j | | j|ddgd}t|j | t| j | |  }tt |j|d ddd}W 5 Q R X |d ksPtt|j |d |d g t| j | |  }tt |j|d ddd}W 5 Q R X |d kstt|j |d |d g t| j | |  }tt |j|ddgdd}W 5 Q R X |d kstt|j | t| j | |  }d	D ]\}tj	t
d
d. tt | jdgd|d W 5 Q R X W 5 Q R X t| j |j dd tj	t
d
d8 tt" | jddddddgd|d W 5 Q R X W 5 Q R X t| j|jdd tj	tdd, tt | jdd|d W 5 Q R X W 5 Q R X t| j |j dd tj	tdd, tt | jdd|d W 5 Q R X W 5 Q R X t| j|jdd q<d S )Nc                 S   s   g | ]}d d |D qS )c                 S   s   g | ]}|d  qS rP   r   rQ   r   r   r   rD      s     z.test_set_levels.<locals>.<listcomp>.<listcomp>r   rR   r   r   r   rD      s     z#test_set_levels.<locals>.<listcomp>TrE   r   rG   r   rI   )TFz^Onr   r"   r   r   r         z^Levelsz^Codes)rW   
set_levelsr   copyr   assert_produces_warningFutureWarningr   r   r   rL   	set_codesrX   	TypeError)r   rW   rZ   ind2inplace_returnZoriginal_indexrF   r   r   r   test_set_levels   sn    

&0$$ri   c           
   	   C   s  | j }|\}}dd |D }dd |D }||g}| |}t|j | t| j | |  }tt |j|dd}W 5 Q R X |d kstt|j | | j|d dd}t|j |d |d g t| j | | j|d dd}t|j |d |d g t| j | | j|ddgd}t|j | t| j | |  }tt |j|d ddd	}W 5 Q R X |d ksntt|j |d |d g t| j | |  }tt |j|d ddd	}W 5 Q R X |d kstt|j |d |d g t| j | |  }tt |j|ddgdd	}W 5 Q R X |d ks6tt|j | t| j | t	d
d t
dD }t
ddd}t	dd |D }|j|dd}	|	|st| }	tt |	j|ddd W 5 Q R X |	|std S )Nc                 S   s   g | ]}|d  d qS rS   r   rT   r   r   r   rD      s     z"test_set_codes.<locals>.<listcomp>c                 S   s   g | ]}|d  d  qS rV   r   rT   r   r   r   rD      s     TrE   r   rG   r   rI   c                 S   s   g | ]}d |fqS r   r   rB   ir   r   r   rD      s           c                 S   s   g | ]}d |fqS rj   r   rk   r   r   r   rD   "  s     )rX   rH   rX   rH   rF   )rX   re   r   rb   r   rc   rd   r   r   from_tuplesrangeequals)
r   rX   r[   r\   r]   rg   rh   rM   r   r<   r   r   r   test_set_codes   sb    
rt   c              	   C   s  | j | j }}| j}tjtdd | |d g W 5 Q R X tjtdd | |d g W 5 Q R X tjtdd | |d g W 5 Q R X tjt	dd | |d  W 5 Q R X tjt	dd | |d  W 5 Q R X tjt	dd | |d  W 5 Q R X tjt	dd | j|d ddgd	 W 5 Q R X tjt	dd | j|dd	 W 5 Q R X tjt	dd | j|d ddgd	 W 5 Q R X tjt	dd | j|dd	 W 5 Q R X tjtdd | j|d ddgd	 W 5 Q R X tjt	d
d | j|dd	 W 5 Q R X d S )NzLength of levelsr   r   zLength of codeszLength of nameszlist of lists-likez	list-liker   rG   zNames must be a)
rW   rX   r   r   r   rL   ra   re   rK   rf   )r   rW   rX   r   r   r   r   %test_set_levels_codes_names_bad_input/  s4       ru   rF   Tc                 C   sT   t ddggddggdgd}t ddgg}|jdd| d}| rD|}t|| d S )Nr   r   firstrW   rX   r   rI   )r   r0   rK   r   assert_index_equal)rF   r   mr<   r   r   r   test_set_names_with_nlevel_1^  s    rz   c               	   C   s   t ddgddgg} d}tjt|d | ddgd }W 5 Q R X t ddgddggd	d	d
d
gd	d
d	d
ggddgd}t|| d S )NpythonZcobrai  i  zxIn a future version of pandas all arguments of MultiIndex.set_names except for the argument 'names' will be keyword-onlyr   kindyearr   r   rw   )r   r0   r   rc   rd   rK   rx   r   r   r<   r   r   r   r   )test_multi_set_names_pos_args_deprecationm  s    r   orderedc                 C   s   t tdddddgg}ttd| d}|j|dd}t |ddddgg|jd	}t|| |d}ttd
|j	|j
d}t|| d S )NZxyzxr   r   r   r   Zbac)r   rG   )rW   rX   Zbacb)
categoriesr   )r   from_arrayslistr   ra   rX   r   rx   Zget_level_valuesr   r   )r   r;   Zcidxr<   r   Z
result_lvlZexpected_lvlr   r   r   test_set_levels_categorical~  s    
  r   c                  C   s   ddddddg} dddgd }t j| |gddgd	}tjtjd
dddddg|d}| }|jd kslt	|j
jdks|t	d|jd< |jd kst	|j
jdkst	d S )NZhansgrethe123r   NameNumberr*      r_   onetwothreefour)columnsr;   )r   r   g(\X@))r   4r   )r   r   r1   Z	DataFramer
   randomZrandnZ
sort_indexZ_is_copyr   r;   r   at)Zlev1Zlev2r   Zdfr   r   r   test_set_value_keeps_names  s     
 
r   c                  C   st   dddg} dgd }t j| |gddgd}|jttdd	d
gdd}dddg}t j||gddgd}t|| d S )Nr   r   r   blacksizecolorr*   r   r   r   rG   )r   r   ra   mapr'   r   rx   )sizescolorsr;   r<   Zexpected_sizesr   r   r   r   test_set_levels_with_iterable  s    


r   c              	   C   s>   | j d d d d }tt | j|d|d W 5 Q R X d S )Nr   ro   rp   )rX   r   rc   rd   re   )r   rF   r]   r   r   r   !test_set_codes_inplace_deprecated  s    r   c              	   C   s8   | j d  }tt | j|d|d W 5 Q R X d S )Nr   )rW   rH   rF   )rW   rb   r   rc   rd   ra   )r   rF   Z	new_levelr   r   r   "test_set_levels_inplace_deprecated  s    r   c               	   C   sp   t jdddgddgd} d}tjt|d | d	d
dgd}W 5 Q R X t jdddgddgd}t|| d S )N)r   r   )r   r   )r   r   foobarr*   zzIn a future version of pandas all arguments of MultiIndex.set_levels except for the argument 'levels' will be keyword-onlyr   r    r!   r"   r   )r    r   )r!   r   )r"   r   )r   rq   r   rc   rd   ra   rx   r~   r   r   r   $test_set_levels_pos_args_deprecation  s&    	r   c              
   C   sx   d}t jt|d2 | ddddddgddddddggddg}W 5 Q R X tjddd	d
ddgddgd}t || d S )NzxIn a future version of pandas all arguments of MultiIndex.set_codes except for the argument 'codes' will be keyword-onlyr   r   r   r   r   )r   r   )r   r   )r   r   )Zbazr   )quxr   )r   r   rv   secondr*   )r   rc   rd   re   r   rq   rx   r~   r   r   r   $test_set_codes_pos_args_depreciation  s    6r   )F)"numpyr
   r   Zpandas.core.dtypes.dtypesr   Zpandasr1   r   r   Zpandas._testingZ_testingr   r   r   r8   r9   r=   r@   rO   r^   ri   rt   ru   markZparametrizerz   r   r   r   r   r   r   r   r   r   r   r   r   <module>   s:   

	"PJ/



