o
    iUB                     @   st  d dl Zd dlZd dlmZmZ d dlmZ d dlm	Z	m
Z
 d dlmZ dddZd	diZd
d Zdd Zdd Zdd Zdd Zdd Zdd Zdd ZdAddZdd Zd d! Zd"d# Zd$d% Zd&d' Zd(d) Zd*d+ Zd,d- Zd.d/ Z d0d1 Z!G d2d3 d3e
e	Z"d4d5 Z#d6d7 Z$G d8d9 d9e
e	Z%G d:d; d;e	Z&G d<d= d=e
e	Z'G d>d? d?e Z(e)d@kre*  dS dS )B    N)jitnjit)types)TestCaseMemoryLeakMixin)test_factoryFT)nopythonforceobjr   c                        fdd}|S )Nc                    s   d} | D ]}||7 }q|S )Ng         )xresygen_funcr   X/home/ubuntu/veenaModal/venv/lib/python3.10/site-packages/numba/tests/test_generators.pyconsumer   s   
zmake_consumer.<locals>.consumerr   )r   r   r   r   r   make_consumer      r   c                 c   s    t | D ]}|V  qd S Nrange)r   ir   r   r   gen1   s   r   c                 c   s4    t | D ]}|V  t ddD ]}|| V  qqd S )N      r   )r   r   jr   r   r   gen2   s   r   c                 c   s     | V  | d V  | d V  d S )Ng      ?              ?r   r   r   r   r   gen3#   s   
r    c                 c   s$    t dD ]
}|V  || V  qd S )Nr   r   )r   r   zr   r   r   r   gen4*   s
   r"   c                   c   s    d S r   r   r   r   r   r   gen52   s   r#   c                 c   s     | d }	 |d }|| V  q)Nr   T   r   )abr   r   r   r   r   gen6:   s   
