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   j/var/www/html/eduruby.in/lip-sync/lip-sync-env/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&   Zbbr   r   r   gen8I   s   r-   c                 c   s    t   | V  d S r   )objectr   r   r   r   genobjR      
r/   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      r3   c                 c       t | D ]}|V  qd S r   )npZndindex)shapeindr   r   r   gen_ndindex[      r9   c                 c   s    | j D ]}|V  qd S r   )Zflat)r+   valr   r   r   gen_flat`   s   
r<   c                 c   r5   r   )r6   Zndenumerate)r+   tupr   r   r   gen_ndenumeratee   r:   r>   c                   c   s    dV  d S )NTr   r   r   r   r   gen_boolj      
r?   c                   c   s    d V  dV  d S )Nr   r   r   r   r   r   gen_unification_errorn   r0   rA   c                  c   s    d} dV  	 | V } q)Nr   r   r   )r   r   r   r   'gen_optional_and_type_unification_errors   s   rB   c                   c   s    dV  dV  d S )N)r   r$   )r      r   r   r   r   r   gen_changing_tuple_type{   s   
rD   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   
rE   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   r1   r   r   r   r   
<listcomp>       z2TestGenerators.check_generator.<locals>.<listcomp>c                 S   rG   r   r   rH   r   r   r   rI      rJ   )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   int32rT   rO   kwargspyfunccrrP   rQ   r   r   r   
check_gen1   
   zTestGenerators.check_gen1c                 C      | j di t d S Nr   )r]   nopython_flagsrO   r   r   r   	test_gen1      zTestGenerators.test_gen1c                 C   r_   r`   )r]   forceobj_flagsrb   r   r   r   test_gen1_objmode   rd   z TestGenerators.test_gen1_objmodec                 K   rU   rV   )r   r   r   rX   rT   rY   r   r   r   
check_gen2   r^   zTestGenerators.check_gen2c                 C   r_   r`   )rg   ra   rb   r   r   r   	test_gen2   rd   zTestGenerators.test_gen2c                 C   r_   r`   )rg   re   rb   r   r   r   test_gen2_objmode   rd   z TestGenerators.test_gen2_objmodec                 K   rU   rV   )r    r   r   rX   rT   rY   r   r   r   
check_gen3   r^   zTestGenerators.check_gen3c                 C   r_   r`   )rj   ra   rb   r   r   r   	test_gen3   rd   zTestGenerators.test_gen3c                 C   r_   r`   )rj   re   rb   r   r   r   test_gen3_objmode   rd   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   rX   rT   rY   r   r   r   
check_gen4   s
   zTestGenerators.check_gen4c                 C   r_   r`   )rp   ra   rb   r   r   r   	test_gen4   rd   zTestGenerators.test_gen4c                 C   r_   r`   )rp   re   rb   r   r   r   test_gen4_objmode   rd   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   ra   r#   assertInstr	exception)rO   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   rs   )r   re   r#   rK   listrM   rN   rL   )rO   rQ   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$   rm   rn   r      )r   r   rX   r'   r   appendrL   rK   )rO   rZ   r\   rQ   lr   r   r   r   
