
    xKg                        d dl Z 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
mZmZ d dlmZmZ d dlmZ d dlmZmZmZmZ ddiZd	diZd
 Zd Zd Zd Zd Zd Zd Zd Z d Z!d Z"d Z#d Z$d Z%d Z&d Z'd Z(d Z)d Z*d Z+d Z,d Z-ed        Z.ed         Z/d! Z0d" Z1d# Z2 G d$ d%e      Z3 G d& d'e      Z4 G d( d)e      Z5e6d*k(  r ejn                          yy)+    N)jitnjittypeof)utilstypeserrors)TestCasetag)	arraydecl)intpellipsisslice2_typeslice3_typeforceobjTnopythonc                     | |||   S N )astartstopsteps       ]/home/alanp/www/video.onchill/myenv/lib/python3.12/site-packages/numba/tests/test_indexing.pyslicing_1d_usecaser      s    U4_    c                 p    | |||   }d}t        |j                  d         D ]  }|||   |dz   z  z  } |S Nr      rangeshape)r   r   r   r   btotalis          r   slicing_1d_usecase2r%      sI    	%T/AE1771:1Q Lr   c                 j    | || }d}t        |j                  d         D ]  }|||   |dz   z  z  } |S r   r   )r   r   r   r"   r#   r$   s         r   slicing_1d_usecase3r'      sD    	%AE1771:1Q Lr   c                 j    | d d  }d}t        |j                  d         D ]  }|||   |dz   z  z  } |S r   r   )r   r"   r#   r$   s       r   slicing_1d_usecase4r)   $   sB    	!AE1771:1Q Lr   c                 j    | |d  }d}t        |j                  d         D ]  }|||   |dz   z  z  } |S r   r   r   r   r"   r#   r$   s        r   slicing_1d_usecase5r,   +   sD    	%&	AE1771:1Q Lr   c                 j    | d | }d}t        |j                  d         D ]  }|||   |dz   z  z  } |S r   r   )r   r   r"   r#   r$   s        r   slicing_1d_usecase6r.   2   sD    	%4AE1771:1Q Lr   c                 p    | |d d   }d}t        |j                  d         D ]  }|||   |dz   z  z  } |S Nr   r   r   r+   s        r   slicing_1d_usecase7r2   9   sH    	%))AE1771:1Q Lr   c                 p    | d d d   }d}t        |j                  d         D ]  }|||   |dz   z  z  } |S r0   r   r+   s        r   slicing_1d_usecase8r4   A   sF    	$B$AE1771:1Q Lr   c                     | ||||||f   S r   r   )r   start1stop1step1start2stop2step2s          r   slicing_2d_usecaser<   J   s     VE%e!3344r   c                     | ||||f   S r   r   )r   r6   r7   r8   indexs        r   slicing_2d_usecase3r?   N   s    VE%&''r   c                 t    | ||d |f   }d}t        |j                  d         D ]  }|||   |dz   z  z  } |S r   r   )r   index0r6   index2r"   r#   r$   s          r   slicing_3d_usecaserC   R   sL    	&&'6
!"AE1771:1Q Lr   c                 t    | |d ||f   }d}t        |j                  d         D ]  }|||   |dz   z  z  } |S r   r   )r   rA   r7   rB   r"   r#   r$   s          r   slicing_3d_usecase2rE   Y   sN    	&&5&&
 !AE1771:1Q Lr   c                 j    | |   }d}t        |j                  d         D ]  }|||   |dz   z  z  } |S r   r   )r   r>   r"   r#   r$   s        r   partial_1d_usecaserG   `   sB    	%AE1771:1Q Lr   c                     | |   S r   r   )r   r$   s     r   integer_indexing_1d_usecaserI   g   s    Q4Kr   c                     | ||f   S r   r   r   i1i2s      r   integer_indexing_2d_usecaserN   j   s    RU8Or   c                     | |   |   S r   r   rK   s      r   integer_indexing_2d_usecase2rP   m   s    R59r   c                     | ||df   S N.r   r   r$   js      r   ellipsis_usecase1rU   p   s    QqS#X;r   c                     | d||f   S rR   r   rS   s      r   ellipsis_usecase2rW   s   s    S!A#X;r   c                     | |d|f   S rR   r   rS   s      r   ellipsis_usecase3rY   v   s    QQY<r   c                     | d    S r   r   r   s    r   none_index_usecaser\   y   s    T7Nr   c                     | d   S )Nr   r   r[   s    r   empty_tuple_usecaser^   |   s    R5Lr   c                     || |<   y r   r   )r   r>   values      r   setitem_usecasera      s    AeHr   c                     || d d  y r   r   )r   r`   s     r   setitem_broadcast_usecaserc      s    AaDr   c                     || |||<   | S r   r   )r   r"   r   r   r   s        r   slicing_1d_usecase_setre      s    AeDoHr   c                 "    | ||xxx |z  ccc | S r   r   )r   r"   r   r   s       r   slicing_1d_usecase_addrg      s     eDMQMHr   c                      || ||||||f<   | S r   r   )r   r"   r   r   r   r9   r:   r;   s           r   slicing_2d_usecase_setri      s#    ,-AeDofU5(()Hr   c                      e Zd ZdZefdZd ZefdZd ZefdZ	d Z
