o
    i9F                     @   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	 d dl
mZ d dlmZmZ dd Zdd	 Zd
d Zdd ZG dd deeZG dd deeZedkrVe  dS dS )    N)jittypeofnjit)types)TypingError)MemoryLeakMixinTestCasec                 C      | | S N )abr   r   n/var/www/html/eduruby.in/lip-sync/lip-sync-env/lib/python3.10/site-packages/numba/tests/test_fancy_indexing.pygetitem_usecase   s   r   c                 C      || |< d S r
   r   )r   idxr   r   r   r   setitem_usecase      r   c                 C   s   t | |S r
   npZtake)Aindicesr   r   r   np_take   r   r   c                 C   s   t j| ||dS )Naxisr   )r   r   r   r   r   r   np_take_kws   s   r   c                   @   s   e Zd Zd#ddZd#ddZd#ddZ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"S )$TestFancyIndexingTc                 C   sD   t d|d dgg}|r |t dd|d gt g dg7 }|S )Nr      )r   r   r   r   )r   Zint16Zuint16Zbool_)selfNmanychoicesr   r   r   generate_advanced_indices   s   z+TestFancyIndexing.generate_advanced_indicesc              	   c   s    |r,t dddt d|d dt dddt |d ddt | d ddt d| dg}nt d|d dt d| dg}t|d D ]}tj||dD ]}|V  qJqAdS )zG
        Generate basic index tuples with 0 to *maxdim* items.
        Nr   r      r   )repeat)slicerange	itertoolsproduct)r   r    maxdimr!   r"   ndimtupr   r   r   generate_basic_index_tuples!   s"   

z-TestFancyIndexing.generate_basic_index_tuplesc                 c   sl    t | j||d}t|d D ]#}| ||d |D ]}|D ]}|d| |f ||d  V  qqqdS )z
        Generate advanced index tuples by generating basic index tuples
        and adding a single advanced index item.
        r!   r   N)listr#   r(   r.   )r   r    r+   r!   r"   ir-   Zadvr   r   r   generate_advanced_index_tuples8   s   "z0TestFancyIndexing.generate_advanced_index_tuplesc                 c   sP    |  |||D ]}tt|d D ]}|d| tf ||d  V  qqdS )zr
        Same as generate_advanced_index_tuples(), but also insert an
        ellipsis at various points.
        r   N)r2   r(   lenEllipsis)r   r    r+   r!   r-   r1   r   r   r   ,generate_advanced_index_tuples_with_ellipsisE   s   "z>TestFancyIndexing.generate_advanced_index_tuples_with_ellipsisc           
      C   s   t }tdd|}| }|jp|}|D ]9}|||}|j|us"J |||}	| |	j|j | |	j|j tj	|	| |	j
rM|	d tj	|| qd S )NTZnopython*   )r   r   copybaseassertEqualshapedtyper   testingassert_equalsizefill)
r   arrr   pyfunccfuncorig	orig_baseindexexpectedgotr   r   r   check_getitem_indicesN   s    



z'TestFancyIndexing.check_getitem_indicesc                 C   F   d}d}t || |f| t j}| ||}| || d S N      )r   arangereshapeastypeint32r2   rI   r   r    r,   rA   r   r   r   r   test_getitem_tupled   s
   "z$TestFancyIndexing.test_getitem_tuplec                 C   J   d}d}t || |f| t j}| j||dd}| || d S NrL   rM   Fr/   )r   rN   rO   rP   rQ   r5   rI   rR   r   r   r   test_getitem_tuple_and_ellipsism      "z1TestFancyIndexing.test_getitem_tuple_and_ellipsisc                 C   s>   t dddd }td}||d | |d |d  d S )NTr6   c                 S   s   | d | d< d S )N).r   ).r   r   )rA   vr   r   r   fooy   s   z7TestFancyIndexing.test_ellipsis_getsetitem.<locals>.foor$   r   r   )r   r   rN   r:   )r   rY   rA   r   r   r   test_ellipsis_getsetitemw   s
   


z*TestFancyIndexing.test_ellipsis_getsetitemc                 C   sD   d}d}t || |f| t j}| |}| || d S rK   )r   rN   rO   rP   rQ   r#   rI   rR   r   r   r   test_getitem_array   s
   "