check_gen6   s   
zTestGenerators.check_gen6c                 C   r_   r`   )r   ra   rb   r   r   r   	test_gen6   rd   zTestGenerators.test_gen6c                 C   r_   r`   )r   re   rb   r   r   r   test_gen6_objmode   rd   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
   ro   )	r,   r   r   ArrayZfloat64r6   linspacecopyrT   )rO   rZ   r[   r\   r+   rP   rQ   r   r   r   
check_gen7   s   "zTestGenerators.check_gen7c                 C   r_   r`   )r   ra   rb   r   r   r   	test_gen7   rd   zTestGenerators.test_gen7c                 C   r_   r`   )r   re   rb   r   r   r   test_gen7_objmode   rd   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   )rT   )argsrZ   cfuncr[   rO   r   r   check   s   z(TestGenerators.check_gen8.<locals>.checkr$   r   rC   rm   )r   rn   T)r   r&   r   )r-   r   )rO   Zjit_argsr   r   r   r   
check_gen8   s   

zTestGenerators.check_gen8c                 C      | j dd d S )NTr   r   rb   r   r   r   	test_gen8      zTestGenerators.test_gen8c                 C   r   )NT)r	   r   rb   r   r   r   test_gen8_objmode   r   z TestGenerators.test_gen8_objmodec                 K   s2   t }tdi ||}| }| }| || d S rz   )r?   r   rT   rY   r   r   r   
check_gen9  s
   zTestGenerators.check_gen9c                 C   r_   r`   )r   ra   rb   r   r   r   	test_gen9	  rd   zTestGenerators.test_gen9c                 C   r_   r`   )r   re   rb   r   r   r   test_gen9_objmode  rd   z TestGenerators.test_gen9_objmodec                 C   sH   t dd|}t ddt|}t|}|d}|d}| || d S )NTr   rm   )r   r   ZassertPreciseEqual)rO   r   rQ   r   r[   rR   rS   r   r   r   check_consume_generator  s   z&TestGenerators.check_consume_generatorc                 C      |  t d S r   )r   r   rb   r   r   r   test_consume_gen1  r4   z TestGenerators.test_consume_gen1c                 C   r   r   )r   r   rb   r   r   r   test_consume_gen2  r4   z TestGenerators.test_consume_gen2c                 C   r   r   )r   r    rb   r   r   r   test_consume_gen3  r4   z TestGenerators.test_consume_gen3c                 K   sH   t }tttjdffi ||}d}||}||}| || d S )Nr$   r$   r   )r9   r   r   ZUniTupleintprT   )rO   rZ   r[   r\   r7   rP   rQ   r   r   r   check_ndindex"  s    zTestGenerators.check_ndindexc                 C   r_   r`   )r   ra   rb   r   r   r   test_ndindex*  rd   zTestGenerators.test_ndindexc                 C   r_   r`   )r   re   rb   r   r   r   test_ndindex_objmode-  rd   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   rn   Zdtyper   A)	r   r   r   rX   r6   arangeZreshaperT   T)rO   r[   rZ   r\   r+   ZcrAr   r   r   check_np_flat0  s   ""zTestGenerators.check_np_flatc                 C      | j tfi t d S r   )r   r<   ra   rb   r   r   r   test_np_flat8     zTestGenerators.test_np_flatc                 C   r   r   )r   r<   re   rb   r   r   r   test_np_flat_objmode;  r   z#TestGenerators.test_np_flat_objmodec                 C   r   r   )r   r>   ra   rb   r   r   r   test_ndenumerate>  r   zTestGenerators.test_ndenumeratec                 C   r   r   )r   r>   re   rb   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, noners   )rA   rt   r   ra   ru   rv   rw   rO   r[   rx   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, noners   )
rB   rt   r   ra   ru   r   r   Zbitwidthrv   rw   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   )rD   r{   r   rK   rO   r[   rR   rS   r   r   r   test_changing_tuple_typeV     
z'TestGenerators.test_changing_tuple_typec                 C   r   r   )rE   r{   r   rK   r   r   r   r   test_changing_number_type]  r   z(TestGenerators.test_changing_number_typeN).__name__
__module____qualname__rT   r]   rc   rf   rg   rh   ri   rj   rk   rl   rp   rq   rr   ry   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   rF      sX    
		rF   c                 c   s    | D ]}|V  qd S r   r   )Zaryelemr   r   r   nrt_gen0e  s   r   c                 c   s&    t | |D ]
\}}|V  |V  qd S r   )zip)Zary1Z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   r6   r   r   r{   testingassert_equalrK   assertRefCountEqual)rO   rP   rQ   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   r6   r   r   r{   r   r   rK   r   )	rO   rP   rQ   Zpy_ary1Zpy_ary2Zc_ary1Z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   rb   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   r6   r   r   rL   rM   rN   r   r   rK   r   )	rO   rP   rQ   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   r6   r   r   r   r   r   )rO   rP   rQ   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   )r6   Z
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   )	r6   r   r   r   assertIsNotassertIsr   r   r   )
rO   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     r@   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   )r6   r   r   r   r   r   r   )rO   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$   )r6   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   )Zdecorr   r   )r   r   r   r     s   
z>TestNrtNestedGen.test_nrt_nested_nopython_gen.<locals>.factoryTr   r   N)r   r{   rK   )rO   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   unittestZ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   r6   empty)r   r   Zvecr   r   r   random_directions:  s   
z?TestGeneratorWithNRT.test_issue_1254.<locals>.random_directionsrm   r   r   N)r   r{   rK   lenr6   r   r   r   )rO   r   ZoutputsexpectrS   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   )r6   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   r6   r   Z
complex128	enumerater   r   r   r   Z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{   rK   r6   r   r   )rO   r   r   r   r   r   Z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   rK   )rO   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   Z
yield_typeZ	arg_typesZstate_typesZhas_finalizer)
r   r   r   r   	GeneratorrX   Zint64Zfloat32r   Zfe_typer   r   r   r   r    s    


r  __main__)r   r$   F)+numpyr6   r   Znumbar   r   Z
numba.corer   Znumba.tests.supportr   r   Znumba.core.datamodel.testingr   re   ra   r   r   r   r    r"   r#   r'   r,   r-   r/   r3   r9   r<   r>   r?   rA   rB   rD   rE   rF   r   r   r   r   r   r  r   r   r   r   r   r   <module>   sL    


	 ^d`Y