efdZd	 Zd
 ZefdZd ZefdZd ZefdZd ZefdZd ZefdZd ZefdZd ZefdZd ZefdZd ZefdZd ZefdZ d Z!efdZ"d  Z#ee$fd!Z%d" Z&d# Z'd$ Z(efd%Z)efd&Z*d' Z+efd(Z,d) Z-efd*Z.d+ Z/efd,Z0d- Z1d. Z2d/ Z3efd0Z4d1 Z5efd2Z6d3 Z7y4)5TestGetItemz
    Test basic indexed load from an array (returning a view or a scalar).
    Note fancy indexing is tested in test_fancy_indexing.
    c                 ^   t         }t        j                  t        j                  dd      }|t        j                  t        j                  t        j                  f} t	        |fi ||      }t        j                  dd      }dD ]$  } ||g| }| j                   ||g| |       & y )Nr   C
   i4dtype)r   rn   r         r   rn   r   r   r   rn   r   rn   rt   )	   r   rx   r1   r   r   rx   r   )r   r   Arrayint32r   nparangeassertPreciseEqual)	selfflagspyfunc	arraytypeargtyscfuncr   indicesexpecteds	            r   test_1d_slicingzTestGetItem.test_1d_slicing   s    #KKQ4	U[[%++u{{C$F$e$V,IIb%G a*'*H##E!$6g$6Ar   c                 0    | j                  t               y N)r   )r   Noflagsr   s    r   test_1d_slicing_npmzTestGetItem.test_1d_slicing_npm       7+r   c                 8   t         }t        j                  t        j                  dd      }|t        j                  t        j                  t        j                  f} t	        |fi ||      }t        j                  dd      }g d}|D ]"  }| j                   ||g|  ||g|        $ t        j                  t        j                  dd      }|t        j                  t        j                  t        j                  f} t	        |fi ||      }t        j                  dd      d d d	   }| j                  |j                  d
          | j                  |j                  d          g d}|D ]"  }| j                   ||g|  ||g|        $ y )Nr   rm   rn   ro   rp   rr   rs   rv   rw   ry   A   rt   C_CONTIGUOUSF_CONTIGUOUS)
r%   r   r~   r   r   r   r   assertEqualassertFalser   	r   r   r   r   r   r   r   argsargs	            r   test_1d_slicing2zTestGetItem.test_1d_slicing2   sR   $KKQ4	U[[%++u{{C$F$e$V,IIb% CVA__eAnn= 
 KKQ4	U[[%++u{{C$F$e$V,IIb%cc*0101 CVA__eAnn= r   c                 0    | j                  t               y r   )r   r   r   s    r   test_1d_slicing2_npmz TestGetItem.test_1d_slicing2_npm       G,r   c                    t         }t        j                  t        j                  dd      }|t        j                  t        j                  f} t	        |fi ||      }t        j                  dd      }g d}|D ]"  }| j                   ||g|  ||g|        $ t        j                  t        j                  dd      }|t        j                  t        j                  f} t	        |fi ||      }t        j                  dd      d d d	   }| j                  |j                  d
          | j                  |j                  d          |D ]"  }| j                   ||g|  ||g|        $ y )Nr   rm   rn   ro   rp   ))ru   rn   )rt   ru   )rn   r   r   rn   )   rn   r   r   rt   r   r   )
r'   r   r~   r   r   r   r   r   r   r   r   s	            r   test_1d_slicing3zTestGetItem.test_1d_slicing3   s=   $KKQ4	U[[%++6$F$e$V,IIb% CVA__eAnn= 
 KKQ4	U[[%++6$F$e$V,IIb%cc*0101CVA__eAnn= r   c                 0    | j                  t               y r   )r   r   r   s    r   test_1d_slicing3_npmz TestGetItem.test_1d_slicing3_npm   r   r   c                 P   t         }t        j                  t        j                  dd      }|f} t	        |fi ||      }t        j                  dd      }| j                   ||       ||             t        j                  t        j                  dd      }|f} t	        |fi ||      }t        j                  dd      d d d   }| j                  |j                  d	          | j                  |j                  d
          | j                   ||       ||             y )Nr   rm   rn   ro   rp   r   r   rt   r   r   )
r)   r   r~   r   r   r   r   r   r   r   r   r   r   r   r   r   r   s          r   test_1d_slicing4zTestGetItem.test_1d_slicing4   s    $KKQ4	$F$e$V,IIb%E!H- KKQ4	$F$e$V,IIb%cc*0101E!H-r   c                 0    | j                  t               y r   )r   r   r   s    r   test_1d_slicing4_npmz TestGetItem.test_1d_slicing4_npm  r   r   c                    t        t        dd            }t        j                  t        j                  dd      }|t        j                  f} t        |fi ||      }t        j                  dd      }|D ]"  }| j                   |||       |||             $ t        j                  t        j                  dd      }|t        j                  f} t        |fi ||      }t        j                  dd      d d d	   }| j                  |j                  d
          | j                  |j                  d          |D ]"  }| j                   |||       |||             $ y )Nirn   r   rm   ro   rp   r   r   rt   r   r   )listr    r   r~   r   r   r   r   r   r   r   )	r   r   r   r   r   r   r   r   r   s	            r   check_1d_slicing_with_argz%TestGetItem.check_1d_slicing_with_arg  s+   E"bM"KKQ4	U[[)$F$e$V,IIb%CVAs^U1c];  KKQ4	U[[)$F$e$V,IIb%cc*0101CVAs^U1c]; r   c                 4    t         }| j                  ||       y r   )r,   r   r   r   r   s      r   test_1d_slicing5zTestGetItem.test_1d_slicing5)      $&&vu5r   c                 0    | j                  t               y r   )r   r   r   s    r   test_1d_slicing5_npmz TestGetItem.test_1d_slicing5_npm-  r   r   c                 4    t         }| j                  ||       y r   )r.   r   r   s      r   test_1d_slicing6zTestGetItem.test_1d_slicing60  r   r   c                 0    | j                  t               y r   )r   r   r   s    r   test_1d_slicing6_npmz TestGetItem.test_1d_slicing6_npm4  r   r   c                 4    t         }| j                  ||       y r   )r2   r   r   s      r   test_1d_slicing7zTestGetItem.test_1d_slicing77  r   r   c                 0    | j                  t               y r   )r   r   r   s    r   test_1d_slicing7_npmz TestGetItem.test_1d_slicing7_npm;  r   r   c                 4    t         }| j                  ||       y r   )r4   r   r   s      r   test_1d_slicing8zTestGetItem.test_1d_slicing8>  r   r   c                 0    | j                  t               y r   )r   r   r   s    r   test_1d_slicing8_npmz TestGetItem.test_1d_slicing8_npmB  r   r   c                    t         }t        j                  t        j                  dd      }|t        j                  t        j                  t        j                  f} t	        |fi ||      }t        j                  dd      j                  dd      }dD ]'  }| j                   ||g|  ||g| d|	       ) y
)z
        arr_2d[a:b:c]
        rt   rm   d   ro   rp   rn   r   z	for args )msgN)	r   r   r~   r   r   r   r   reshaper   )r   r   r   r   r   r   r   r   s           r   test_2d_slicingzTestGetItem.test_2d_slicingE  s     $KKQ4	U[[%++u{{C$F$e$V,IIc&..r26.D##F1$4t$4eAoo9=(? $ A.r   c                 0    | j                  t               y r   )r   r   r   s    r   test_2d_slicing_npmzTestGetItem.test_2d_slicing_npmT  r   r   c                     t         }t        j                  t        j                  dd      }|t        j                  t        j                  t        j                  t        j                  t        j                  t        j                  f} t	        |fi ||      }t        j                  dd      j                  dd      }g d}t        j                  ||      D 	cg c]
  \  }}	||	z    }
}}	|
D ]$  } ||g| }| j                   ||g| |       & t        j                  t        j                  dd      }|t        j                  t        j                  t        j                  t        j                  t        j                  t        j                  f} t	        |fi ||      }t        j                  d	d      j                  d
d
      ddddddf   }|
