
    ܙd                        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	m
Z
 d dlmZmZmZ d dlmZ  e            Zde_         e            Zde_         e            Zd Zd	 Zd
 Zd Zd Zd Zd Zd Zd Z ej        dej         fdej!        fg          Z" G d dee          Z# G d dee          Z$e%dk    r ej&                     dS dS )    N)njit)compile_isolatedFlags)typeserrors)TestCaseMemoryLeakMixintag)numpy_supportTc                      d} dD ]}| |z  } | S )Nr   )      c       resis     :lib/python3.11/site-packages/numba/tests/test_iteration.pyint_tuple_iter_usecaser      s&    
C  qJ    c                      d} dD ]}| |z  } | S )N        )      ?       @33333X@333333@r   r   s     r   float_tuple_iter_usecaser      s&    
C"  qJr   c                  2    d} dD ]}|D ]}| |z  } | dz  } | S Nr   ))r   r   )r   r   )g?g      @r   r   )r   r   js      r   tuple_tuple_iter_usecaser"      s@    
C2   	 	A1HCCAgJr   c                  X    d} t          d          D ]\  }}|D ]
}| ||z  z  } | dz  } | S r    	enumerate)r   r   r!   ls       r   enumerate_nested_tuple_usecaser'   (   sQ    
C?@@  1 	 	A1q5LCCAgJr   c                  ~    d} t          t          d                    D ]\  }\  }}|D ]}| ||z  |z  z  } | dz  } | S r    r$   )r   r   r!   kr&   s        r   nested_enumerate_usecaser*   0   sb    
Cy)NOOPP  	6Aq 	 	A1q519CCAgJr   c                      d} t          j        d          t          j        d          f}t          |          D ]\  }}| |                                z  } | S )Nr         )nponesr%   sum)r   arraysr   vs       r   enumerate_array_usecaser3   9   sT    
Cgajj"'!**%F&!!  1quuwwJr   c                     d}| D ]}||z  }|S Nr   r   iterabler   xs      r   scalar_iter_usecaser9   A   s&    
C  qJr   c                 8    d}| D ]}||j         |j        z  z  }|S r5   abr6   s      r   record_iter_usecaser>   G   s/    
C  qsQSyJr   c                 8    | D ]}|j         |j        z   |_         d S Nr;   )r7   r8   s     r   record_iter_mutate_usecaserA   M   s,      cACi r   r<   r=   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efd
Z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d Zd Zd ZdS )IterationTestc                     t          |d|          }|j        } |            }|                      |            |           d S )Nr   flags)r   entry_pointassertPreciseEqual)selfpyfuncrF   crcfuncexpecteds         r   run_nullary_funczIterationTest.run_nullary_funcY   sJ    fb66668822222r   c                 <    |                      t          |           d S r@   )rN   r   rI   rF   s     r   test_int_tuple_iterz!IterationTest.test_int_tuple_iter_   s    4e<<<<<r   c                 <    |                      t                     d S NrE   )rQ   no_pyobj_flagsrI   s    r   test_int_tuple_iter_npmz%IterationTest.test_int_tuple_iter_npmb   s      ~ 66666r   c                 <    |                      t          |           d S r@   )rN   r   rP   s     r   test_float_tuple_iterz#IterationTest.test_float_tuple_iterh       6>>>>>r   c                 <    |                      t                     d S rS   )rX   rT   rU   s    r   test_float_tuple_iter_npmz'IterationTest.test_float_tuple_iter_npmk       """88888r   c                 <    |                      t          |           d S r@   )rN   r"   rP   s     r   test_tuple_tuple_iterz#IterationTest.test_tuple_tuple_itern   rY   r   c                 <    |                      t                     d S rS   )r^   rT   rU   s    r   test_tuple_tuple_iter_npmz'IterationTest.test_tuple_tuple_iter_npmq   r\   r   c                 <    |                      t          |           d S r@   )rN   r'   rP   s     r   test_enumerate_nested_tuplez)IterationTest.test_enumerate_nested_tuplet   s    <eDDDDDr   c                 <    |                      t                     d S rS   )rb   rT   rU   s    r   test_enumerate_nested_tuple_npmz-IterationTest.test_enumerate_nested_tuple_npmw   s    ((~(>>>>>r   c                 <    |                      t          |           d S r@   )rN   r*   rP   s     r   test_nested_enumeratez#IterationTest.test_nested_enumeratez   rY   r   c                 <    |                      t                     d S rS   )rf   rT   rU   s    r   test_nested_enumerate_npmz'IterationTest.test_nested_enumerate_npm}   r\   r   c                     t           }t          |d          }|j        } |            }|                      |            |           d S )Nr   )r3   r   rG   rH   )rI   rJ   rK   rL   rM   s        r   test_enumerate_refctz"IterationTest.test_enumerate_refct   sJ    (fb))68822222r   c                     t           }t          |t          j        |dd          f||          }|j        }|                      ||           ||                     d S )Nr   ArE   )r9   r   r   ArrayrG   rH   )rI   	item_typeargrF   rJ   rK   rL   s          r   run_array_1dzIterationTest.run_array_1d   sh    $fu{9a'E'E&G'u6 6 6c