z$TestFancyIndexing.test_getitem_arrayc           	      C   s   t }tdd|}|D ]3}|| }t|}t|}|||| |||| | |j|j | |j|j tj|| qd S )NTr6   )	r   r   r   
zeros_liker:   r;   r<   r=   r>   )	r   rA   r   rB   rC   rF   srcrG   rH   r   r   r   check_setitem_indices   s   

z'TestFancyIndexing.check_setitem_indicesc                 C   rJ   rK   )r   rN   rO   rP   rQ   r2   r^   rR   r   r   r   test_setitem_tuple   s
   "z$TestFancyIndexing.test_setitem_tuplec                 C   rT   rU   )r   rN   rO   rP   rQ   r5   r^   rR   r   r   r   test_setitem_tuple_and_ellipsis   rW   z1TestFancyIndexing.test_setitem_tuple_and_ellipsisc                 C   sH   d}d}t || |f| t jd }| |}| || d S )NrL   rM   
   )r   rN   rO   rP   rQ   r#   r^   rR   r   r   r   test_setitem_array   s
   &
z$TestFancyIndexing.test_setitem_arrayc                 C   s  t }tdd|}tdtdftdtdftjdtjdtjdtjdftjdtjdtjdtjdftjdd	dtd
ftjdd	dtdftjg dd	dtjdddftjdtdtjdtdfg}|D ]\}}| }||d| ||d| | 	|| qod S )NTr6   rM   gQ	@r$   r<   r      z<U3abc)rM   r   )re   defZghiZWXYZz<U4y       @      @r   )
r   r   r   Zzerosarrayint64Zfloat64complexr8   assertPreciseEqual)r   rB   rC   Zinpsx1rX   Zx2r   r   r   test_setitem_0d   s&   z!TestFancyIndexing.test_setitem_0dc           
   	      s  t tdd  fdd}g }|d |tg d |tdgdggdgdggg |g d |d	 |d
 tjtjfD ]}tjd|dd}|D ]}||| qYqK|j	}|| d t|t| d |g| d gg}|D ]}
t  || W d    n1 sw   Y  q
t  |dg W d    n1 sw   Y  
t tddt}	|	|dd W d    n1 sw   Y  
t tddt}	|	|ddd W d    n1 sw   Y    d S )NTr6   c                    sD   | |} | |} || t|dr |j|jk d S d S )Norder)rj   hasattrr:   rm   )rA   indrG   rH   rC   rB   r   r   r   check   s   


z-TestFancyIndexing.test_np_take.<locals>.checkr   )r   rd   r      rM   rd   rr   r   rd   r   )rs   )rr   rM   r$      rc   )rL   rM   g333333?r   )r   r   appendr   rg   rh   Z
complex128rN   rO   r?   assertRaises
IndexErrorr   r   Zdisable_leak_check)
r   rq   Ztest_indicesdtr   ro   ZszAZillegal_indicesxZtake_kwsr   rp   r   test_np_take   sH   
$

 zTestFancyIndexing.test_np_takec                 C   s  t dd }t dd }tdd}td tjftjtd ftdtjtg dftjtg dtd ftdttjtg dftg dtjtftjtdtjtg dftg dtd tjftjtdttjtg dftg dtjtjtftjtg dtjtftd	tg dtjd ftjtg dtd fg}|j}|}|j}|}|D ]3}	|||	}
|||	}tj	
|
| t|}||	 }|| |	|}
|| |	|}tj	
|
| qd S )
Nc                 S   r	   r
   r   )r   r   r   r   r   np_new_axis_getitem     z;TestFancyIndexing.test_newaxis.<locals>.np_new_axis_getitemc                 S   s   || |< | S r
   r   )r   r   itemr   r   r   np_new_axis_setitem  s   z;TestFancyIndexing.test_newaxis.<locals>.np_new_axis_setitemiH  )rL   rd         r   )r   r$   r   rM   )r   r   rN   rO   r'   Znewaxisrg   r4   py_funcr=   r>   r\   r8   )r   r{   r~   r   Z	idx_casesZpyfunc_getitemZcfunc_getitemZpyfunc_setitemZcfunc_setitemr   rG   rH   Za_emptyr}   r   r   r   test_newaxis  sB   




zTestFancyIndexing.test_newaxisN)T)__name__
__module____qualname__r#   r.   r2   r5   rI   rS   rV   rZ   r[   r^   r_   r`   rb   rl   rz   r   r   r   r   r   r      s"    



		