D ]$  } ||g| }| j                   ||g| |       & yc c}	}w )z&
        arr_2d[a:b:c, d:e:f]
        rt   rm   r   ro   rp   rn   )	rr   rs   rv   rw   ry   )rn   r   rx   )rz   r   r1   r{   r}   r     r   N)r<   r   r~   r   r   r   r   r   	itertoolsproductr   )r   r   r   r   r   r   r   r   tup1tup2r   r   r   s                r   test_2d_slicing2zTestGetItem.test_2d_slicing2W  s   
 $KKQ4	U[[%++u{{++u{{EKK9$F$e$V,IIc&..r26	 %.$5$5gw$GI$GLT4 t$G 	 ICa#H##E!NcNH= 
 KKQ4	U[[%++u{{++u{{EKK9$F$e$V,IIc&..r26ssCaCx@Ca#H##E!NcNH= Is   H
c                 0    | j                  t               y r   )r   r   r   s    r   test_2d_slicing2_npmz TestGetItem.test_2d_slicing2_npm  r   r   c                 F   t         }t        j                  t        j                  dd      }|t        j                  t        j                  t        j                  t        j                  f} t	        |fi ||      }t        j                  dd      j                  dd      }g d}|D ]$  } ||g| }	| j                   ||g| |	       & t        j                  t        j                  dd      }|t        j                  t        j                  t        j                  t        j                  f} t	        |fi ||      }t        j                  d	d      j                  d
d
      ddddddf   }|D ]$  } ||g| }	| j                   ||g| |	       & y)z"
        arr_2d[a:b:c, d]
        rt   rm   r   ro   rp   rn   ))r   rn   r   r   )rt   ru   r   r   )rn   r   rx      )rz   r   r1      )r   rn   rt   ru   )r   rx   ru   r   r   r   r   N)	r?   r   r~   r   r   r   r   r   r   )
