
    xKgA5                         d dl Zd dl Z d dlZd dlmZmZ d dlmZmZ d dl	m
Z
 d dlmZmZ  G d de      Z G d d	e      Zed
k(  r ej"                          yy)    N)njitjit)TypingErrorUnsupportedError)ir)TestCaseIRPreservingTestPipelinec                   T    e 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y)TestClosurec                     dfd} t        di ||      }| j                   |d      d       d| j                   |d      d       y )N
   c                     | z   S N xYs    \/home/alanp/www/video.onchill/myenv/lib/python3.12/site-packages/numba/tests/test_closure.pyadd_Yz3TestClosure.run_jit_closure_variable.<locals>.add_Y       q5L             zi4(i4)r   assertEqual)selfjitargsr   c_add_Yr   s       @r   run_jit_closure_variablez$TestClosure.run_jit_closure_variable   sR    	 +#*'*51R( R(r   c                 (    | j                  d       y NT)forceobjr!   r   s    r   test_jit_closure_variablez%TestClosure.test_jit_closure_variable       %%t%4r   c                 (    | j                  d       y NTnopythonr%   r&   s    r   test_jit_closure_variable_npmz)TestClosure.test_jit_closure_variable_npm   r(   r   c                     dfd} t        di ||      }| j                   |d      d       d t        di ||      }| j                   |d      d       d| j                   |d      d       | j                   |d      d       y )Nr   c                     | z   S r   r   r   s    r   r   z0TestClosure.run_rejitting_closure.<locals>.add_Y%   r   r   r   r   r      r   r   )r   r   r   r    	c_add_Y_2r   s        @r   run_rejitting_closurez!TestClosure.run_rejitting_closure"   s    	 +#*'*51R( ,C,G,U3	1r*1r*R(r   c                 (    | j                  d       y r#   r2   r&   s    r   test_rejitting_closurez"TestClosure.test_rejitting_closure4       ""D"1r   c                 (    | j                  d       y r*   r4   r&   s    r   test_rejitting_closure_npmz&TestClosure.test_rejitting_closure_npm7   r6   r   c                 n    ddfd} t        di ||      }| j                   |d      d       y )Nr      c                     | z   z  S r   r   )r   r   Zs    r   add_Y_mult_ZzDTestClosure.run_jit_multiple_closure_variables.<locals>.add_Y_mult_Z>   s    EQ;r   r      r   r   )r   r   r=   c_add_Y_mult_Zr   r<   s       @@r   "run_jit_multiple_closure_variablesz.TestClosure.run_jit_multiple_closure_variables:   s=    	 211,?*B/r   c                 (    | j                  d       y r#   r@   r&   s    r   #test_jit_multiple_closure_variablesz/TestClosure.test_jit_multiple_closure_variablesD       ///>r   c                 (    | j                  d       y r*   rB   r&   s    r   'test_jit_multiple_closure_variables_npmz3TestClosure.test_jit_multiple_closure_variables_npmG   rD   r   c                    d } t        di ||      j                          fd} t        di ||      }|j                          | j                  |      5  | j                   |d      d       d d d        y # 1 sw Y   y xY w)Nc                     | dz  S )Nr   r   )as    r   mult_10z3TestClosure.run_jit_inner_function.<locals>.mult_10K   s    r6Mr   c                      | dz         S N   r   )r   	c_mult_10s    r   do_mathz3TestClosure.run_jit_inner_function.<locals>.do_mathQ   s    QU##r   r   2   )z
