U
    ,d9                     @   s   d dl Zd dlZd dlZd dlmZmZ d dlm	Z	 d dl
mZmZmZ e	 Zde_dd Zdd	 Zd
d Zdd Zdd Zdd Zdd ZG dd deZG dd deeZG dd deeZdS )    N)jittyped)Flags)TestCaseCompilationCacheMemoryLeakMixinTc                 C   s
   t | S N)hqheapify)x r   :/tmp/pip-unpacked-wheel-eu7e0c37/numba/tests/test_heapq.pyr
      s    r
   c                 C   s
   t | S r   )r	   heappop)heapr   r   r   r      s    r   c                 C   s   t | |S r   )r	   heappushr   itemr   r   r   r      s    r   c                 C   s   t | |S r   )r	   heappushpopr   r   r   r   r      s    r   c                 C   s   t | |S r   )r	   heapreplacer   r   r   r   r      s    r   c                 C   s   t | |S r   )r	   	nsmallestniterabler   r   r   r   "   s    r   c                 C   s   t | |S r   )r	   nlargestr   r   r   r   r   &   s    r   c                       s   e Zd Z f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%d& Zd'd( Zd)d* Zd+d, Zd-d. Zd/d0 Zd1d2 Zd3d4 Z  ZS )5
_TestHeapqc                    s(   t t|   t | _tjd| _d S )N*   )	superr   setUpr   ZccachenprandomZRandomStaternd)self	__class__r   r   r   ,   s    z_TestHeapq.setUpc              
   C   s.  t }tdd|}ddddddd	d
ddg
}| |}|| || | |t| dddtjtj tjg}t	|d
D ]6}t|}| |}|| || | |t| qxt
t|D ]8}|| g}| |}|| || | |t| qdddg}| |}|| || | |t| d S )NTZnopython            	               r   A`"	@g      $g      @r&   !   )r%      )r*      )r
   r   listimplassertPreciseEquallistr   naninf	itertoolscombinations_with_replacementrangelen)r!   pyfunccfuncabZelement_poolr   ir   r   r   test_heapify_basic_sanity1   s2    





z$_TestHeapq.test_heapify_basic_sanityc                 C   s8   t |D ]*\}}|r|d d? }| || |k qd S )Nr%   )	enumerate
assertTrue)r!   r   posr   	parentposr   r   r   check_invariantX   s    z_TestHeapq.check_invariantc                 C   s   t }tdd|}t}tdd|}| dg}| dg}| | tdD ]4}| jdd}|	| ||| | | qNg }	|r||}| | |		| q|d d  }
|

  | t|
|	 | |	 d S )NTr$            r%   r   )r   r   r   r3   rF   r:   r    Zrandnr   appendsortr4   r5   )r!   Zpyfunc_heappushcfunc_heappushZpyfunc_heappopcfunc_heappopr   datar@   r   resultsZdata_sortedr   r   r   test_push_pop^   s*    



z_TestHeapq.test_push_popc                 C   sT   t }tdd|}ttdddg D ](}| | j|}|| | | q&d S )NTr$   r%      i N  )r
   r   r5   r:   r3   r    Zrandom_samplerF   )r!   r<   r=   sizer   r   r   r   test_heapifyy   s    z_TestHeapq.test_heapifyc              	   C   s   t }tdd|}|   |  }|d W 5 Q R X d}| |t|j |  }|| ddg W 5 Q R X d}| |t|j d S )NTr$   r%   r'   r+   heap argument must be a listy      ?      ?y       @      z>'<' not supported between instances of 'complex' and 'complex')r
   r   disable_leak_checkassertTypingErrorassertInstr	exceptionr3   r!   r<   r=   emsgr   r   r   test_heapify_exceptions   s    

z"_TestHeapq.test_heapify_exceptionsc           	      C   sz   t }tdd|}dd }| D ]T}t| | |}tt|D ]0}||}||}| |t| | || qBq d S )NTr$   c                
   s   sT   ddddddddd	d
g
V  dddgV  t jdt jd V  t ddd V  d S )Nr%   r&   r'   r(   r)   r*   r+   r,   r-   r   r/   )r%   o   )r*   i  )Z
fill_valued   )r   fullr6   tolistlinspacer   r   r   r   a_variations   s    z:_TestHeapq.test_heappop_basic_sanity.<locals>.a_variations)r   r   r
   r3   r:   r;   r4   r5   )	r!   r<   r=   re   r>   r?   r@   Zval_pyZval_cr   r   r   test_heappop_basic_sanity   s    

z$_TestHeapq.test_heappop_basic_sanityc              	   C   sP   t }tdd|}|   |  }|d W 5 Q R X d}| |t|j d S )NTr$   rS   rT   )r   r   rU   rV   rW   rX   rY   rZ   r   r   r   test_heappop_exceptions   s    
z"_TestHeapq.test_heappop_exceptionsc                 c   sl   |  ddddddddd	d
g
V  tddd}|  |V  |  |d d d V  | j| |  |V  d S )Nr%   r&   r'   r(   r)   r*   r+   r,   r-   r   r_      )r3   r   rd   r    shuffle)r!   r>   r   r   r   	iterables   s     z_TestHeapq.iterablesc                    s   t }tdd|}t}tdd| |  D ]Z}t|}| |dg|D ]}|| qN fddttD }| 	|| q,d S )NTr$   r   c                    s   g | ]} qS r   r   .0_Z	cfunc_popr   r   r   
<listcomp>   s     z2_TestHeapq.test_heappush_basic.<locals>.<listcomp>)
r   r   r   rk   sortedr3   popr:   r;   r4   )r!   Zpyfunc_pushZ
cfunc_pushZ
pyfunc_popr   expectedvaluegotr   ro   r   test_heappush_basic   s    z_TestHeapq.test_heappush_basicc              	   C   s   t }tdd|}|   |  }|dd W 5 Q R X d}| |t|j |  }|| dddgd	 W 5 Q R X d
}| |t|j d S )NTr$   rS   r,   rT   r%   r'   r+   g      @'heap type must be the same as item type)r   r   rU   rV   rW   rX   rY   r3   rZ   r   r   r   test_heappush_exceptions   s    

 z#_TestHeapq.test_heappush_exceptionsc                 C   s   t }tdd|}|  D ]<}tdt|d D ]$}|d|}|d|}| || q0q|d| dddg}| |g  |d| dddg}| |dg |dd}| |ddg |dtd	}| |d
ddg d S )NTr$   r`   r&   r%   Fr*   r,   r'   r+   r&   r*   r%   r,   r   )	r   r   rk   r:   r;   r4   r3   r   aranger!   r<   r=   r   r   rs   ru   outr   r   r   test_nsmallest_basic   s    