r   r   r   r   r   r   r   r   r   r   s
             r   test_2d_slicing3zTestGetItem.test_2d_slicing3  s[   
 %KKQ4	U[[%++u{{++$F$e$V,IIc&..r26
 Ca#H##E!NcNH= 
 KKQ4	U[[%++u{{++$F$e$V,IIc&..r26ssCaCx@Ca#H##E!NcNH= r   c                 0    | j                  t               y r   )r   r   r   s    r   test_2d_slicing3_npmz TestGetItem.test_2d_slicing3_npm  r   r   c                    t         }t        j                  t        j                  dd      }|t        j                  t        j                  t        j                  f} t	        |fi ||      }t        j                  dd      j                  ddd      }g d}|D ]"  }| j                   ||g|  ||g|        $ t        j                  t        j                  dd      }|t        j                  t        j                  t        j                  f} t	        |fi ||      }t        j                  d	d      d d d
   j                  ddd      }|D ]"  }| j                   ||g|  ||g|        $ y Nru   rm   i  ro   rp   rn   ))r   rz   r   rs   )rz   r   r   )r   rz   rx   )r   rz   rt   r   i  rt   )	rC   r   r~   r   r   r   r   r   r   r   s	            r   test_3d_slicingzTestGetItem.test_3d_slicing  s9   #KKQ4	U[[%++u{{C$F$e$V,IId$'//B;
 CVA__eAnn=  KKQ4	U[[%++u{{C$F$e$V,IId$'!,44RR@CVA__eAnn= r   c                 0    | j                  t               y r   )r   r   r   s    r   test_3d_slicing_npmzTestGetItem.test_3d_slicing_npm  r   r   c                    t         }t        j                  t        j                  dd      }|t        j                  t        j                  t        j                  f} t	        |fi ||      }t        j                  dd      j                  ddd      }g d}|D ]"  }| j                   ||g|  ||g|        $ t        j                  t        j                  dd      }|t        j                  t        j                  t        j                  f} t	        |fi ||      }t        j                  d	d      d d d
   j                  ddd      }|D ]"  }| j                   ||g|  ||g|        $ y r   )	rE   r   r~   r   r   r   r   r   r   r   s	            r   test_3d_slicing2zTestGetItem.test_3d_slicing2  s9   $KKQ4	U[[%++u{{C$F$e$V,IId$'//B;
 CVA__eAnn=  KKQ4	U[[%++u{{C$F$e$V,IId$'!,44RR@CVA__eAnn= r   c                 0    | j                  t               y r   )r   r   r   s    r   test_3d_slicing2_npmz TestGetItem.test_3d_slicing2_npm  r   r   c                 t   t         }t        j                  t        j                  dd      }|t        j                  f} t	        |fi ||      }t        j                  dd      }| j                   ||d       ||d             | j                   ||d       ||d             | j                   ||d       ||d             t        j                  t        j                  dd	      }|t        j                  f} t	        |fi ||      }t        j                  dd      d d d
   }| j                  |j                  d          | j                  |j                  d          | j                   ||d       ||d             | j                   ||d
       ||d
             | j                   ||d       ||d             t        j                  t        j                  dd      }t        j                  t        j                  dd      }||f} t	        |fi ||      }t        j                  ddt
        j                        }dD ]e  }t        j                  |      j                  t
        j                        }	|	j                  dk(  sJ | j                   |||	       |||	             g y )Nr   rm   rn   ro   rp   r   rz   rx   r   rt   r   r   ru      )r   rz   r1   )rI   r   r~   r   r   r   r   r   r   r   int16arrayastypendim)
r   r   r   r   r   r   r   	indextyper$   idxs
             r   test_1d_integer_indexingz$TestGetItem.test_1d_integer_indexing  s2   ,KKQ4	U[[)$F$e$V,IIb%1uQ{31uQ{32a5 KKQ4	U[[)$F$e$V,IIb%cc*01011uQ{31uQ{32a5 KKQ4	KKQ4	Y'$F$e$V,IIa288,A((1+$$RXX.C88q= =VAs^U1c]; r   c                 0    | j                  t               y r   )r   r   r   s    r   test_1d_integer_indexing_npmz(TestGetItem.test_1d_integer_indexing_npm      %%G%4r   c                    t         }t        j                  t        j                  dd      }|t        j                  f} t	        |fi ||      }t        j                  dd      j                  dd      }| j                   ||d       ||d             | j                   ||d       ||d             | j                   ||d	       ||d	             t        j                  t        j                  dd
      }|t        j                  f} t	        |fi ||      }t        j                  dd      j                  dd      d d d   }| j                   ||d       ||d             y )Nrt   rm   r   ro   rp   rn   r   rz   rx   r   r   r   r   )	rI   r   r~   r   r   r   r   r   r   r   s          r   test_integer_indexing_1d_for_2dz+TestGetItem.test_integer_indexing_1d_for_2d  s.   ,KKQ4	U[[)$F$e$V,IIc&..r26q!eAqk:q!eAqk:q"uQ|<KKQ4	U[[)$F$e$V,IIb%--a3CaC8q!eAqk:r   c                 0    | j                  t               y r   )r   r   r   s    r   #test_integer_indexing_1d_for_2d_npmz/TestGetItem.test_integer_indexing_1d_for_2d_npm%  s    ,,7,;r   c           	      V   t        j                  dd      j                  dd      }t        j                  t        j
                  dd      }|t        j
                  t        j
                  f} t        |fi ||      }| j                   ||dd       ||dd             | j                   ||d	d	       ||d	d	             | j                   ||d
d       ||d
d             t        j                  dd      j                  dd      d d dd d df   }| j                  |j                  d          | j                  |j                  d          t        j                  t        j
                  dd      }|t        j
                  t        j
                  f} t        |fi ||      }| j                   ||dd       ||dd             | j                   ||dd       ||dd             | j                   ||d
d       ||d
d             t        j                  dd      j                  dd      }t        j                  t        j
                  dd      }t        j                  t        j
                  dd      }|||f} t        |fi ||      }dD ]  \  }}	t        j                  |      j                  t         j
                        }t        j                  |	      j                  t         j
                        }	| j                   ||||	       ||||	              y )Nr   ro   rp   rn   rt   rm   r   ru   rz   r1   rx   r   r   r   r   ))r   ru   )r   rz   )r1   rx   )r   r   r   r   r~   r   r   r   r   r   r   r   )
r   r   r   r   r   r   r   r   r$   rT   s
             r   test_2d_integer_indexingz$TestGetItem.test_2d_integer_indexing(  s    IIc&..r26KKQ4	U[[%++6$F$e$V,1a%1a.91a%1a.92r*E!R,<= IIc&..r26ssCaCx@0101KKQ4	U[[%++6$F$e$V,1a%1a.91a%1a.92r*E!R,<= IIc&..r26KKQ4	KKQ4	Y	2$F$e$V,.DAq""288,A""288,AVAq!_eAq!n= /r   c                 0    | j                  t               y r   )r   r   r   s    r   test_2d_integer_indexing_npmz(TestGetItem.test_2d_integer_indexing_npmM  r   r   c                 f    | j                  t               | j                  t        t               y )N)r   )r   r   )r   rP   r   r   s    r   test_2d_integer_indexing2z%TestGetItem.test_2d_integer_indexing2P  s-    %%-I%J%%G-I 	& 	Kr   c                 D   t         d        t         fd       }t        j                  d      j                  dd      }| j	                   ||dd      |d   d          | j	                   ||dd      |d   d          | j	                   ||dd      |d   d          y )	Nc                     | |   S r   r   )Xi0s     r   index1z=TestGetItem.test_2d_integer_indexing_via_call.<locals>.index1V  s    R5Lr   c                      | |   |      S r   r   )r   r   rL   r   s      r   rB   z=TestGetItem.test_2d_integer_indexing_via_call.<locals>.index2Y  s    !B%$$r   rn   rt   r   r   r   rx   )r   r   r   r   r   )r   rB   r   r   s      @r   !test_2d_integer_indexing_via_callz-TestGetItem.test_2d_integer_indexing_via_callU  s    		 
			% 
	%IIbM!!!Q'1a!A$q'21a!A$q'22r*AbE"I6r   c                    t        j                  dd      j                  dd      }t        }t	        j
                  t        j                  dd      }|t        j                  t        j                  f} t        |fi ||      }| j                   ||dd       ||dd             | j                   ||dd       ||dd             | j                   ||d	d	       ||d	d	             y )