r'   c                 c   s     t | jD ]}| | V  qd S r   r   sizearrr   r   r   r   gen7B   s   r,   r   r$   c                 c   s.    | }| V  |r|V  |r| | V  d S d S r   r   )r   r   r&   bbr   r   r   gen8I   s   r.   c                 c   s    t   | V  d S r   )objectr   r   r   r   genobjR      
r0   c                 C   s   dd | D S )Nc                 s   s    | ]}|d  V  qdS )r$   Nr   ).0r   r   r   r   	<genexpr>X   s    z(return_generator_expr.<locals>.<genexpr>r   r   r   r   r   return_generator_exprW      r4   c                 c       t | D ]}|V  qd S r   )npndindex)shapeindr   r   r   gen_ndindex[      r;   c                 c   s    | j D ]}|V  qd S r   )flat)r+   valr   r   r   gen_flat`   s   
r?   c                 c   r6   r   )r7   ndenumerate)r+   tupr   r   r   gen_ndenumeratee   r<   rB   c                   c   s    dV  d S )NTr   r   r   r   r   gen_boolj      
rC   c                   c   s    d V  dV  d S )Nr   r   r   r   r   r   gen_unification_errorn   r1   rE   c                  c   s    d} dV  	 | V } q)Nr   r   r   )r   r   r   r   'gen_optional_and_type_unification_errors   s   rF   c                   c   s    dV  dV  d S )N)r   r$   )r      r   r   r   r   r   gen_changing_tuple_type{   s   
rH   c                   c   s    dV  dV  dV  d S )Nr   g      @y        33333P@r   r   r   r   r   gen_changing_number_type   s   
rI   c                   @   sd  e Zd Zdd Zdd Zdd Zdd Zd	d
 Zdd Zdd Z	dd Z
dd Zdd Zdd Zdd Zdd Zdd Zdd Zdd  Zd!d" Zd#d$ Zd%d& Zd'd( Zd)d* Zd+d, Zd-d. Zd/d0 Zd1d2 Zd3d4 Zd5d6 Zd7d8 Zd9d: Zd;d< Z d=d> Z!d?d@ Z"dAdB Z#dCdD Z$dEdF Z%dGdH Z&dIdJ Z'dKdL Z(dMdN Z)dOdP Z*dQdR Z+dSdT Z,dUdV Z-dWS )XTestGeneratorsc                 C   st   |  t|t| dd |D }dd |D }|  || | t t| W d    d S 1 s3w   Y  d S )Nc                 S      g | ]}|qS r   r   r2   r   r   r   r   
<listcomp>       z2TestGenerators.check_generator.<locals>.<listcomp>c                 S   rK   r   r   rL   r   r   r   rM      rN   )assertEqualnextassertRaisesStopIteration)selfpygencgenexpectedgotr   r   r   check_generator   s   
"zTestGenerators.check_generatorc                 K   <   t }ttjffi ||}|d}|d}| || d S N   )r   r   r   int32rX   rS   kwargspyfunccrrT   rU   r   r   r   
check_gen1   
   zTestGenerators.check_gen1c                 C      | j di t d S Nr   )ra   nopython_flagsrS   r   r   r   	test_gen1      zTestGenerators.test_gen1c                 C   rc   rd   )ra   forceobj_flagsrf   r   r   r   test_gen1_objmode   rh   z TestGenerators.test_gen1_objmodec                 K   rY   rZ   )r   r   r   r\   rX   r]   r   r   r   
check_gen2   rb   zTestGenerators.check_gen2c                 C   rc   rd   )rk   re   rf   r   r   r   	test_gen2   rh   zTestGenerators.test_gen2c                 C   rc   rd   )rk   ri   rf   r   r   r   test_gen2_objmode   rh   z TestGenerators.test_gen2_objmodec                 K   rY   rZ   )r    r   r   r\   rX   r]   r   r   r   
check_gen3   rb   zTestGenerators.check_gen3c                 C   rc   rd   )rn   re   rf   r   r   r   	test_gen3   rh   zTestGenerators.test_gen3c                 C   rc   rd   )rn   ri   rf   r   r   r   test_gen3_objmode   rh   z TestGenerators.test_gen3_objmodec                 K   sH   t }ttjfd fi ||}|ddd}|ddd}| || d S )Nr            )r"   r   r   r\   rX   r]   r   r   r   
check_gen4   s
   zTestGenerators.check_gen4c                 C   rc   rd   )rt   re   rf   r   r   r   	test_gen4   rh   zTestGenerators.test_gen4c                 C   rc   rd   )rt   ri   rf   r   r   r   test_gen4_objmode   rh   z TestGenerators.test_gen4_objmodec                 C   sP   |   }tdi tt W d    n1 sw   Y  | dt|j d S )Nr   z2Cannot type generator: it does not yield any valuer   )assertTypingErrorr   re   r#   assertInstr	exception)rS   raisesr   r   r   	test_gen5   s   
zTestGenerators.test_gen5c                 C   s\   t di tt }| t|g  | t t| W d    d S 1 s'w   Y  d S Nr   rw   )r   ri   r#   rO   listrQ   rR   rP   )rS   rU   r   r   r   test_gen5_objmode   s
   
"z TestGenerators.test_gen5_objmodec                 K   s\   t tjfd fi |t}|dd}g }tdD ]	}|t| q| |dgd  d S )Nr$   rq   rr   r      )r   r   r\   r'   r   appendrP   rO   )rS   r^   r`   rU   lr   r   r   r   
check_gen6   s   
zTestGenerators.check_gen6c                 C   rc   rd   )r   re   rf   r   r   r   	test_gen6   rh   zTestGenerators.test_gen6c                 C   rc   rd   )r   ri   rf   r   r   r   test_gen6_objmode   rh   z TestGenerators.test_gen6_objmodec                 K   sX   t }tttjddffi ||}tddd}|| }||}| || d S )Nr   C
   rs   )	r,   r   r   Arrayfloat64r7   linspacecopyrX   )rS   r^   r_   r`   r+   rT   rU   r   r   r   
check_gen7   s   "zTestGenerators.check_gen7c                 C   rc   rd   )r   re   rf   r   r   r   	test_gen7   rh   zTestGenerators.test_gen7c                 C   rc   rd   )r   ri   rf   r   r   r   test_gen7_objmode   rh   z TestGenerators.test_gen7_objmodec                    sR   t tdi |  fdd}|dd |d |dd |dd	d
 d S )Nc                     s$    | i | | i | d S r   )rX   )argsr^   cfuncr_   rS   r   r   check   s   z(TestGenerators.check_gen8.<locals>.checkr$   r   rG   rq   )r   rr   T)r   r&   r   )r.   r   )rS   jit_argsr   r   r   r   
check_gen8   s   

zTestGenerators.check_gen8c                 C      | j dd d S )NTr   r   rf   r   r   r   	test_gen8      zTestGenerators.test_gen8c                 C   r   )NT)r	   r   rf   r   r   r   test_gen8_objmode   r   z TestGenerators.test_gen8_objmodec                 K   s2   t }tdi ||}| }| }| || d S r~   )rC   r   rX   r]   r   r   r   
check_gen9  s
   zTestGenerators.check_gen9c                 C   rc   rd   )r   re   rf   r   r   r   	test_gen9	  rh   zTestGenerators.test_gen9c                 C   rc   rd   )r   ri   rf   r   r   r   test_gen9_objmode  rh   z TestGenerators.test_gen9_objmodec                 C   sH   t dd|}t ddt|}t|}|d}|d}| || d S )NTr   rq   )r   r   assertPreciseEqual)rS   r   rU   r   r_   rV   rW   r   r   r   check_consume_generator  s   z&TestGenerators.check_consume_generatorc                 C      |  t d S r   )r   r   rf   r   r   r   test_consume_gen1  r5   z TestGenerators.test_consume_gen1c                 C   r   r   )r   r   rf   r   r   r   test_consume_gen2  r5   z TestGenerators.test_consume_gen2c                 C   r   r   )r   r    rf   r   r   r   test_consume_gen3  r5   z TestGenerators.test_consume_gen3c                 K   sH   t }tttjdffi ||}d}||}||}| || d S )Nr$   r$   r   )r;   r   r   UniTupleintprX   )rS   r^   r_   r`   r9   rT   rU   r   r   r   check_ndindex"  s    zTestGenerators.check_ndindexc                 C   rc   rd   )r   re   rf   r   r   r   test_ndindex*  rh   zTestGenerators.test_ndindexc                 C   rc   rd   )r   ri   rf   r   r   r   test_ndindex_objmode-  rh   z#TestGenerators.test_ndindex_objmodec                 K   s   t ttjddffi ||}tjdtjdd}| |||| t ttjddffi ||}|j}| |||| d S )Nr$   r   rr   dtyper   A)	r   r   r   r\   r7   arangereshaperX   T)rS   r_   r^   r`   r+   crAr   r   r   check_np_flat0  s   ""zTestGenerators.check_np_flatc                 C      | j tfi t d S r   )r   r?   re   rf   r   r   r   test_np_flat8     zTestGenerators.test_np_flatc                 C   r   r   )r   r?   ri   rf   r   r   r   test_np_flat_objmode;  r   z#TestGenerators.test_np_flat_objmodec                 C   r   r   )r   rB   re   rf   r   r   r   test_ndenumerate>  r   zTestGenerators.test_ndenumeratec                 C   r   r   )r   rB   ri   rf   r   r   r   test_ndenumerate_objmodeA  r   z'TestGenerators.test_ndenumerate_objmodec                 C   sX   t }|  }tdi t| W d    n1 sw   Y  d}| |t|j d S )Nr   zACan't unify yield type from the following types: complex128, nonerw   )rE   rx   r   re   ry   rz   r{   rS   r_   r|   msgr   r   r   test_type_unification_errorD  s   