z_TestHeapq.test_nsmallest_basicc                 C   s   t }tdd|}|  D ]<}tdt|d D ]$}|d|}|d|}| || q0q|d| dddg}| |g  |d| dddg}| |dg |dd}| |d	d
g |dtd	}| |d
ddg d S )NTr$   r`   r&   r%   Fr*   ry   r,   r'   r+   )	r   r   rk   r:   r;   r4   r3   r   rz   r{   r   r   r   test_nlargest_basic   s    


z_TestHeapq.test_nlargest_basicc              	   C   s   |    |  }|d| dddg W 5 Q R X d}| |t|j |  }|dd W 5 Q R X d}| |t|j d S )Ng@r&   r*   r%   z%First argument 'n' must be an integerra   z+Second argument 'iterable' must be iterable)rU   rV   r3   rW   rX   rY   )r!   r=   r[   r\   r   r   r   _assert_typing_error  s    
 
z_TestHeapq._assert_typing_errorc                 C   s    t }tdd|}| | d S NTr$   )r   r   r   r!   r<   r=   r   r   r   test_nsmallest_exceptions&  s    z$_TestHeapq.test_nsmallest_exceptionsc                 C   s    t }tdd|}| | d S r   )r   r   r   r   r   r   r   test_nlargest_exceptions+  s    z#_TestHeapq.test_nlargest_exceptionsc              
   C   s   t }tdd|}ddddddd	d
ddg
}t| | |}dD ](}||| ||| | |t| q@tddd}tj|d	< tj	|d< |
 }t| | |}ddtj	 tj	fD ](}||| ||| | |t| qd S )NTr$   r%   r&   r'   r(   r)   r*   r+   r,   r-   r   )r+            ri   g      r.   )r   r   r
   r3   r4   r5   r   rd   r6   r7   rc   )r!   r<   r=   r>   r?   r   r   r   r   test_heapreplace_basic0  s&    







z!_TestHeapq.test_heapreplace_basicc              	   C   s   t }tdd|}|   |  }|dd W 5 Q R X d}| |t|j |  }|| dddgd	 W 5 Q R X d
}| |t|j d S )NTr$   rS   ri   rT   r%   r'   r+   rG   rw   )r   r   rU   rV   rW   rX   rY   r3   rZ   r   r   r   test_heapreplace_exceptionsK  s    

 z&_TestHeapq.test_heapreplace_exceptionsc                 c   s*   zt |V  qW n tk
r$   Y nX d S r   )r   
IndexError)r!   r   r   r   r   heapiter^  s    z_TestHeapq.heapiterc                 C   s   t ddt}t ddt}| jtdd }| |d d }|| |dd  D ]}||d krX||| qX| t	| 
t	|t|dd   d S )NTr$       
   r   r_   )r   r
   r   r    choicer:   rc   r3   r4   r5   r   rq   )r!   cfunc_heapifyZcfunc_heapreplacerM   r   r   r   r   r   
test_nbeste  s    z_TestHeapq.test_nbestc                    s   t ddt}t ddt}t ddt tdD ]}tjdtjd}| | j	
|d}|d@ rx|d d  | n,| |d g|dd  D ]}|| q fd	d
tdD }| |t| q2d S )NTr$   ra   r'   Zdtyper   r%   r   c                    s   g | ]} qS r   r   rl   rL   r   r   r   rp     s     z,_TestHeapq.test_heapsort.<locals>.<listcomp>)r   r
   r   r   r:   r   rz   float64r3   r    r   r4   rq   )r!   r   rK   ZtrialvaluesrM   r   Zheap_sortedr   r   r   test_heapsortv  s    
z_TestHeapq.test_heapsortc                 C   sZ   t }tdd|}| | jtdd}dD ]&}| t|||t|d |  q.d S )NTr$   r   r   	r   r%   r*   r   ra   i  i  r   iL  )	r   r   r3   r    r   r:   r4   r5   rq   r!   r<   r=   rM   r   r   r   r   test_nsmallest  s
    z_TestHeapq.test_nsmallestc                 C   s^   t }tdd|}| | jtdd}dD ]*}| t|||t|ddd |  q.d S )NTr$   r   r   r   )reverse)	r   r   r3   r    r   r:   r4   r5   rq   r   r   r   r   test_nlargest  s    z_TestHeapq.test_nlargestc           	      C   s   t }tdd|}t}tdd|}tjdtjd}| | j|d}|d d }|| |dd  D ]}||| qh| 	t
| t
|t|dd   d S )NTr$   r   r   r   r   r_   )r   r   r
   r   rz   r   r3   r    r   r4   r5   r   rq   )	r!   Zpyfunc_heappushpopZcfunc_heappushpopZpyfunc_heapifyr   r   rM   r   r   r   r   r   test_nbest_with_pushpop  s    z"_TestHeapq.test_nbest_with_pushpopc                 C   s   t }tdd|}| dg}||d}| t||fdgdf | t|d t | t|t | dg}||d}| t||fdgdf | dg}||d}| t||fdgdf d S )	NTr$   g      ?g      $@r   r   r)   r1   )r   r   r3   r4   r5   typefloat)r!   r<   r=   hr   r   r   r   test_heappushpop  s    


z_TestHeapq.test_heappushpopc              	   C   s   t }tdd|}|   |  }|dd W 5 Q R X d}| |t|j |  }|| dddgd	 W 5 Q R X d
}| |t|j d S )NTr$   rS   ri   rT   r%   r'   r+   Frw   )r   r   rU   rV   rW   rX   rY   r3   rZ   r   r   r   test_heappushpop_exceptions  s    

 z&_TestHeapq.test_heappushpop_exceptions)__name__
__module____qualname__r   rA   rF   rO   rR   r]   rf   rg   rk   rv   rx   r}   r~   r   r   r   r   r   r   r   r   r   r   r   r   r   __classcell__r   r   r"   r   r   *   s4   'r   c                   @   s   e Zd ZdZeZdS )TestHeapqReflectedListzTest heapq with reflected listsN)r   r   r   __doc__r5   r3   r   r   r   r   r     s   r   c                   @   s   e Zd ZdZejZdS )TestHeapqTypedListzTest heapq with typed listsN)r   r   r   r   r   Listr3   r   r   r   r   r     s   r   )heapqr	   r8   Znumpyr   Znumbar   r   Znumba.core.compilerr   Znumba.tests.supportr   r   r   Zno_pyobj_flagsZnrtr
   r   r   r   r   r   r   r   r   r   r   r   r   r   <module>   s(      9