
    sh;^                     d   S SK r S SKJr  S SKrS SKJrJr  S SKJr  S SK	J
r
JrJrJrJrJrJr  S SKJrJr  S SKJr  S SKJrJrJrJr  S S	KJrJrJrJrJ r J!r!J"r"  S SK#r#S SK$r$S SK%r% " S
 S\&5      r'S r(S r)S r*S r+S r,\" \(5      r-S r.S r/S r0S r1S r2S r3S r4 " S S\5      r5 " S S\\5      r6 " S S\\5      r7 " S S\\%R0                  5      r8 " S  S!\\%R0                  5      r9 " S" S#\\%R0                  5      r: " S$ S%\\%R0                  5      r;\<S&:X  a  \%Rz                  " 5         gg)'    N)StringIO)njit	vectorize)typeof)utilstypestypingircompilercpucgutils)CompilerFlags)
cpu_target)MemoryLeakMixinTestCasetemp_directorycreate_temp_module)overloadmodelslower_builtinregister_modelmake_attribute_wrappertype_callabletypeof_implc                   (   ^  \ rS rSrU 4S jrSrU =r$ )	Namespace   c                 <   > X;   a  X   $ [         [        U ]  U5      $ N)superr   __getattr__)sk	__class__s     p/Users/tiagomarins/Projetos/claudeai/copy_bank/venv/lib/python3.13/site-packages/numba/tests/test_array_exprs.pyr"   Namespace.__getattr__   s     vqtE5A#B1#EE     )__name__
__module____qualname____firstlineno__r"   __static_attributes____classcell__)r%   s   @r&   r   r      s    F Fr(   r   c                     X-  U-   $ r    r)   axys      r&   axyr5           519r(   c                     X-  U-   $ r    r)   r1   s      r&   ax2r8   #   r6   r(   c                 8    U* US-  SU -  U-  -
  S-  -   SU -  -  $ )N       @g      @      ?r)   AsBsCss      r&   pos_rootr@   &   s.    CR2X"r'B,/C78R"WEEr(   c                 D    SU -  nSU-  U-  nUS-  U-
  nU* US-  -
  U-  $ )Nr:   r;   r)   r=   r>   r?   _2As_4AsCs
_Bs2_4AsCss         r&   neg_root_common_subexprrF   )   s=    7D$Y^F(V#JC:$%--r(   c                 J    SU -  nSU-  U-  nUS-  U-
  S-   nU* US-  -
  U-  $ )Nr:   y                r;   r)   rB   s         r&   neg_root_complex_subexprrH   /   sB    7D$Y^F(V#r)JC:$%--r(   c                 r    [         R                  " [        U [         R                  " U5      S-
  S5      5      $ )N      ?)npcosvaxysin)a0a1s     r&   