9r   c                       s  e Zd ZdZedddeg ddfdeg dededfedeg dfeg ddefedeg ddeddfeg ddeeddfedddeg ddfdeg d	ededfgZ fd
dZ	dd Z
dd Zdd Zdd Zdd Zdd Z  ZS )TestFancyIndexingMultiDim)rd   r   r      	   ra   rL   rd   rM   )r   r   rM   rL   r$   r   N)TFTFTFF)TFTFTFc                    s   t    tjd| _d S )Nr   )supersetUpr   randomZdefault_rngrng)r   	__class__r   r   r   I  s   
zTestFancyIndexingMultiDim.setUpc                 C   s  t | j}td d d td|d d tdd dt|d d dt| d dd td| dtd|d d td| dg}tt|}g }d}t|D ]&}| jjdddd	}| jj	|d
d	
 }| j	d
}	|||	< |t| qIt|D ]&}| jjdddd	}| jj	|d
d	
 }| j	d
}	|||	< |t| qtt|D ]1}| jjdddd	}| jj	|d
d	
 }| jj	d
ddd}	|||	d < t||	d < |t| qt|D ]6}| jjdddd	}| jj	|d
d	
 }| j	d
}	| j|	 }
tj| jj	d|
d	td||	< |t| q|S )Nr   r   r$   r   r%      rd      r?   rL   F)r?   replacerc   )minr;   r'   r0   r   rN   r(   r   Zintegerschoicetolistru   tupler4   rg   bool)r   r    Zslice_choicesZinteger_choicesr   K_Z	array_idxZcurr_idxZ
_array_idxZbool_arr_shaper   r   r   generate_random_indicesM  sV   


	

z1TestFancyIndexingMultiDim.generate_random_indicesc           	      C   s   t dd }tjjdd|d}|j}|jp|}|||}|||}| |j| | |j|j | |j	|j	 tj
|| | t|| d S )Nc                 S   r	   r
   r   )rg   r   r   r   r   numba_get_item  r|   zGTestFancyIndexingMultiDim.check_getitem_indices.<locals>.numba_get_itemr   rr   r   )r   r   r   randintr   r9   ZassertIsNotr:   r;   r<   r=   r>   ZassertFalseZmay_share_memory)	r   	arr_shaperF   r   rA   Zget_itemrE   rG   rH   r   r   r   rI     s   



z/TestFancyIndexingMultiDim.check_getitem_indicesc                 C   s   t dd }tjjdd|d}|| }t|}t|}|||| |||| | |j|j | |j|j tj	
|| d S )Nc                 S   r   r
   r   )rg   r   r}   r   r   r   set_item  s   zATestFancyIndexingMultiDim.check_setitem_indices.<locals>.set_itemr   rr   r   )r   r   r   r   r\   r   r:   r;   r<   r=   r>   )r   r   rF   r   rA   r]   rG   rH   r   r   r   r^     s   


z/TestFancyIndexingMultiDim.check_setitem_indicesc              	   C   ^   | j  }||  7 }|D ]}| j|d | | j| W d    n1 s'w   Y  qd S N)r   )indexing_casesr8   r   subTestrI   r;   r   r   r   r   r   r   test_getitem     
z&TestFancyIndexingMultiDim.test_getitemc              	   C   r   r   )r   r8   r   r   r^   r;   r   r   r   r   test_setitem  r   z&TestFancyIndexingMultiDim.test_setitemc              	   C   s   dddt ddgddggffdddt ddgt ddgffddt ddgtd ddffg}|D ])\}}| t}| | j| W d    n1 sNw   Y  | |t|j	 q3d S )	Nz,Multi-dimensional indices are not supported.r   rM   r   r$   z:Using more than one non-scalar array index is unsupported.zUsing more than one indexing subspace is unsupported. An indexing subspace is a group of one or more consecutive indices comprising integer or array types.rL   )
r   rg   r'   rv   r   rI   r;   ZassertInstr	exception)r   Zerr_idx_caseserrr   Zraisesr   r   r   %test_unsupported_condition_exceptions  s&   z?TestFancyIndexingMultiDim.test_unsupported_condition_exceptions)r   r   r   r;   r'   r   rg   r4   r   r   r   rI   r^   r   r   r   __classcell__r   r   r   r   r   0  s.    
@r   __main__)r)   numpyr   ZunittestZnumbar   r   r   Z
numba.corer   Znumba.core.errorsr   Znumba.tests.supportr   r   r   r   r   r   r   r   r   mainr   r   r   r   <module>   s&       8