Nr   ro   rp   rn   rt   rm   r   rz   rx   )
r   r   r   rN   r   r~   r   float32r   r   )r   r   r   r   r   r   r   s          r   test_2d_float_indexingz"TestGetItem.test_2d_float_indexinga  s    IIc&..r26,KKQ4	U]]EKK8$F$e$V,1a%1a.91a%1a.92r*E!R,<=r   c                 <    t          fd}t        j                  dd      j                  d      }t	        j
                  t        j                  dd      } |||       t        j                  dd      j                  d      j                  }t	        j
                  t        j                  dd	      } |||       t        j                  dd      j                  d      d d d   }t	        j
                  t        j                  dd
      } |||       y )Nc                 2   |t         j                  f} t        |fi       }j                   | d       || d             | j                  d   dz
  }j                   | |       || |             j                   | d       || d             y )Nr   r   rx   )r   r   r   r   r!   )arrr   r   r   nr   r   r   s        r   checkz3TestGetItem.test_partial_1d_indexing.<locals>.checko  s    -F(C(%(0EVC^U3];		!q AVC^U3];VC_eCn=r      ro   rp   )r   ru   rt   rm   )ru   r   Fr   )rG   r   r   r   r   r~   r   T)r   r   r  r   r   r   s   ``   @r   test_partial_1d_indexingz$TestGetItem.test_partial_1d_indexingl  s    #	> IIb%--f5KKQ4	aIIb%--f577KKQ4	aIIb%--f5cc:KKQ4	ar   c                     fdfd} |t        j                  dd      j                  dd              |t        j                  dd      j                  ddd             y )	Nc                 |    t        |       t        j                  t        j                  f} t        |fi       S r   )r   r   r   r   )r  r   r   r   s     r   compile_funcz0TestGetItem.check_ellipsis.<locals>.compile_func  s2    Sk5::uzz:F'3v''//r   c           	          d} |       }t        j                  ||      D ]?  \  }} || ||      }t        j                  j	                   | ||       || ||             A y )N)r   r   rt   rx   r1   )r   r   r   testingassert_equal)r   boundsr   r$   rT   xr  r   s         r   runz'TestGetItem.check_ellipsis.<locals>.run  s]    &F OE!))&&91!QN

''q!Qq!QH :r      ro   rp   r      ru   )r   r   r   )r   r   r   r  r  s    `` @r   check_ellipsiszTestGetItem.check_ellipsis  sR    	0	I 	BIIb%--a34BIIb%--aA67r   c                 0    | j                  t        |       y r   )r  rU   r   r   s     r   test_ellipsis1zTestGetItem.test_ellipsis1      -u5r   c                 0    | j                  t               y r   )r  r   r   s    r   test_ellipsis1_npmzTestGetItem.test_ellipsis1_npm      '*r   c                 0    | j                  t        |       y r   )r  rW   r  s     r   test_ellipsis2zTestGetItem.test_ellipsis2  r  r   c                 0    | j                  t               y r   )r  r   r   s    r   test_ellipsis2_npmzTestGetItem.test_ellipsis2_npm  r  r   c                 0    | j                  t        |       y r   )r  rY   r  s     r   test_ellipsis3zTestGetItem.test_ellipsis3  r  r   c                 0    | j                  t               y r   )r!  r   r   s    r   test_ellipsis3_npmzTestGetItem.test_ellipsis3_npm  r  r   c                     t         d        }|j                  }d}d}t        j                  ||z        j	                  ||      } ||      } ||      }t        j
                  j                  ||       y )Nc                     t        j                  |       }d}t        j                  | d|f         D ]  \  }}|||   |<    |S )Nr   .)r   
zeros_likendenumerate)r  outr$   r>   vals        r   udtz0TestGetItem.test_ellipsis_issue1498.<locals>.udt  sH    --$CA nnSa[9
s #E
1 : Jr   r   )r   py_funcr   r   r   r  r  )r   r*  r+  	outersize	innersizer  gotr   s           r   test_ellipsis_issue1498z#TestGetItem.test_ellipsis_issue1498  sq     
	 
	 ++		ii	I-.66y)L#h3<


X.r   c                     t         d        }t        j                  d      } ||      }|j                  |      }t        j                  j                  ||       y )Nc                     | d   S )N).r   r   )r  s    r   r*  z0TestGetItem.test_ellipsis_issue1499.<locals>.udt  s    v;r   ru   )r   r   r   r+  r  r  )r   r*  r  r.  r   s        r   test_ellipsis_issue1499z#TestGetItem.test_ellipsis_issue1499  sP     
	 
	 iil#h;;s#