call_stuffrQ   7   s&    66$r266":?B/00r(   c                 $    US-  SU -  U-  -
  S:  $ )N      r   r)   r<   s      r&   are_roots_imaginaryrU   :   s    !Ga"frk!Q&&r(   c                     X-  U-   $ r    r)   r<   s      r&   div_addrW   =   s    7R<r(   c                     U S-  $ )N   r)   )r=   s    r&   cuberZ   @   s    7Nr(   c                 \    [         R                  " X5        [         R                  " X!U5      $ r    )rK   rL   add)r2   bouts      r&   explicit_outputr_   C   s    FF1N66##r(   c                 *    X-   nX@U-  -
  nXB-  U-   nU$ r    r)   )r2   r]   cdus        r&   variable_name_reuserd   G   s$    	A	E	A		AHr(   c                    U R                   S   n[        R                  " X4[        R                  S9n[	        U5       HB  n[	        X15       H0  n[        R
                  " X   X   -
  S-  5      S-  =X#U4'   X$U4'   M2     MD     U$ )Nr   dtyperS   r;   )shaperK   emptyfloat64rangesum)vectors	n_vectorsresultijs        r&   distance_matrixrr   O   s    a IXXy,BJJ?F9q$A(*gj(Q.)036)7 7FQ3K&1+ % 
 Mr(   c                   P    \ rS rSr\SS0 SSS4S j5       r\SS0 S4S j5       rSrg)RewritesTester[   Nc           	          U(       d
  [        5       nSUl        Uc  [        R                  nUc  [        R                  nU " XgXQUX45      $ NT)r   nrtr   typing_contexttarget_context)clsargsreturn_typeflagslocalslibraryry   rz   s           r&   mk_pipelineRewritesTester.mk_pipeline\   sM     GE	!'66N!'66N>7+" 	"r(   c                 \    U(       d
  [        5       nSUl        U R                  " XX4U40 UD6$ rw   )r   no_rewritesr   )r{   r|   r}   r~   r   r   kwss          r&   mk_no_rw_pipeline RewritesTester.mk_no_rw_pipelinei   s.     GE t%PCPPr(   r)   )r*   r+   r,   r-   classmethodr   r   r.   r)   r(   r&   rt   rt   [   sB    +/tB d
" 
" 15T""&Q Qr(   rt   c                       \ rS rSrS rS rS rSS jr\4S jr	\
4S jrS	 rS
 rSS jrS rS rS r\4S jrS rS rS rS rSrg)TestArrayExpressionsr   c                     [         R                  U5      nUR                  U5      nUR                  n[         R	                  U5      nUR                  U5      nUR                  nX5Xh4$ )zD
Compile the given function both without and with rewrites enabled.
)rt   r   compile_extraentry_pointr   )	selffnarg_tyscontrol_pipelinecres_0control_cfunctest_pipelinecres_1
test_cfuncs	            r&   _compile_function&TestArrayExpressions._compile_functiont   sf     *;;GD!//3**&227;,,R0''
IIr(   c                 p   [         R                  " SSS5      n[         R                  " SSS5      n[         R                  " SSS5      nXU4 Vs/ s H  n[        U5      PM     nnU R                  [        U5      u  pgp[
        R                  U5      n
U
R                  [        5      nUR                  nU" XU5      nU	" XU5      nU" XU5      n[         R                  R                  X5        [         R                  R                  X5        UR                  R                  R                  nUR                  R                  R                  nU
R                  R                  R                  nU R                  [!        U5      [!        U5      5        U R                  [!        U5      [!        U5      5        U R#                  [!        US   R$                  5      [!        US   R$                  5      5        U R                  [!        US   R$                  5      [!        US   R$                  5      5        gs  snf )z^
Using a simple array expression, verify that rewriting is taking
place, and is fusing loops.
r      
   rS   N)rK   linspacer   r   r5   rt   r   r   r8   r   testingassert_array_equalstatefunc_irblocksassertEquallenassertGreaterbody)r   AXYargr   r   nb_axy_0r   nb_axy_1control_pipeline2cres_2nb_ctlexpectedactualcontrolir0ir1ir2s                      r&   test_simple_expr%TestArrayExpressions.test_simple_expr   s   
 KK!BKK!BKK!B+,)4)36#;)4 ""30 	<M +<<WE"005##A!$!"q/


%%h7


%%g6$$,,33!!))00%%--44S3s8,S3s8,3s1v{{+SQ-=>SV[[)3s1v{{+;<- 5s   H3c              #      #    U Ho  n[        U[        R                  5      (       d  M$  [        UR                  [        R                  5      (       d  MO  UR                  R
                  S:X  d  Mk  Uv   Mq     g 7f)N	arrayexpr)
isinstancer
   AssignvalueExprop)r   blockinstrs      r&   _get_array_exprs%TestArrayExpressions._get_array_exprs   sJ     E%++ekk27733{{~~4#	 s   #A9'A9A90	A9Nc                 L   Uc
  [        5       n[        U[        5      (       d  [        SR	                  U5      5      eUu  p4/ nU H<  n[        U[        5      (       a  U R                  Xb5      u  pgUR                  U5        M>     U[        U5      4nUR                  U5        X4$ )z;
Convert an array expression tree into a set of operators.
z{0} not a tuple)setr   tuple
ValueErrorformat_array_expr_to_setappendr\   )	r   exprr^   	operationoperandsprocessed_operandsoperand_processed_exprs	            r&   r   'TestArrayExpressions._array_expr_to_set   s     ;%C$&&.55d;<<"	G'5))!44WB
%%g.   #E*<$==""r(   c                    [         R                  R                  S5      n[         R                  R                  S5      S-   n[         R                  R                  S5      nX#U4 Vs/ s H  n[        U5      PM     nn[        R	                  U5      nUR                  U5      nUR                  n	[        R                  U5      n
U
R                  U5      nUR                  nU" X#U5      nU	" X#U5      nU" X#U5      n[         R                  R                  X5        [         R                  R                  X5        [        [        5       5      $ s  snf )Nr   rJ   )rK   randomr   rt   r   r   r   r   r   assert_array_almost_equalr   r   )r   r   r   BCr   r   r   control_cresnb_fn_0r   	test_cresnb_fn_1	np_resultnb_result_0nb_result_1s                   r&   _test_root_function(TestArrayExpressions._test_root_function   s   IIR IIR 2% IIR +,)4)36#;)4);;GD'55b9**&227;!//3	''qQK	aA&aA&


,,YD


,,[F""! 5s   'Ec                    [         R                  " S[         R                  S9n[        U5      4n[        R                  U5      nUR                  U5      nUR                  n[        R                  U5      nUR                  U5      nUR                  n	US-  n
U R                  X" U5      5        U R                  X" U5      5        [        [        5       5      $ )Nr   rf   rY   )rK   arangerj   r   rt   r   r   r   r   assertPreciseEqualr   r   )r   r   r   r   r   r   r   r   r   r   r   s              r&   _test_cube_function(TestArrayExpressions._test_cube_function   s    IIb

+!9,);;GD'55b9**&227;!//3	''6'!*5'!*5""r(   c                 <  ^	^
 [         R                  " S[         R                  S9m	T	S-   m
[        T	5      4S-  nU R	                  X5      u  p4pVU	U
4S jnU" U5      nU R                  X" U5      5        U R                  X" U5      5        [        [        5       5      $ )zg
Test function having a (a, b, out) signature where *out* is
an output array the function writes into.
r   rf   r   rY   c                 H   > [         R                  " T5      nU " TTU5        U$ r    )rK   
zeros_like)r   r^   r   r   s     r&   run_funcETestArrayExpressions._test_explicit_output_function.<locals>.run_func   s!    --"Cq!SMJr(   )rK   r   rj   r   r   r   r   r   )r   r   r   r   r   r   r   r   r   r   r   s            @@r&   _test_explicit_output_function3TestArrayExpressions._test_explicit_output_function   s    
 IIb

+E!9," ""2/ 	C	
 B<(=*AB(:*>?""r(   c                 l    [        [        U R                  U5      5      5      nU R                  X25        g)zH
Assert the *block* has the expected number of array expressions
in it.
N)r   listr   r   )r   r   expected_countrewrite_counts       r&   _assert_array_exprs(TestArrayExpressions._assert_array_exprs   s,    
 D!6!6u!=>?7r(   c                 (   U R                  [        U5      [        U5      5        US   R                  nUS   R                  nU R                  US5        U R                  US5        U(       d%  U R	                  [        U5      [        U5      5        gg)z
Given two dictionaries of Numba IR blocks, check to make sure the
control IR has no array expressions, while the test IR
contains one and only one.
r   r   N)r   r   r   r   r   )r   
control_irtest_irtrivialcontrol_block
test_blocks         r&   _assert_total_rewrite*TestArrayExpressions._assert_total_rewrite  s{     	Z#g,7"1**QZ__
  2  Q/ s=13z?C	 r(   c                 B   U R                  [        U5      [        U5      5        UR                  5        Hg  u  p4UR                  nX#   R                  nU R                  [        U5      [        U5      5        U R	                  US5        U R	                  US5        Mi     g)z~
Given two dictionaries of Numba IR blocks, check to make sure
the control IR and the test IR both have no array expressions.
r   N)r   r   itemsr   r   )r   r   r   r$   vr   r   s          r&   _assert_no_rewrite'TestArrayExpressions._assert_no_rewrite  s~    
 	Z#g,7$$&DAFFM JS/ZA$$]A6$$Z3 'r(   c                     U R                  5       nU R                  UR                  R                  R                  R
                  UR                  R                  R                  R
                  SS9  g)zu
Ensure even a non-nested expression is rewritten, as it can enable
scalar optimizations such as rewriting `x ** 2`.
T)r   N)r   r   r   r   r   r   r   r   nss     r&   test_trivial_expr&TestArrayExpressions.test_trivial_expr$  s]    
 %%'""2#6#6#<#<#D#D#K#K#%#3#3#9#9#A#A#H#H+/ 	# 	1r(   c                     U R                  5       nU R                  UR                  R                  R                  R
                  UR                  R                  R                  R
                  5        g)z
Using the polynomial root function, ensure the full expression is
being put in the same kernel with no remnants of intermediate
array expressions.
N)r   r   r   r   r   r   r   r   s     r&   test_complicated_expr*TestArrayExpressions.test_complicated_expr.  sV     %%'""2#6#6#<#<#D#D#K#K#%#3#3#9#9#A#A#H#H	Jr(   c           
      l  ^  T R                  U5      nUR                  R                  R                  R                  nUR
                  R                  R                  R                  nT R                  [        U5      [        U5      5        T R                  [        US   R                  5      [        US   R                  5      5        T R                  [        [        T R                  US   R                  5      5      5      S5        [        T R                  US   R                  5      5      nT R                  [        U5      S5        [        U 4S jU 5       5      n[        USS USS 5       H?  u  pxUR                  U5      n	U	(       d  M  T R                  SR                  U	5      5        MA     g)zX
Attempt to verify that rewriting will incorporate user common
subexpressions properly.
r   r   c              3   t   >#    U  H-  nTR                  UR                  R                  5      S    v   M/     g7f)r   N)r   r   r   ).0r   r   s     r&   	<genexpr>BTestArrayExpressions.test_common_subexpressions.<locals>.<genexpr>K  s2      :'8e 11%++2B2BCAF'8s   58Nz9Common subexpressions detected in array expressions ({0}))r   r   r   r   r   r   r   r   r   r   r   r   zipintersectionfailr   )
r   r   r  r   r   array_expr_instrs
array_sets
expr_set_0
expr_set_1intersectionss
   `         r&   test_common_subexpressions/TestArrayExpressions.test_common_subexpressions8  s_   
 %%b)!!''//66$$,,33S3s8,3s1v{{+SQ-=>T$"7"7A"DEFJ
 !!6!6s1v{{!CD30115  :'8: :
&)*Sb/:ab>&J"J&33J?M}		 ..4f].CE 'Kr(   c                 ,    U R                  [        5      $ r    )r  rH   )r   s    r&   test_complex_subexpression/TestArrayExpressions.test_complex_subexpressionS  s    ../GHHr(   c                 l   [         R                  R                  S5      n[         R                  R                  S5      nX4 Vs/ s H  n[        U5      PM     nn[        R                  R
                  n[        R                  UUR                  UR                  S9nUR                  [        5      nUR                  n[        R                  UUR                  UR                  S9n	U	R                  [        5      n
U
R                  n[        X5      nU" X5      nU" X5      n[         R                  R                  X5        [         R                  R                  X5        U R!                  UR"                  R$                  R&                  U	R"                  R$                  R&                  5        gs  snf )zV
Verify that ufunc and DUFunc calls are being properly included in
array expressions.
r   )ry   rz   N)rK   r   r   rM   _dispatchertargetdescrrt   r   ry   rz   r   rQ   r   r   r   r   r   r   r   r   )r   r   r   r   r   
vaxy_descrr   r   nb_call_stuff_0r   r   nb_call_stuff_1r   r   r   s                  r&   test_ufunc_and_dufunc_calls0TestArrayExpressions.test_ufunc_and_dufunc_callsV  sf   
 IIR IIR +,&1&36#;&1%%11
);;%44%44 < 6 "//
; ,,&22%44%44 3 6 ,,Z8 ,,a#!!' &


,,X?


,,X>""#3#9#9#A#A#H#H#0#6#6#>#>#E#E	G/ 2s   F1c                     U R                  [        5      nU R                  UR                  R                  R
                  R                  UR                  R                  R
                  R                  5        g)zA
Verify that comparison operators are supported by the rewriter.
N)r   rU   r   r   r   r   r   r   r   s     r&   test_cmp_op TestArrayExpressions.test_cmp_opw  sY     %%&9:""2#6#6#<#<#D#D#K#K#%#3#3#9#9#A#A#H#H	Jr(   c                     U R                  [        5      nU R                  UR                  R                  R
                  R                  UR                  R                  R
                  R                  5        g)zA
Check that ufunc calls with explicit outputs are not rewritten.
N)r   r_   r   r   r   r   r   r   r   s     r&   test_explicit_output)TestArrayExpressions.test_explicit_output  sX     00A 3 3 9 9 A A H H " 0 0 6 6 > > E E	Gr(   r)   r    )F)r*   r+   r,   r-   r   r   r   r   r@   r   rZ   r   r   r   r   r   r  r  rF   r  r  r   r#  r&  r.   r)   r(   r&   r   r   r   sq    J=@$#$ &. #0 &* #$#.8D$41J -D E6IGBJGr(   r   c                   >    \ rS rSrS rS rS rS rS rS r	S r
S	rg
)TestRewriteIssuesi  c                    ^ SSK Jn  SS KnU" SS9S 5       mU" SS9U4S j5       nUR                  " S5      nU" U5      u  pVU R	                  XV5        U R	                  XT5        g )Nr   )jitT)nopythonc                     U $ r    r)   )arrs    r&   foo.TestRewriteIssues.test_issue_1184.<locals>.foo  s    Jr(   c                 *   > T" U 5      nT" U 5      nX4$ r    r)   )r.  ra   rb   r/  s      r&   bar.TestRewriteIssues.test_issue_1184.<locals>.bar  s    CACA4Kr(   r   )numbar+  numpyr   assertIs)r   r+  rK   r2  r.  out_cout_dr/  s          @r&   test_issue_1184!TestRewriteIssues.test_issue_1184  sk    	d		 
	 
d		 
	
 iim3xe#e!r(   c                    Sn[         R                  R                  US-  S9R                  US45      n[	        U5      n[        [        5      " U5      n[         R                  R                  X45        [        R                  " 5         g )Nd   rY   )size)
rK   r   uniformreshaperr   r   r   r   gccollect)r   nr3   r   r   s        r&   test_issue_1264!TestRewriteIssues.test_issue_1264  se    II1Q3'//16"1%o&q)


,,X>


r(   c                    SSK Jn  US 5       n[        R                  R	                  S5      n[        R                  R	                  S5      nUR                  X45      nU" X45      n[        R                  R                  Xe5        g)z*Test array expression with duplicated termr   )r   c                 <    [         R                  " U5      nX-   U -   $ r    )rK   rN   r2   r]   s     r&   r/  .TestRewriteIssues.test_issue_1372.<locals>.foo  s    q	A519r(   r   N)r4  r   rK   r   r>  py_funcr   assert_allclose)r   r   r/  r2   r]   expectgots          r&   test_issue_1372!TestRewriteIssues.test_issue_1372  sj    		 
	 IIb!IIb!Q"!i


""3/r(   c                     [         S 5       nSn[        R                  " S[        R                  S9nUR	                  X25      nU" X25      nU R                  XT5        g)z@
Typing of unary array expression (np.negate) can be incorrect.
c                     X-
  U * -   $ r    r)   rG  s     r&   r/  3TestRewriteIssues.test_unary_arrayexpr.<locals>.foo  s    5A2:r(   g      ?r   rf   N)r   rK   r   int32rI  r   )r   r/  r]   r2   rK  rL  s         r&   test_unary_arrayexpr&TestRewriteIssues.test_unary_arrayexpr  sV     
	 
	 IIb)Q"!i,r(   c                     [         S 5       n[        R                  " / SQ5      n[        R                  " / SQ5      nUR                  X#5      nU" X#5      nU R	                  XT5        g)zL
Typing of bitwise boolean array expression can be incorrect
(issue #1813).
c                     X) -  ) $ r    r)   rG  s     r&   r/  5TestRewriteIssues.test_bitwise_arrayexpr.<locals>.foo  s    "X;r(   )TTFF)FTFTN)r   rK   arrayrI  r   )r   r/  r2   r]   rK  rL  s         r&   test_bitwise_arrayexpr(TestRewriteIssues.test_bitwise_arrayexpr  s[    
 
	 
	 HH/0HH/0Q"!i,r(   c                    [        [        5      n[        R                  " SSS5      nU" X"X"5        [	        5       nUR                  U5        UR                  5       nU R                  SU5        U R                  SU5        g)zW
Type annotation of array expressions with disambiguated
variable names (issue #1466).
r   r   r   z
#   u.1 = z
#   u.2 = N)r   rd   rK   r   r   inspect_typesgetvalueassertIn)r   cfuncr2   bufress        r&   test_annotations"TestRewriteIssues.test_annotations  sg    
 ()KK1b!aAjC llnlC(lC(r(   c                     [         S 5       nU" S5      nUR                  S5      n[        R                  R	                  X#5        g )Nc                 ~    [         R                  " U 5      n[        S5       H  nX-  nUR                  5       nM     U$ )NrS   )rK   onesrk   copy)r3   r.  r   vals       r&   f;TestRewriteIssues.test_issue_5599_name_collision.<locals>.f  s6    ''!*C1Xyhhj  Jr(      )r   rI  rK   r   r   )r   ri  rL  rK  s       r&   test_issue_5599_name_collision0TestRewriteIssues.test_issue_5599_name_collision  s>    		 
	 d1


%%c2r(   r)   N)r*   r+   r,   r-   r9  rC  rM  rS  rY  rb  rl  r.   r)   r(   r&   r)  r)    s%    "&0-- ) 3r(   r)  c                       \ rS rSrS rSrg)TestSemanticsi  c                 (   [         n[        U5      n[        R                  " SS[	        S5      /5      n[        R                  " / SQ5      n[        R
                  " U5      nU" X4U5      nU" X4U5      n[        R                  R                  Xg5        g )N        rJ   inf)rq  rq  rJ   )rW   r   rK   rj   float	ones_liker   r   )r   pyfuncr_  r2   r]   ra   rK  rL  s           r&   test_division_by_zero#TestSemantics.test_division_by_zero  sn     VJJS%,/0JJ'LLOaA!n


%%f2r(   r)   N)r*   r+   r,   r-   rv  r.   r)   r(   r&   ro  ro    s    3r(   ro  c                   $    \ rS rSrSrS rS rSrg)TestOptionalsi  z`Tests the arrival and correct lowering of Optional types at a arrayexpr
derived ufunc, see #3972c                   ^ [         S 5       m[         U4S j5       n[        R                  " S5      S4nU" U6 nUR                  " U6 n[        R                  R                  X45        TR                  nUS   S   nU R                  [        U[        R                  5      5        U R                  [        UR                  [        R                  5      5        g )Nc                 
    X-   $ r    r)   r3   r4   s     r&   arr_expr9TestOptionals.test_optional_scalar_type.<locals>.arr_expr  	    5Lr(   c                 *   > US:  a  S nOUnT" X5      $ Nr   r)   r3   r4   zr}  s      r&   do_call8TestOptionals.test_optional_scalar_type.<locals>.do_call      1uA>!r(   rk  g333333r   r   )r   rK   r   rI  r   rJ  
signatures
assertTruer   r   OptionaltypeFloatr   r  r|   ra  r   r#   otyr}  s          @r&   test_optional_scalar_type'TestOptionals.test_optional_scalar_type	  s    		 
	 
	" 
	" 		!d# tn??D)


""31 d1g
378
388U[[9:r(   c                 T  ^ [         S 5       m[         U4S j5       n[        R                  " S5      [        R                  " S5      4nU" U6 nUR                  " U6 n[        R                  R                  X45        TR                  nUS   S   nU R                  [        U[        R                  5      5        U R                  [        UR                  [        R                  5      5        U R                  [        UR                  R                  [        R                  5      5        g )Nc                 
    X-   $ r    r)   r|  s     r&   r}  8TestOptionals.test_optional_array_type.<locals>.arr_expr&  r  r(   c                 0   > US   S:  a  S nOUnT" X5      $ r  r)   r  s      r&   r  7TestOptionals.test_optional_array_type.<locals>.do_call*  "    taxA>!r(   rk        @r   r   )r   rK   r   rI  r   rJ  r  r  r   r   r  r  Arrayrg   r  r  s          @r&   test_optional_array_type&TestOptionals.test_optional_array_type$  s    		 
	 
	" 
	" 		!biim, tn??D)


""31 d1g
378
388U[[9:
388>>5;;?@r(   r)   N)r*   r+   r,   r-   __doc__r  r  r.   r)   r(   r&   ry  ry    s     ;6Ar(   ry  c                        \ rS rSrS rS rSrg)TestOptionalsExceptionsiA  c                   ^ U R                  5         [        S 5       m[        U4S j5       n[        R                  " S5      S4nU R	                  [
        5       nU" U6   S S S 5        U R                  S[        WR                  5      5        TR                  nUS   S   nU R                  [        U[        R                  5      5        U R                  [        UR                  [        R                  5      5        g ! , (       d  f       N= f)Nc                 
    X-   $ r    r)   r|  s     r&   r}  UTestOptionalsExceptions.test_optional_scalar_type_exception_on_none.<locals>.arr_exprH  r  r(   c                 *   > US:  a  S nOUnT" X5      $ r  r)   r  s      r&   r  TTestOptionalsExceptions.test_optional_scalar_type_exception_on_none.<locals>.do_callL  r  r(   rk  rJ   zexpected float64, got Noner   r   )disable_leak_checkr   rK   r   assertRaises	TypeErrorr^  str	exceptionr  r  r   r   r  r  r  )r   r  r|   raisesr#   r  r}  s         @r&   +test_optional_scalar_type_exception_on_noneCTestOptionalsExceptions.test_optional_scalar_type_exception_on_noneD  s    !		 
	 
	" 
	" 		!c" y)VTN * 	2C8H8H4IJ d1g
378
388U[[9: *)s   C>>
Dc                   ^ U R                  5         [        S 5       m[        U4S j5       n[        R                  " S5      [        R                  " SS5      4nU R	                  [
        5       nU" U6   S S S 5        [        WR                  5      nU R                  SU5        U R                  SU5        TR                  nUS   S	   nU R                  [        U[        R                  5      5        U R                  [        UR                  [        R                  5      5        U R                  [        UR                  R                   [        R"                  5      5        g ! , (       d  f       N= f)
Nc                 
    X-   $ r    r)   r|  s     r&   r}  TTestOptionalsExceptions.test_optional_array_type_exception_on_none.<locals>.arr_exprf  r  r(   c                 0   > US   S:  a  S nOUnT" X5      $ r  r)   r  s      r&   r  STestOptionalsExceptions.test_optional_array_type_exception_on_none.<locals>.do_callj  r  r(   rk  rJ   r  zexpected array(float64,zgot Noner   r   )r  r   rK   r   r  r  r  r  r^  r  r  r   r   r  r  r  rg   r  )r   r  r|   r  excstrr#   r  r}  s          @r&   *test_optional_array_type_exception_on_noneBTestOptionalsExceptions.test_optional_array_type_exception_on_noneb  s   !		 
	 
	" 
	" 		!biiB/0 y)VTN * V%%&/8j&) d1g
378
388U[[9:
388>>5;;?@ *)s   /E$$
E2r)   N)r*   r+   r,   r-   r  r  r.   r)   r(   r&   r  r  A  s    ;<Ar(   r  c                   H    \ rS rSrSr\R                  " S5      rS rS r	Sr
g)TestExternalTypesi  zETests RewriteArrayExprs with external (user defined) types,
see #5157z
        from numba.core import types

        class FooType(types.Type):
            def __init__(self):
                super(FooType, self).__init__(name='Foo')
        c                 8  ^  " S S[         5      n[        T5       " S S[        R                  5      5       n[	        TSS5        [        U5      U4S j5       n[        U[        R                  5      S 5       n[        R                  " U5      U4S j5       nUT4$ )	Nc                       \ rS rSrS rSrg),TestExternalTypes.make_foo_type.<locals>.Fooi  c                     Xl         g r    r   )r   r   s     r&   __init__5TestExternalTypes.make_foo_type.<locals>.Foo.__init__  s    "
r(   r  Nr*   r+   r,   r-   r  r.   r)   r(   r&   Foor    s    #r(   r  c                       \ rS rSrS rSrg)1TestExternalTypes.make_foo_type.<locals>.FooModeli  c                 j    S[         R                  4/n[        R                  R	                  XX#5        g )Nr   )r   intpr   StructModelr  )r   dmmfe_typememberss       r&   r  :TestExternalTypes.make_foo_type.<locals>.FooModel.__init__  s)    #UZZ01""++DwHr(   r)   Nr  r)   r(   r&   FooModelr    s    Ir(   r  r   c                    > U4S jnU$ )Nc                    > T" 5       $ r    r)   )r   FooTypes    r&   typer@TestExternalTypes.make_foo_type.<locals>.type_foo.<locals>.typer  s    y r(   r)   )contextr  r  s     r&   type_foo1TestExternalTypes.make_foo_type.<locals>.type_foo  s    ! Lr(   c                     UR                   nUu  n[        R                  " U5      " X5      nXVl        UR	                  5       $ r    )r}   r   create_struct_proxyr   	_getvalue)r  buildersigr|   typr   r/  s          r&   impl_foo1TestExternalTypes.make_foo_type.<locals>.impl_foo  s8    //CGU--c27DCI==?"r(   c                    > T" 5       $ r    r)   )rh  ra   r  s     r&   
typeof_foo3TestExternalTypes.make_foo_type.<locals>.typeof_foo  s    9r(   )objectr   r   r  r   r   r   r   r  r   register)r   r  r  r  r  r  r  s    `     r&   make_foo_typeTestExternalTypes.make_foo_type  s    	#& 	# 
	 	Iv)) 	I 
!	I
 	w9	s		 
	 
sEJJ	'	# 
(	# 
		c	"	 
#	 G|r(   c                 2  ^^ [        U R                  5       nU R                  UR                  5      u  mm[	        [
        R                  5      U4S j5       n[	        [
        R                  5      U4S j5       n[	        [
        R                  5      U4S j5       n[        U4S j5       n[        U4S j5       n[        U4S j5       n[        R                  R                  U" SS5      [        R                  " SS/5      5        [        R                  R                  U" S	S
5      [        R                  " S	S
/5      5        [        R                  R                  U" S5      [        R                  " S/5      5        S S S 5        g ! , (       d  f       g = f)Nc                 r   > [        U T5      (       a%  [        U[        R                  5      (       a  S nU$ g g )Nc                 L    [         R                  " U R                  US   /5      $ r  rK   rX  r   lhsrhss     r&   impKTestExternalTypes.test_external_type.<locals>.overload_foo_add.<locals>.imp  s    !xxCF(;<<r(   )r   r   r  r  r  r  r  s      r&   overload_foo_add>TestExternalTypes.test_external_type.<locals>.overload_foo_add  s5    c7++
30L0L= J	 1M+r(   c                 V   > [        U T5      (       a  [        UT5      (       a  S nU$ g g )Nc                 Z    [         R                  " U R                  UR                  /5      $ r    r  r  s     r&   r  r    s    !xxCII(>??r(   r   r  s      r&   r  r    s2    c7++
30H0H@ J	 1I+r(   c                 2   > [        U T5      (       a  S nU$ g )Nc                 F    [         R                  " U R                  * /5      $ r    r  )r3   s    r&   r  KTestExternalTypes.test_external_type.<locals>.overload_foo_neg.<locals>.imp  s    !xx!''
33r(   r  )r3   r  r  s     r&   overload_foo_neg>TestExternalTypes.test_external_type.<locals>.overload_foo_neg  s!    a))4 J	 *r(   c                 D   > T" U 5      [         R                  " U/5      -   $ r    )rK   rX  r3   r4   r  s     r&   arr_expr_sum1;TestExternalTypes.test_external_type.<locals>.arr_expr_sum1  s    1v!--r(   c                 &   > T" U 5      T" U5      -   $ r    r)   r  s     r&   arr_expr_sum2;TestExternalTypes.test_external_type.<locals>.arr_expr_sum2  s    1vA&r(   c                    > T" U 5      * $ r    r)   )r3   r  s    r&   arr_expr_neg:TestExternalTypes.test_external_type.<locals>.arr_expr_neg  s    Awr(   r   r   rS   rY   rT   )r   source_linesr  r  r   operatorr\   negr   rK   r   r   rX  )	r   test_moduler  r  r  r  r  r  r  s	          @@r&   test_external_type$TestExternalTypes.test_external_type  sA    1 12k--k.A.ABLC hll# $ hll# $ hll# $ . . ' '   JJ))-1*=rxxA?OPJJ))-1*=rxxA?OPJJ)),q/288RD>J_ 322s   E'F
Fr)   N)r*   r+   r,   r-   r  textwrapdedentr  r  r  r.   r)   r(   r&   r  r    s(     ?? $ L D0Kr(   r  __main__)>r@  ior   r5  rK   r4  r   r   r   
numba.corer   r   r	   r
   r   r   r   numba.core.compilerr   r   numba.core.registryr   numba.tests.supportr   r   r   r   numba.extendingr   r   r   r   r   r   r   r  r  unittestdictr   r5   r8   r@   rF   rH   rM   rQ   rU   rW   rZ   r_   rd   rr   rt   r   r)  ro  ry  r  r  r*   mainr)   r(   r&   <module>r     sD   	   !  G G G / *5 5     F FF.. ~1'	QX Q.SG?H SGli3 i3X3OX%6%6 3"9AOX%6%6 9Ax@Aox/@/@ @AF^K):): ^KB zMMO r(   