FF3KK88888r   c                 l    |                      t          j        t          j        d          |           d S )N      @)rp   r   float64r.   arangerP   s     r   test_array_1d_floatz!IterationTest.test_array_1d_float   s*    %-3?????r   c                 :    |                      t                     d S r@   )ru   rT   rU   s    r   test_array_1d_float_npmz%IterationTest.test_array_1d_float_npm   s      00000r   c                 r    |                      t          j        t          j        d          dz  |           d S )Nrr   y              ?)rp   r   
complex128r.   rt   rP   s     r   test_array_1d_complexz#IterationTest.test_array_1d_complex   s0    %*BIcNNT,A5IIIIIr   c                 :    |                      t                     d S r@   )rz   rT   rU   s    r   test_array_1d_complex_npmz'IterationTest.test_array_1d_complex_npm   s    "">22222r   c                    t           }t          j        t                    }t	          |t          j        |dd          f|          }|j        }t          j	        dt                    }t          d          D ]/}t          |dz            ||         _        |dz   ||         _        0 ||          }|                      ||          |           d S Nr   rl   rE   r   dtyper   )r>   r   
from_dtyperecord_dtyper   r   rm   rG   r.   recarrayrangefloatr<   r=   rH   )	rI   rF   rJ   rn   rK   rL   arrr   gots	            r   test_array_1d_recordz"IterationTest.test_array_1d_record   s    $!,\::	fu{9a'E'E&G$)+ + +k!<000q 	 	AQU||CFH1uCFHHfSkkc

C00000r   c                 :    |                      t                     d S r@   )r   rT   rU   s    r   test_array_1d_record_npmz&IterationTest.test_array_1d_record_npm   s    !!.11111r   c                    t           }t          j        t                    }t	          |t          j        |dd          f|          }|j        }t          j	        dt                    }t          d          D ]/}t          |dz            ||         _        |dz   ||         _        0|                                } ||           |                                }	 ||	           |                     ||	           d S r~   )rA   r   r   r   r   r   rm   rG   r.   r   r   r   r<   r=   copyrH   )
rI   rF   rJ   rn   rK   rL   r   r   rM   r   s
             r   test_array_1d_record_mutate_npmz-IterationTest.test_array_1d_record_mutate_npm   s    +!,\::	fu{9a'E'E&G$)+ + +k!<000q 	 	AQU||CFH1uCFHH88::xhhjjc