intp(intp))r   disable_compileassertRefCountr   )r   r   rJ   rO   	c_do_mathrN   s        @r   run_jit_inner_functionz"TestClosure.run_jit_inner_functionJ   s    	 1C009	!!#	$ 1C009	!!#  I6Yq\2. 766s   B  B	c                 (    | j                  d       y r#   rT   r&   s    r   test_jit_inner_functionz#TestClosure.test_jit_inner_functionZ       ##T#2r   c                 (    | j                  d       y r*   rV   r&   s    r   test_jit_inner_function_npmz'TestClosure.test_jit_inner_function_npm]   rX   r   N)__name__
__module____qualname__r!   r'   r-   r2   r5   r8   r@   rC   rF   rT   rW   rZ   r   r   r   r   r      s>    )55)$220??/ 33r   r   c                   d    e Zd ZdZd Zd Zd Zd Zd Zd Z	d Z
d	 Zej                  d
        Zy)TestInlinedClosurez
    Tests for (partial) closure support in njit. The support is partial
    because it only works for closures that can be successfully inlined
    at compile time.
    c                 \    d }t        |      }| j                   |d       |d             y )Nc                 *    d } ||        ||       z   S )Nc                     | | z  S r   r   )r   s    r   innerzDTestInlinedClosure.test_inner_function.<locals>.outer.<locals>.innerl       1ur   r   r   rc   s     r   outerz5TestInlinedClosure.test_inner_function.<locals>.outerj   s     8eAh&&r   r   r   r   r   rf   cfuncs      r   test_inner_functionz&TestInlinedClosure.test_inner_functionh   s*    	' UrE"I.r   c                 \    d }t        |      }| j                   |d       |d             y )Nc                 :    | dz   fd} ||        ||       z   S )Nr   c                     | | z  z   S r   r   r   ys    r   rc   zQTestInlinedClosure.test_inner_function_with_closure.<locals>.outer.<locals>.innery   s    1uqy r   r   r   rc   ro   s     @r   rf   zBTestInlinedClosure.test_inner_function_with_closure.<locals>.outerv   s&    AA! 8eAh&&r   r   rg   rh   s      r    test_inner_function_with_closurez3TestInlinedClosure.test_inner_function_with_closuret   s*    	' UrE"I.r   c                 \    d }t        |      }| j                   |d       |d             y )Nc                 >    | dz   fd} ||        ||       z   S )Nr   c                     | z  S r   r   rn   s    r   rc   zSTestInlinedClosure.test_inner_function_with_closure_2.<locals>.outer.<locals>.inner       1ur   r   rp   s     @r   rf   zDTestInlinedClosure.test_inner_function_with_closure_2.<locals>.outer   s*    AA aAuQx<r   r   rg   rh   s      r   "test_inner_function_with_closure_2z5TestInlinedClosure.test_inner_function_with_closure_2   s*    	  UrE"I.r   c                     d}i }t        |j                         |       t        |d         }| j                   |d       |d   d             y )Na  
            def outer(x):
                y = x + 1
                z = 0

                def inner(x):
                    nonlocal z
                    z += x * x
                    return z + y

                return inner(x) + inner(x) + z
        rf   r   )execstripr   r   )r   codensri   s       r   "test_inner_function_with_closure_3z5TestInlinedClosure.test_inner_function_with_closure_3   sL     TZZ\2R[!rKBwKO4r   c                 \    d }t        |      }| j                   |d       |d             y )Nc                 $      fd} |  z        S )Nc                 P      fd}d}t               D ]  }| ||      z  } |S )Nc                     z   | z   S r   r   zr   ro   s    r   	innermostz^TestInlinedClosure.test_inner_function_nested.<locals>.outer.<locals>.inner.<locals>.innermost   s    q519$r   r   range)ro   r   sir   s   `   r   rc   zKTestInlinedClosure.test_inner_function_nested.<locals>.outer.<locals>.inner   s0    % qA1%A "r   r   re   s   ` r   rf   z<TestInlinedClosure.test_inner_function_nested.<locals>.outer   s     Q<r   r   rg   rh   s      r   test_inner_function_nestedz-TestInlinedClosure.test_inner_function_nested   s*    	  UrE"I.r   c                 .   t         fd       fd}fd}d }d }d }d }dfd	}dfd
}d }	d }
d }d }d!d}d!d}d!d}d }d }d }d }d }d }d }|||||||	|
|||||||g}|D ]-  }t        |      }d}| j                   ||       ||             / | j                  t              5 } t	        d      |      } |       ddd       d}| j                  |t        j                               | j                  t              5 } t	        d      |      } |       ddd       d}| j                  |t        |j                               | j                  t              5 } t	        d      |      } |       ddd       d}| j                  |t        |j                               | j                  t              5 } t	        d      |      } |       ddd       d}| j                  |t        |j                               | j                  t              5 } t	        d      |      } |       ddd       d}| j                  |t        |j                               | j                  t              5 } t	        d      |      } |       ddd       d}| j                  |t        |j                               | j                  t              5 } t	        d      |      } |        ddd       d }| j                  |t        |j                               y# 1 sw Y   kxY w# 1 sw Y   xY w# 1 sw Y   xY w# 1 sw Y   uxY w# 1 sw Y   #xY w# 1 sw Y   xY w# 1 sw Y   ~xY w)"z3 Tests the large number of use cases defined below c                 @    | dk  r| S  | dz
         | dz
        z   S )Nr:   r   r   )nfib3s    r   r   z4TestInlinedClosure.test_bulk_use_cases.<locals>.fib3   s*    1uA;a!e,,r   c                     fd} ||       S )z, Test calling recursive function from inner c                      |       S r   r   )r   r   s    r   rc   zETestInlinedClosure.test_bulk_use_cases.<locals>.outer1.<locals>.inner   s    Awr   r   )r   rc   r   s     r   outer1z6TestInlinedClosure.test_bulk_use_cases.<locals>.outer1   s    8Or   c                 ,    | dz   fd} ||       S )z. Test calling recursive function from closure r   c                     |        z   S r   r   )r   r   r   s    r   rc   zETestInlinedClosure.test_bulk_use_cases.<locals>.outer2.<locals>.inner   s    47{"r   r   )r   rc   r   r   s     @r   outer2z6TestInlinedClosure.test_bulk_use_cases.<locals>.outer2   s    AA#8Or   c                     fd |       S )z Test recursive inner c                 (    | dk  ry | dz
         y Nr:   r   r   r   re   s    r   rc   zETestInlinedClosure.test_bulk_use_cases.<locals>.outer3.<locals>.inner   s    q5!a%Lr   r   re   s    @r   outer3z6TestInlinedClosure.test_bulk_use_cases.<locals>.outer3   s    !
 8Or   c                 ,    | dz   fd |       S )z Test recursive closure r   c                 .    | z   dk  ry | dz
         y r   r   rp   s    r   rc   zETestInlinedClosure.test_bulk_use_cases.<locals>.outer4.<locals>.inner   s    q519!a%Lr   r   rp   s    @@r   outer4z6TestInlinedClosure.test_bulk_use_cases.<locals>.outer4   s    AA!
 8Or   c                 (    | dz   fd} ||       S )z Test nested closure r   c                 6    | z   dz   fd} ||       z   S )Nr:   c                     | z   S r   r   r   r   s    r   inner2zVTestInlinedClosure.test_bulk_use_cases.<locals>.outer5.<locals>.inner1.<locals>.inner2   s    q5Lr   r   )r   r   r   ro   s     @r   inner1zFTestInlinedClosure.test_bulk_use_cases.<locals>.outer5.<locals>.inner1   s%    EAI! ay1}$r   r   r   r   ro   s     @r   outer5z6TestInlinedClosure.test_bulk_use_cases.<locals>.outer5   s    AA% !9r   c                 (    | dz   fd} ||       S )z. Test closure with list comprehension in body r   c                 P    | z   dz   }t        |      D cg c]  }| c}S c c}w )Nr:   r   )r   r   tro   s      r   r   zFTestInlinedClosure.test_bulk_use_cases.<locals>.outer6.<locals>.inner1   s,    EAI#(8,8a8,,,s   	#r   r   s     @r   outer6z6TestInlinedClosure.test_bulk_use_cases.<locals>.outer6   s    AA- !9r   	   c                     | dz   }| |z   z   S )z) Test use of outer scope var, no closure r   r   )r   r   _OUTER_SCOPE_VARs     r   outer7z6TestInlinedClosure.test_bulk_use_cases.<locals>.outer7   s    AAq5+++r   c                 ,    | dz   fd} ||       S )z+ Test use of outer scope var, with closure r   c                     | z   z   S r   r   )r   r   r   s    r   rc   zETestInlinedClosure.test_bulk_use_cases.<locals>.outer8.<locals>.inner  s    1u///r   r   )r   rc   r   r   s     @r   outer8z6TestInlinedClosure.test_bulk_use_cases.<locals>.outer8  s    AA08Or   c                 ,    | dz   fd}|} ||       S )z Test closure assignmentr   c                     | z   S r   r   r   s    r   rc   zETestInlinedClosure.test_bulk_use_cases.<locals>.outer9.<locals>.inner  ru   r   r   )r   rc   fr   s      @r   outer9z6TestInlinedClosure.test_bulk_use_cases.<locals>.outer9
  s    AAAQ4Kr   c                 4    | dz   fdfd} ||       S )z! Test two inner, one calls other r   c                     | z   S r   r   r   s    r   rc   zFTestInlinedClosure.test_bulk_use_cases.<locals>.outer10.<locals>.inner  ru   r   c                      |       S r   r   re   s    r   r   zGTestInlinedClosure.test_bulk_use_cases.<locals>.outer10.<locals>.inner2  s    Qxr   r   )r   r   rc   r   s     @@r   outer10z7TestInlinedClosure.test_bulk_use_cases.<locals>.outer10  s"    AA  !9r   c                     | dz   fd}|S )z return the closure r   c                     | z   S r   r   r   s    r   rc   zFTestInlinedClosure.test_bulk_use_cases.<locals>.outer11.<locals>.inner#  ru   r   r   r   rc   r   s     @r   outer11z7TestInlinedClosure.test_bulk_use_cases.<locals>.outer11  s    AALr   c                 *    | dz   dfd	} ||       S )z closure with kwargr   c                     | z   |z   S r   r   r   kwr   s     r   rc   zFTestInlinedClosure.test_bulk_use_cases.<locals>.outer12.<locals>.inner+      1urz!r      r   r   s     @r   outer12z7TestInlinedClosure.test_bulk_use_cases.<locals>.outer12'  s    AA"8Or   c                     | dz   |z   }|S )z outer with kwarg no closurer   r   r   s      r   outer13z7TestInlinedClosure.test_bulk_use_cases.<locals>.outer13/  s    A
AHr   c                 ,    | dz   fd} ||       S )z! outer with kwarg used in closurer   c                     | z   z   S r   r   r   s    r   rc   zFTestInlinedClosure.test_bulk_use_cases.<locals>.outer14.<locals>.inner8  r   r   r   r   r   rc   r   s    ` @r   outer14z7TestInlinedClosure.test_bulk_use_cases.<locals>.outer144  s    AA"8Or   c                 *    | dz   fd} || |      S )z# outer with kwarg as arg to closurer   c                     | z   |z   S r   r   r   s     r   rc   zFTestInlinedClosure.test_bulk_use_cases.<locals>.outer15.<locals>.inner@  r   r   r   r   s      @r   outer15z7TestInlinedClosure.test_bulk_use_cases.<locals>.outer15<  s    AA"B<r   c                 :    | dz   fd}t         ||             S )z( closure is generator, consumed locally r   c              3      K   | z    y wr   r   r   s    r   rc   zFTestInlinedClosure.test_bulk_use_cases.<locals>.outer16.<locals>.innerH       !e   	)listr   s     @r   outer16z7TestInlinedClosure.test_bulk_use_cases.<locals>.outer16D  s!    AA a>!r   c                 (    | dz   fd} ||       S )z  closure is generator, returned r   c              3      K   | z    y wr   r   r   s    r   rc   zFTestInlinedClosure.test_bulk_use_cases.<locals>.outer17.<locals>.innerQ  r   r   r   r   s     @r   outer17z7TestInlinedClosure.test_bulk_use_cases.<locals>.outer17M  s    AA 8Or   c                 :    | dz   fd} ||       D ]  }|} S )z( closure is generator, consumed in loop r   c              3      K   | z    y wr   r   r   s    r   rc   zFTestInlinedClosure.test_bulk_use_cases.<locals>.outer18.<locals>.innerZ  r   r   r   )r   rc   r   r   r   s       @r   outer18z7TestInlinedClosure.test_bulk_use_cases.<locals>.outer18V  s-    AA 1X  Hr   c                 @    | dz   | dz   fd}fd} |||       S )z# closure as arg to another closure r   r:   c                     | z   S r   r   )r   z1s    r   rc   zFTestInlinedClosure.test_bulk_use_cases.<locals>.outer19.<locals>.innerg  s    2vr   c                      | |      z   S r   r   )r   r   z2s     r   r   zGTestInlinedClosure.test_bulk_use_cases.<locals>.outer19.<locals>.inner2j  s    tby r   r   )r   rc   r   r   r   s      @@r   outer19z7TestInlinedClosure.test_bulk_use_cases.<locals>.outer19b  s.    QBQB! %##r   c                 (    | dz   fd} ||       S )z Test calling numpy in closure r   c                 4    | t        j                        z   S r   )numpycosr   s    r   rc   zFTestInlinedClosure.test_bulk_use_cases.<locals>.outer20.<locals>.inners  s    599Q<''r   r   r   s     @r   outer20z7TestInlinedClosure.test_bulk_use_cases.<locals>.outer20o  s    AA(8Or   c                 (    | dz   fd} ||       S )z) Test calling numpy import as in closure r   c                 4    | t        j                        z   S r   )npr   r   s    r   rc   zFTestInlinedClosure.test_bulk_use_cases.<locals>.outer21.<locals>.inner{  s    266!9}$r   r   r   s     @r   outer21z7TestInlinedClosure.test_bulk_use_cases.<locals>.outer21w  s    AA%8Or   c                      d } d} | |  y)z6Test to ensure that unsupported *args raises correctlyc                      y r   r   )rI   bs     r   barzDTestInlinedClosure.test_bulk_use_cases.<locals>.outer22.<locals>.bar  s    r   )r   r:   Nr   )r   r   s     r   outer22z7TestInlinedClosure.test_bulk_use_cases.<locals>.outer22  s    AGr   r   Tr+   Nz.Unsupported use of op_LOAD_CLOSURE encounteredz%Cannot capture the non-constant valuez-The use of yield in a closure is unsupported.z,Calling a closure with *args is unsupported.r   )
r   r   assertRaisesNotImplementedErrorr   assertInstr	exceptionr   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   refri   varraisesmsgr   r   s                                @@r   test_bulk_use_casesz&TestInlinedClosure.test_bulk_use_cases   sR    
	- 
	-
							 	,
 		
				
		 	"	
		$			 VVVVVWgwgw1 CIECU3ZS2  23v&C&v.E#J 4 ?c3v//0123v&C&v.E#J 4 ?c3v//01{+v&C&w/E#J , 6c3v//01/0F&C&w/E#J 1 >c3v//01/0F&C&w/E#J 1 >c3v//01/0F&C&w/E#J 1 >c3v//01/0F&C&w/EG 1 =c3v//01Q 43 43 ,+ 10 10 10 10sT   :L>M8MM%6M2M?4N>MMM"%M/2M<?NNc                    t        t              d        }| j                   |dd      d       |j                  |j                  d      j
                  d   }g }|j                  j                         D ]~  }|j                  D ]m  }t        |t        j                        st        |j                  t        j                        sC|j                  j                  dk(  s]|j                  |       o  | j                  t        |      d	       |D ],  }|j                   j"                  }d
}| j%                  ||       . y )N)pipeline_classc                 (    d } ||        ||      fS )Nc                     d}d}||z   | z   S )N   r   r   r   s      r   r   zITestInlinedClosure.test_closure_renaming_scheme.<locals>.foo.<locals>.bar  s    1uqy r   r   )rI   r   r   s      r   fooz<TestInlinedClosure.test_closure_renaming_scheme.<locals>.foo  s    ! q63q6>!r   r      )   #   r   preserved_irr   r:   zclosure__locals__bar_v[0-9]+.x)r   r	   r   	overloads
signaturesmetadatablocksvaluesbody
isinstancer   AssignvalueConstappendlentargetnameassertRegex)	r   r   func_irstoreblkstmtr   r  regexs	            r   test_closure_renaming_schemez/TestInlinedClosure.test_closure_renaming_scheme  s    
5	6	" 
7	" 	Rh/ --q 12;;NK>>((*CdBII.!$**bhh7::++q0!LL.	 ! + 	UQ'A88==D5ET5) r   c                     t         d        t         fd       }t        j                  j                   |       d       y )Nc                      dd} | S )Nc                     | |z   S r   r   rn   s     r   r   z;TestInlinedClosure.test_issue9222.<locals>.foo.<locals>.bar  rd   r   )g?r   )r   s    r   r   z.TestInlinedClosure.test_issue9222.<locals>.foo  s    Jr   c                                d      S rL   r   )r   s   r   consumez2TestInlinedClosure.test_issue9222.<locals>.consume  s    358Or   gffffff@)r   r   testingassert_allclose)r   r  r   s     @r   test_issue9222z!TestInlinedClosure.test_issue9222  sC     
	 
	
 
	 
	
 	

""79g6r   c                 *    t         d        } |        y )Nc                      d} t        d      D ]1  }t        j                  t        | d      D cg c]  }d c}       d} 3 y c c}w )Nr   r   r   )r   r   array)range_start_s     r   _innerz2TestInlinedClosure.test_issue_9577.<locals>._inner  sA    K1XU;%:;%:!%:;< ;s   	A
)r   )r   r   s     r   test_issue_9577z"TestInlinedClosure.test_issue_9577  s    		  
	  	r   N)r[   r\   r]   __doc__rj   rq   rv   r|   r   r   r  r  r   run_test_in_subprocessr!  r   r   r   r_   r_   a   sL    
///5(/&B2H *D7$ $$ %r   r_   __main__)r   r   unittestnumbar   r   numba.core.errorsr   r   
numba.corer   numba.tests.supportr   r	   r   r_   r[   mainr   r   r   <module>r+     sR        ;  BR3( R3jV Vr zHMMO r   