X.r   c                    t         }t        j                  t        j                  dd      }|f} t	        |fi ||      }t        j                  dd      j                  dd      }| j                   ||       ||             y )Nrt   rm   r   ro   rp   rn   )	r\   r   r~   r   r   r   r   r   r   r   s          r   test_none_indexzTestGetItem.test_none_index  so    #KKQ4	$F$e$V,IIc&..r26q	584r   c                 z    | j                         5  | j                  t               d d d        y # 1 sw Y   y xY wr   )assertTypingErrorr4  r   r   s    r   test_none_index_npmzTestGetItem.test_none_index_npm  s+    ##%  w / &%%s   1:c                    t         }t        j                  t        j                  dd      }|f} t	        |fi ||      }t        j                  dd      j                  d      }| j                   ||       ||             y )Nr   rm   r   ro   rp   r   )	r^   r   r~   r   r   r   r   r   r   r   s          r   test_empty_tuple_indexingz%TestGetItem.test_empty_tuple_indexing  sm    $KKQ4	$F$e$V,IIat$,,R0q	584r   c                 0    | j                  t               y r   )r9  r   r   s    r   test_empty_tuple_indexing_npmz)TestGetItem.test_empty_tuple_indexing_npm      &&W&5r   N)8__name__
__module____qualname____doc__enable_pyobj_flagsr   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   rN   r   r   r   r   r   r  r  r  r  r  r  r!  r#  r/  r2  r4  r7  r9  r;  r   r   r   rk   rk      s   
 %7 B&, &8 ">H- &8 ><- &8 .&-<, &8 6- &8 6- &8 6- &8 6- %7 A, &8 '>R- &8 #>J- %7 >:, &8 >:- .@ "<H5 5G ;&< .@(C#>J5K

7 ,> 	> .@ .8 $6 6+ $6 6+ $6 6+/*
/ %7 50 /A 56r   rk   c                       e Zd ZdZefdZefdZd ZefdZefdZ	efdZ
efdZd	 Zd
 Zd Zd Zd ZefdZefdZd Zd Zd Zd Zd Zd Zd Zy)TestSetItemzk
    Test basic indexed store into an array.
    Note fancy indexing is tested in test_fancy_indexing.
    c                     ddl m} d }  |di ||      }t        j                  dd      }|j	                         } ||        ||       | j                  ||       y)	z9 this used to work, and was used in one of the tutorials r   )r   c                 l    t        t        |             D ]  }|t        j                  d      z  | |<    y )Nr   )r    lendecimalDecimal)r   r>   s     r   r   z3TestSetItem.test_conversion_setitem.<locals>.pyfunc  s,    s5z*$ws';;e +r   r   rL   rp   N)zvoid(i8[:]))numbar   r   r   copyr   )r   r   r   r   r   r*  controls          r   test_conversion_setitemz#TestSetItem.test_conversion_setitem  sY    	< ,+U+F3ii4(((*wc
W-r   c                 6   t         }t        j                  t        j                  dd      }t        j                  t        j                  dd      }||t        j                  t        j                  t        j                  f} t        |fi ||      }d}t        j                  |d      dz   dd	|d	z
  ||dz   |d
z   d| d	z   | | dz
  | d
z
  g}fd}	t        j                  ||      D ]O  \  }
}dD ]E  }|
||f}t        | } | |	       |   g| } | |	       |   g| }| j                  ||       G Q | j                  t              5   |t        j                  t        j                        ddd       ddd       y# 1 sw Y   yxY w)z+
        1d to 1d slice assignment
        r   rm   r   rn   rM   rp   (   r   rt   ru   r1   c                  2    t        j                   d      S )Nro   rp   )r   r&  )r   s   r   	make_destz2TestSetItem.test_1d_slicing_set.<locals>.make_dest  s    ==D11r   r   rt   rx   r1   N)re   r   r~   r   r   r   r   r   r   r   slicer   assertRaises
ValueErrorr&  )r   r   r   	dest_typesrc_typer   r   Nr  rP  r   r   r   r   r>   pyleftcleftr   s                    @r   test_1d_slicing_setzTestSetItem.test_1d_slicing_set  s{    ( KKQ4	;;u{{As3Xu{{EKKM$F$e$V,ii&+QAq!a%Qqb1fqb1"q&1"q&2	2$,,VV<KE4&dD(t	SZ?$?ik3u:==''6 ' = z*"--2884c1aC +**s   0FFc                    t         }t        j                  t        j                  dd      }||t        j                  t        j                  t        j                  f} t	        |fi ||      j
                  |   j                  }d}t        |      }	t        j                  |t        j                        }
|d| |	z   dz   df} ||
j                         g| } ||
j                         g| }| j                  ||       |d| |	z   df}| j                  t              5 } ||
j                         g|  ddd       y# 1 sw Y   yxY w)z9
        Generic sequence to 1d slice assignment
        r   rm   rn   rp   N)re   r   r~   r   r   	overloadsentry_pointrF  r   r   rJ  r   rS  rT  )r   r   seqtyseqr   rU  r   r   rW  kr   r   r   r.  raisess                  r   check_1d_slicing_set_sequencez)TestSetItem.check_1d_slicing_set_sequence  s    (KKQ4	UEKKekkJ %F$e$V,66v>JJHii*QQ
A&#((*,t,CHHJ&&#.QQ"z*f#((*$t$ +**s   $EEc                 n    | j                  |t        j                  t        j                  d      d       y)z.
        Tuple to 1d slice assignment
        rt   )r   N)rb  r   UniTupler   r  s     r   test_1d_slicing_set_tuplez%TestSetItem.test_1d_slicing_set_tuple,  s(     	**5>>%++q18	=r   c                 p    | j                  |t        j                  t        j                        ddg       y)z-
        List to 1d slice assignment
        r   rd  N)rb  r   Listr   r  s     r   test_1d_slicing_set_listz$TestSetItem.test_1d_slicing_set_list3  s*     	**5::ekk*QH	6r   c                 N   t         }t        j                  t        j                  dd      }|t        j                  t        j                  t        j                  t        j                  f} t        |fi ||      }d}t        j                  |d      }d}dd|dz
  ||dz   |d	z   d
| dz   | | dz
  | d	z
  g}	t        j                  |	|	      D ]R  \  }
}dD ]H  }||
||f} ||j                         g| } ||j                         g| }| j                  ||       J T y)z/
        scalar to 1d slice assignment
        r   rm   rn   ro   rp   *   r   rt   ru   r1   rQ  N)re   r   r~   r   r   r   r   r   r   r   rJ  r   )r   r   r   r   r   r   rW  r   r)  r  r   r   r   r   rX  rY  s                   r   test_1d_slicing_broadcastz%TestSetItem.test_1d_slicing_broadcast:  s    (KKQ4	 U[[%++u{{EKKP$F$e$V,ii&QAq!a%Qqb1fqb1"q&1"q&2$,,VV<KE4&E4-
2T2chhj040''6	 ' =r   c                    t         }t        j                  t        j                  dd      }||t        j                  t        j                  f} t	        |fi ||      }t        j                  dd      }dD ]`  } |t        j                  |      |t        |    g| } |t        j                  |      |t        |    g| }	| j                  ||	       b y )Nr   rm   rn   ro   rp   )r   )rt   r   )
rg   r   r~   r   r   r   r   r&  rR  r   )
r   r   r   r   r   r   r   testrX  rY  s
             r   test_1d_slicing_addzTestSetItem.test_1d_slicing_addQ  s    'KKQ4	YU[[A$F$e$V,ii$'%DBMM#.E4L0AIDIF"--,c%,.?G$GE##FE2 &r   c                 0    | j                  t               y r   )rZ  r   r   s    r   test_1d_slicing_set_npmz#TestSetItem.test_1d_slicing_set_npm]        w /r   c                 0    | j                  t               y r   )ri  r   r   s    r   test_1d_slicing_set_list_npmz(TestSetItem.test_1d_slicing_set_list_npm`  r   r   c                 0    | j                  t               y r   )rf  r   r   s    r   test_1d_slicing_set_tuple_npmz)TestSetItem.test_1d_slicing_set_tuple_npmc  r<  r   c                 0    | j                  t               y r   )rl  r   r   s    r   test_1d_slicing_broadcast_npmz)TestSetItem.test_1d_slicing_broadcast_npmf  r<  r   c                 0    | j                  t               y r   )ro  r   r   s    r   test_1d_slicing_add_npmz#TestSetItem.test_1d_slicing_add_npmi  rr  r   c                    t         }t        j                  t        j                  dd      }||t        j                  t        j                  t        j                  t        j                  t        j                  t        j                  f} t	        |fi ||      }t        j                  dd      j                  dd      }g d}|D ]|  } |t        j                  |      |t        |dd	  t        |d	d
  f   g| }	 |t        j                  |      |t        |dd	  t        |d	d
  f   g| }
| j                  |
|	       ~ y)z+
        2d to 2d slice assignment
        rt   r   r   ro   rp   rn   )r   rn   r   r   rn   r   )rt   ru   r   rt   ru   r   )rn   r   r   rn   r   r   )r   rn   rx   r   rn   rx   )r   rn   rt   r   rn   rt   r   ru      N)ri   r   r~   r   r   r   r   r   r&  rR  r   )r   r   r   r   r   r   r   testsrn  rX  rY  s              r   test_2d_slicing_setzTestSetItem.test_2d_slicing_setl  s     (KKQ4	YU[[%++++u{{EKK9$F$e$V,iiT*222b9
 DBMM#.E4!94EudSTUViGX4X0Ya\`aF"--,c%a2CUDQRSTIEV2V.W_Z^_E##E62 r   c                 B   t         }t        j                  t        j                  dd      }|t        j                  t        j                  t        j                  t        j                  t        j                  t        j                  t        j                  f} t        |fi ||      }t        j                  dd      j                  dd      }d}g d}|D ]D  }	 ||j                         |g|	 }
 ||j                         |g|	 }| j                  ||
       F y	)
z/
        scalar to 2d slice assignment
        rt   rm   r   ro   rp   rn   rk  r|  N)ri   r   r~   r   r   r   r   r   r   rJ  r   )r   r   r   r   r   r   r   r)  r~  rn  rX  rY  s               r   test_2d_slicing_broadcastz%TestSetItem.test_2d_slicing_broadcast  s     (KKQ4	 U[[%++u{{EKK++u{{EKK9$F$e$V,iiT*222b9
 DCHHJ3d3F#((*c1D1E##E62 r   c                 0    | j                  t               y r   )r  r   r   s    r   test_2d_slicing_set_npmz#TestSetItem.test_2d_slicing_set_npm  rr  r   c                 0    | j                  t               y r   )r  r   r   s    r   test_2d_slicing_broadcast_npmz)TestSetItem.test_2d_slicing_broadcast_npm  r<  r   c                     t        j                  d      }t        |dd       | j                  |j	                         g d       t        |t        j
                  d      j                  t         j                        d       | j                  |j	                         g d       t        j                  d      j                  dd      }t        |dd       | j                  |j	                         g d	g d
g dg       y)z+
        scalar indexed assignment
        r   r   rk  )r   rk  rt   ru   r   ru   r   )r   rk  rt   r   r   rz   )r   r   rt   )rk  rk  rk  )r}     r   N)	r   r   ra   r   tolistr   r   uint16r   )r   r  s     r   test_setitemzTestSetItem.test_setitem  s     iilQ#'78RXXa[//		:A>'78iil""1a(Q#	<'KLr   c           	      D   t        j                  d      }t        |d       | j                  |j	                         dgdz         t        j                  d      j                  dd      }t        |t        j                  dd             | j                  |j	                         g dg dg       t        j                  d      j                  dd      }t        |t        j                  dd      j                  dd             | j                  |j	                         g dg dg       t        j                  d	      j                  dddd      }t        |t        j                  dd      j                  dd             g dg dg}| j                  |j	                         |ggdz         t        j                  d      }t        |t        j                  dd      j                  dd             | j                  |j	                         g d
       t        j                  d      j                  dd      }t        |t        j                  dd|j                  z         j                  dddd             | j                  |j	                         g dg dg       y)z.
        broadcasted array assignment
        r   rk  r}  rt   ru   r   r   )r   rt   ru   r  )r   rt   ru   r   r   )r   r   r}  N)r   r   rc   r   r  r   size)r   dstinner2s      r   test_setitem_broadcastz"TestSetItem.test_setitem_broadcast  s   
 iil!#r*tax0iil""1a(!#ryyA7	9'=>iil""1a(!#ryyA'>'>q!'DE	9'=>iim##Aq!Q/!#ryyA'>'>q!'DEY'zA~6iil!#ryyA'>'>q!'DE7iil""1a(!#ryyAL'A'I'I!QPQST'UV	9'=>r   c                    t        j                  d      }t        j                  d      j                  dd      }| j                  t              5 }t        ||       d d d        t        j                        }| j                  d|       t        j                  d      j                  dd      }t        j                  d      j                  ddd      }| j                  t              5 }t        ||       d d d        t        |j                        }| j                  d|       t        j                  d      j                  dd      }t        j                  d      }| j                  t              5 }t        ||       d d d        t        |j                        }| j                  d|       y # 1 sw Y   YxY w# 1 sw Y   xY w# 1 sw Y   JxY w)Nr   rn   rt   z,cannot broadcast source array for assignmentr   z0cannot assign slice from input of different sizer   )	r   r   r   rS  rT  rc   str	exceptionr   )r   r  srcra  errmsgs        r   test_setitem_broadcast_errorz(TestSetItem.test_setitem_broadcast_error  sp    iiliim##Aq)z*f%c3/ +V%%&G	! iil""1a(iim##Aq!,z*f%c3/ +V%%&K	! iim##Aq)iilz*f%c3/ +V%%&K	!) +* +* +*s$   F4,G7G4F>G
Gc                    t        j                  d      j                  dd      }t        j                  dd      }t        ||ddd       | j	                  |j                         ddgddgddgg       t        j                  d      j                  dd      }t        j                  dd      }t        ||dd d       | j	                  |j                         ddgddgddgg       t        j                  d      j                  dd      }t        j                  dd      j                  dd      }t        ||d dd       | j	                  |j                         ddgddgddgg       y )Nr}  ru   rt   r   r   r   r   )r   r   r   re   r   r  )r   r  r  s      r   test_slicing_1d_broadcastz%TestSetItem.test_slicing_1d_broadcast  s5   iil""1a(ii1osCAq1AAA'?@iil""1a(ii1osCD!4AAA'?@iil""1a(ii1o%%a+sCq!4AAA'?@r   c                 (   t        j                  d      }d|j                  _        | j	                  t
        t        j                  f      5 }t        |dd       d d d        | j                  dt        j                               y # 1 sw Y   /xY w)Nr   Fr   rk  z%Cannot modify readonly array of type:)r   r   r   	writeablerS  	TypeErrorr   TypingErrorra   assertInr  r  )r   r  ra  s      r   test_setitem_readonlyz!TestSetItem.test_setitem_readonly  so    iil#			6+=+=>?6CB' @=&**+	- @?s   BBN)r=  r>  r?  r@  rA  rL  rZ  rb  rf  ri  rl  ro  rq  rt  rv  rx  rz  r  r  r  r  r  r  r  r  r  r   r   r   rC  rC    s    
 -? .  ); D<%0 /A = .@ 6 /A 7. ); 
305660 ); 3. /A 3406M?<!8A"-r   rC  c                       e Zd ZdZd Zy)
TestTypingz3
    Check typing of basic indexing operations
    c                    t         j                  }t        j                  t        j                  dd      }t        j                  t        j                  dd      }t        j                  t        j                  dd      }dt
        fddft        fddft        fddft        t        fddft        t        fddft        t        fddft        t        fddft        fddft
        t        fddft
        t        fddft
        t        t        fddft
        t        t        fddft
        t        t        fddft
        t        fddft
        t        t        fddft        t
        t        fddft        t
        t        fddft        t        t        fddft        t
        t        t        fddft        t        t        fddft        t        t
        t        fddft        t        t        fddft        t        t        fddft        t        t        fddfg}|D ]  \  }}}t        j                  |      }	 |||	      }
| j                  t        |
j                        |       | j                  |
j                  j                  |rdnd|       | j                  |
j                           |D ]  \  }}}t        j                  |      }	 |||	      }
| j                  t        |
j                        |       | j                  |
j                  j                  |rdnd|       | j                  |
j                           |D ]  \  }}}t        j                  |      }	 |||	      }
| j                  t        |
j                        |       | j                  |
j                  j                  d       | j                  |
j                           y)	zc
        Check an appropriate layout is inferred for the result of array
        indexing.
        ru   rm   r  r   )r   TTTFN)r   get_array_index_typer   r~   float64r   r   r   r   Tupler   tupler>   resultlayoutr   advanced)r   funcctyftyatyr   index_tuplekeep_c_r>   rkeep_fs               r   test_layoutzTestTyping.test_layout  s>    --kk%--C0kk%--C0kk%--C0 [$% WdE"^T5)K $.4 %/;'6K %/^UE* t,%ud3k+UE:T+UD9[2E5A%ue4T+UE: Hd$eU38[2E5AD+'u5HdK14?4'54415$?K'64'6D+'6G$L '."KKK,ES% AU177^[9QXX__Vc(*QZZ( '. '."KFKK,ES% AU177^[9QXX__Vc(*QZZ( '. ")KAKK,ES% AU177^[9QXX__c2QZZ( ")r   N)r=  r>  r?  r@  r  r   r   r   r  r    s    G)r   r  __main__)8rG  r   numpyr   unittestrI  r   r   r   
numba.corer   r   r   numba.tests.supportr	   r
   numba.core.typingr   numba.core.typesr   r   r   r   rA  r   r   r%   r'   r)   r,   r.   r2   r4   r<   r?   rC   rE   rG   rI   rN   rP   rU   rW   rY   r\   r^   ra   rc   re   rg   ri   rk   rC  r  r=  mainr   r   r   <module>r     s'       # # + + - ' E E !$' t
5(    
D	6( D	6Nc-( c-L	L) L)^ zHMMO r   