z*TestGenerators.test_type_unification_errorc                 C   s`   t }|  }tdi t| W d    n1 sw   Y  d}| |tjj t|j	 d S )Nr   zHCan't unify yield type from the following types: complex128, int%s, nonerw   )
rF   rx   r   re   ry   r   r   bitwidthrz   r{   r   r   r   r   .test_optional_expansion_type_unification_errorM  s   
z=TestGenerators.test_optional_expansion_type_unification_errorc                 C   ,   t }t| }tt| }| || d S r   )rH   r   r   rO   rS   r_   rV   rW   r   r   r   test_changing_tuple_typeV     
z'TestGenerators.test_changing_tuple_typec                 C   r   r   )rI   r   r   rO   r   r   r   r   test_changing_number_type]  r   z(TestGenerators.test_changing_number_typeN).__name__
__module____qualname__rX   ra   rg   rj   rk   rl   rm   rn   ro   rp   rt   ru   rv   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   r   r   r   r   rJ      sX    
		rJ   c                 c   s    | D ]}|V  qd S r   r   )aryelemr   r   r   nrt_gen0e  s   r   c                 c   s&    t | |D ]
\}}|V  |V  qd S r   )zip)ary1ary2e1e2r   r   r   nrt_gen1j  s
   r   c                   @   s4   e Zd Zdd Zdd Zdd Zdd Zd	d
 ZdS )TestNrtArrayGenc                 C   sf   t }tdd|}td}| }t||}t||}tj|| | || | 	|| d S )NTr   r   )
r   r   r7   r   r   r   testingassert_equalrO   assertRefCountEqual)rS   rT   rU   py_aryc_arypy_resc_resr   r   r   test_nrt_gen0q  s   
zTestNrtArrayGen.test_nrt_gen0c           	      C   s   t }tdd|}td}|d }| }| }t|||}t|||}tj|| tj|| | || | 	|| | 	|| d S )NTr   r   d   )
r   r   r7   r   r   r   r   r   rO   r   )	rS   rT   rU   py_ary1py_ary2c_ary1c_ary2r   r   r   r   r   test_nrt_gen1  s   
zTestNrtArrayGen.test_nrt_gen1c                 C   s   |    |   dS )a2  
        Issue #1163 is observed when two generator with NRT object arguments
        is ran in sequence.  The first one does a invalid free and corrupts
        the NRT memory subsystem.  The second generator is likely to segfault
        due to corrupted NRT data structure (an invalid MemInfo).
        N)r   r   rf   r   r   r   test_combine_gen0_gen1  s   z&TestNrtArrayGen.test_combine_gen0_gen1c           	      C   s   t }tdd|}td}| }||}||}t|}t|}| t t|}W d   n1 s6w   Y  | t t|}W d   n1 sOw   Y  ~~tj	|| | 
|| | || dS )z/
        Test cleanup on StopIteration
        Tr   r   N)r   r   r7   r   r   rP   rQ   rR   r   r   rO   r   )	rS   rT   rU   r   r   py_iterc_iterr   r   r   r   r   test_nrt_gen0_stop_iteration  s&   


z,TestNrtArrayGen.test_nrt_gen0_stop_iterationc                 C   sV   t }tdd|}td}| }||}||}~~tj|| | || dS )zj
        Test cleanup for a initialized but never iterated (never call next())
        generator.
        Tr   r   N)r   r   r7   r   r   r   r   r   )rS   rT   rU   r   r   r   r   r   r   r   test_nrt_gen0_no_iter  s   
z%TestNrtArrayGen.test_nrt_gen0_no_iterN)r   r   r   r   r   r   r   r   r   r   r   r   r   p  s    
r   c                   @   s*   e Zd Zdd Zejdd Zdd ZdS )TestNrtNestedGenc           
      C   s   dd }dd }t d}| }|||\}}tdd|tdd|}||\}}	| || | || | |	| t j|| | || d S )Nc                 s   s    t | jD ]}| V  qd S r   r(   r*   r   r   r   gen0  r<   z2TestNrtNestedGen.test_nrt_nested_gen.<locals>.gen0c                    r
   )Nc                    s(   t | } | D ]}|| }q	|| fS r   )r7   
zeros_like)r+   outr   r   r   r   r     s   

zCTestNrtNestedGen.test_nrt_nested_gen.<locals>.factory.<locals>.gen1r   r   r   r   r   r   factory  r   z5TestNrtNestedGen.test_nrt_nested_gen.<locals>.factoryr   Tr   )	r7   r   r   r   assertIsNotassertIsr   r   r   )
rS   r   r   py_arrc_arrr   py_oldc_genr   c_oldr   r   r   test_nrt_nested_gen  s   
	z$TestNrtNestedGen.test_nrt_nested_genc                 C   s   dd }dd }t d}| }|||}tdd|tdd|}||}| || | || | || | || d S )Nc                 s   s    | V  d S r   r   )r+   r   r   r   r     rD   z8TestNrtNestedGen.test_nrt_nested_gen_refct.<locals>.gen0c                    r
   )Nc                    s    | D ]}|  S d S r   r   )r+   r   r   r   r   r     s   zITestNrtNestedGen.test_nrt_nested_gen_refct.<locals>.factory.<locals>.gen1r   r   r   r   r   r     s   z;TestNrtNestedGen.test_nrt_nested_gen_refct.<locals>.factoryr   Tr   )r7   r   r   r   r   r   r   )rS   r   r   r   r   r   r   r   r   r   r   test_nrt_nested_gen_refct  s   
z*TestNrtNestedGen.test_nrt_nested_gen_refctc                 C   sL   dd fdd}| }|t dd}t|d}t|d}| || dS )	z/
        Test nesting three generators
        c                 S   s   | S r   r   r   r   r   r   <lambda>  s    z?TestNrtNestedGen.test_nrt_nested_nopython_gen.<locals>.<lambda>c                    s0   | dd | fdd |  fdd}|S )Nc                 s   s.    t |D ]}| | V  | |  |7  < qd S r   r   )r%   nr   r   r   r   foo  s
   
zKTestNrtNestedGen.test_nrt_nested_nopython_gen.<locals>.factory.<locals>.fooc                 3   sD    t | } || D ]}|d V  qt|jD ]}|| V  qd S )Nr$   )r7   r   r   r)   )r   r%   r   )r   r   r   bar  s   
zKTestNrtNestedGen.test_nrt_nested_nopython_gen.<locals>.factory.<locals>.barc                 3   s     | D ]}|| V  qd S r   r   )r   r   )r   r   r   cat$  s   zKTestNrtNestedGen.test_nrt_nested_nopython_gen.<locals>.factory.<locals>.catr   )decorr   r   )r   r   r   r     s   
z>TestNrtNestedGen.test_nrt_nested_nopython_gen.<locals>.factoryTr   r   N)r   r   rO   )rS   r   py_genr   r   r   r   r   r   test_nrt_nested_nopython_gen  s   z-TestNrtNestedGen.test_nrt_nested_nopython_genN)r   r   r   r   unittestexpectedFailurer   r   r   r   r   r   r     s
    #
r   c                   @   s$   e Zd Zdd Zdd Zdd ZdS )TestGeneratorWithNRTc                 C   s`   t dddd }t|d}| t|d td}d|dd< |D ]	}tj|| q$dS )	z9
        Missing environment for returning array
        Tr   c                 s   s0    t | D ]}td}d|d d < |V  qd S )Nr      )r   r7   empty)r   r   vecr   r   r   random_directions:  s   