#.....r   c                 <    |                      t                     d S rS   )r   force_pyobj_flagsrU   s    r   test_array_1d_record_mutatez)IterationTest.test_array_1d_record_mutate   s     ,,3D,EEEEEr   c                 ,   d }|                      t          j                  5 }t          j        t          j        dd          }t          ||f           d d d            n# 1 swxY w Y   |                     dt          |j	                             d S )Nc                     | D ]}d S r@   r   )r8   r   s     r   fooz/IterationTest.test_array_0d_raises.<locals>.foo   s       r   r   Cz	0-d array)
assertRaisesr   TypingErrorr   rm   int32r   assertInstr	exception)rI   r   raisesarytys       r   test_array_0d_raisesz"IterationTest.test_array_0d_raises   s    	 	 	
 v122 	,fKQ44ES5(+++	, 	, 	, 	, 	, 	, 	, 	, 	, 	, 	, 	, 	, 	, 	, 	k3v'7#8#899999s   2A!!A%(A%c                 $   d }t          j        dt           j                  x}}t          j        t          j        dd          }t          |||f          } |||          }|                    ||          }|                     ||           d S )Nc                 X    d}t          | |          D ]}||d         |d         z   z  }|S Nr   r   )zip)r8   ytotalrows       r   barz4IterationTest.test_tuple_iter_issue1504.<locals>.bar   s;    E1ayy ) )Q#a&(Lr   r   r   r   r   )r.   rt   r   r   rm   r   rG   assertEqualrI   r   r8   r   r   cresexpectr   s           r   test_tuple_iter_issue1504z'IterationTest.test_tuple_iter_issue1504   s    	 	 	 	!28,,,,AEKC00eU^44Qq!$$%%%%%r   c                 b   d }t          j        dt           j                  x}}t          j        t          j        dd          }t          |t          j                            |d          f          } |||f          }|                    ||f          }| 	                    ||           d S )Nc                 *    d}| D ]}||d         z  }|S )Nr   r   )arrsr   r   s      r   r   z4IterationTest.test_tuple_of_arrays_iter.<locals>.bar   s*    E    QLr   r   r   r   r   r   )
r.   rt   r   r   rm   r   
containersUniTuplerG   r   r   s           r   test_tuple_of_arrays_iterz'IterationTest.test_tuple_of_arrays_iter   s    	 	 	 	!28,,,,AEKC00e&6&?&?q&I&I%KLLaV1v&&%%%%%r   N)__name__
__module____qualname__rN   r   rQ   rV   rX   r[   r^   r`   rb   rd   rf   rh   rj   rp   ru   rw   rz   r|   r   r   rT   r   r   r   r   r   r   r   r   rC   rC   W   s       3 3 3 ): = = = =7 7 7 +< ? ? ? ?9 9 9 +< ? ? ? ?9 9 9 1B E E E E? ? ? +< ? ? ? ?9 9 93 3 39 9 9 ): @ @ @ @1 1 1 +< J J J J3 3 3 *; 1 1 1 12 2 2 5C / / / / F F F: : :& & &"& & & & &r   rC   c                       e Zd Zd ZdS )TestIterationRefctc                     t           d             }d t          d          D             }|                      ||          |                    |                     d S )Nc                     d}t          t          t          |                     |           D ]"\  }}||dz   |                                z  z  }#d S r   )r   r   lenr0   )sequencecr<   r=   s       r   r   z4TestIterationRefct.test_zip_with_arrays.<locals>.foo   sT    AE#h--00(;; ' '1a!equuww&&Fr   c                 <    g | ]}t          j        d |z             S )r   )r.   rt   ).0r   s     r   
<listcomp>z;TestIterationRefct.test_zip_with_arrays.<locals>.<listcomp>   s&    888BIa!e$$888r   
   )r   r   r   py_func)rI   r   r   s      r   test_zip_with_arraysz'TestIterationRefct.test_zip_with_arrays   sf    		 	 
	 98eBii888XH(=(=>>>>>r   N)r   r   r   r   r   r   r   r   r      s#        	? 	? 	? 	? 	?r   r   __main__)'numpyr.   numbar   unittestnumba.core.compilerr   r   
numba.corer   r   numba.tests.supportr   r	   r
   numba.npr   enable_pyobj_flagsenable_pyobjectr   force_pyobjectrT   r   r   r"   r'   r*   r3   r9   r>   rA   r   rs   r   r   rC   r   r   mainr   r   r   <module>r      s              7 7 7 7 7 7 7 7 $ $ $ $ $ $ $ $ > > > > > > > > > > " " " " " "UWW %)  "EGG #'                     
 rx#rz*rx  
U& U& U& U& U&OX U& U& U&r
? 
? 
? 
? 
?( 
? 
? 
? zHMOOOOO r   