z?TestGeneratorWithNRT.test_issue_1254.<locals>.random_directionsrq   r   r   N)r   r   rO   lenr7   r   r   r   )rS   r  outputsexpectrW   r   r   r   test_issue_12545  s   

z$TestGeneratorWithNRT.test_issue_1254c                    s   dd t dd tddd}t ddd}| || fdd	}t dd fd
d}d}dd|f}||}||}tj|| dS )zM
        Double-free for locally allocated, non escaping NRT objects
        c                 s   s(    t | ||}|d V  |d V  d S )Nr   r   )r7   r   )rminrmaxnrr%   r   r   r   r   N  s   
z4TestGeneratorWithNRT.test_issue_1265.<locals>.py_genTr   r$   r   c                    @   | \}}}t j|t jd}t |||D ]\}}|||< q|S Nr   r7   r   
complex128	enumerater   r  r  r  pointsr   c)r   r   r   	py_driverZ  s
   

z7TestGeneratorWithNRT.test_issue_1265.<locals>.py_driverc                    r
  r  r  r  )r   r   r   c_driverb  s
   

z6TestGeneratorWithNRT.test_issue_1265.<locals>.c_driverN)r   r   rO   r7   r   r   )rS   r   r   r  r  r   patchesr   )r   r   r   test_issue_1265I  s   
z$TestGeneratorWithNRT.test_issue_1265c                    sF   dt fdd t  fddt fdd}| |  dS )	z-
        Incorrect return data model
        l   >[= c                   3   s     V  d S r   r   r   )magicr   r   	generatorz  s   
z7TestGeneratorWithNRT.test_issue_1808.<locals>.generatorc                      s     S r   r   r   )r  r   r   get_generator~  s   z;TestGeneratorWithNRT.test_issue_1808.<locals>.get_generatorc                     s   d}   D ]}| |7 } q| S )Nr   r   )r   r   )r  r   r   main  s   

z2TestGeneratorWithNRT.test_issue_1808.<locals>.mainN)r   rO   )rS   r  r   )r  r  r  r   test_issue_1808t  s   z$TestGeneratorWithNRT.test_issue_1808N)r   r   r   r  r  r  r   r   r   r   r   4  s    +r   c                	   @   s<   e Zd Zejdejejejgejejddd gddZ	dS )TestGeneratorModelNr   F)r   
yield_type	arg_typesstate_typeshas_finalizer)
r   r   r   r   	Generatorr\   int64float32r   fe_typer   r   r   r   r    s    


r  __main__)r   r$   F)+numpyr7   r   numbar   r   
numba.corer   numba.tests.supportr   r   numba.core.datamodel.testingr   ri   re   r   r   r   r    r"   r#   r'   r,   r.   r0   r4   r;   r?   rB   rC   rE   rF   rH   rI   rJ   r   r   r   r   r   r  r   r  r   r   r   r   <module>   sL    


	 ^d`Y