Ë
    pPêjÄ"  ã                  óÂ   — d Z ddlmZ ddlZddlmZmZmZmZm	Z	m
Z
mZ ddlmZ ddlmZmZ ddlmZ ddlmZ dd	lmZ dd
„Zdd„Zdd„Z G d„ dej4                  «      Zy)z”Tests for promoting generator registers onto their frames.

These don't use the normal data-driven tests, since generators produce very verbose IR.
é    )ÚannotationsN)ÚENV_ATTR_NAMEÚGENERATOR_ATTRIBUTE_PREFIXÚ GENERATOR_FRAME_ATTRIBUTE_PREFIXÚNEXT_LABEL_ATTR_NAMEÚ	SELF_NAMEÚTEMP_ATTR_NAMEÚ generator_frame_attribute_prefix)ÚClassIR)ÚAssignÚGetAttr)Ú	RInstance)Úbuild_ir_for_single_file2)Úregisters_live_across_yieldc                ó  — t        | j                  «       «      \  }}}}|j                  D �cg c]=  }|j                  |k(  s|j                  j	                  |dz   «      r|j
                  �|‘Œ? }}t        |«      dk(  sJ ‚|d   S c c}w )NÚ___é   r   )r   Ú
splitlinesÚclassesÚnameÚ
startswithÚenv_user_functionÚlen)Úsourcer   ÚmoduleÚ_ÚclÚmatchess         úe/home/oem/CyberBrain26-unified/.venv/lib/python3.12/site-packages/mypyc/test/test_generator_spills.pyÚgenerator_classr       s�   € Ü/°×0AÑ0AÓ0CÓD�O€FˆAˆq�!ð —.‘.öàØ�G‰G�tŠO˜rŸw™w×1Ñ1°$¸±,Ô?Ø× Ñ Ð,ò 	ð€Gð ô ˆw‹<˜1ÒÐÐØ�1‰:Ðùòs   ­ABc                ór   — | j                   D �ch c]  }|j                  t        dz   «      sŒ|’Œ c}S c c}w )NÚ3_)Ú
attributesr   r	   )r   r   s     r   Úpromoted_slotsr$   &   s*   € ØŸ]™]ÖU�T¨d¯o©o¼nÈtÑ>SÕ.TŠDÒUÐUùÒUs   �4­4c                óÄ   — t        | j                  | j                  ¬«      }| j                  D �ch c]%  }|j	                  |«      sŒ|j                  |«      ’Œ' c}S c c}w )N)Úis_final_class)r
   Úfullnamer&   r#   r   Úremoveprefix)r   Úprefixr   s      r   Úframe_variablesr*   *   sG   € Ü-¨b¯k©kÈ"×J[ÑJ[Ô\€FØ24·-±-Ö[¨$À4Ç?Á?ÐSYÕCZˆD×Ñ˜fÕ%Ò[Ð[ùÒ[s   °AÁAc                  ód   — 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y)ÚTestGeneratorSpillsc                ó†  — t        dj                  «       «      \  }}}}t        d„ |j                  D «       «      }|j                  t
           }t        |t        «      sJ ‚|j                  }|j                  t           }t        |t        «      sJ ‚|j                  |u sJ ‚dddddt        fD ]  }||j                  vrŒJ ‚ y )Nz–from typing import Iterator

def outer() -> Iterator[str]:
    captured = "value"
    def nested() -> str:
        return captured
    yield nested()
c              3  ó:   K  — | ]  }|j                   sŒ|–— Œ y ­w©N)Úhas_running_flag©Ú.0r   s     r   ú	<genexpr>z[TestGeneratorSpills.test_generator_environment_excludes_helper_arguments.<locals>.<genexpr><   s   è ø€ ÒH˜B°B×4GÓ4G”RÑHùs   ‚”ÚtypeÚvalueÚ	tracebackÚargÚstop_iter_ptr)r   r   Únextr   r#   r   Ú
isinstancer   Úclass_irr   r   )Úselfr   r   ÚframeÚenv_typeÚenvironmentÚ	self_typer   s           r   Ú4test_generator_environment_excludes_helper_argumentszHTestGeneratorSpills.test_generator_environment_excludes_helper_arguments0   sÊ   € ô 4ð 5÷ �JƒLó‰ˆ��1�aô ÑH &§.¡.ÔHÓHˆØ×#Ñ#¤MÑ2ˆÜ˜(¤IÔ.Ð.Ð.Ø×'Ñ'ˆà×*Ñ*¬9Ñ5ˆ	Ü˜)¤YÔ/Ð/Ð/Ø×!Ñ! [Ñ0Ð0Ð0Ø˜W k°5¸/ÔK_Ð`ò 	6ˆDØ˜{×5Ñ5Ò5Ð5Ð5ñ	6ó    c           	     óX  — t        dd«      }|j                  }|€J ‚t        |«      }t        |«      dk(  sJ ‚|j                  D ��cg c]3  }|j
                  D ]"  }t        |t        «      r|j                  |v r|‘Œ$ Œ5 }}}|rt        d„ |D «       «      sJ ‚t        |«      rJ ‚y c c}}w )Nzofrom typing import Iterator

def gen() -> Iterator[int]:
    try:
        yield 1
    finally:
        yield 2
Úgen_genr   c              3  ó4   K  — | ]  }|j                   –— Œ y ­wr/   )Úallow_error_value)r2   Úreads     r   r3   zMTestGeneratorSpills.test_saved_exception_state_is_promoted.<locals>.<genexpr>^   s   è ø€ ÒF¸˜T×3Õ3ÑFùs   ‚)r    r   r$   r   ÚblocksÚopsr:   r   ÚattrÚallr   )r<   r   ÚhelperÚslotsÚblockÚopÚreadss          r   Ú&test_saved_exception_state_is_promotedz:TestGeneratorSpills.test_saved_exception_state_is_promotedG   sÈ   € Üðð ó
ˆð ×%Ñ%ˆØÐ!Ð!Ð!Ü˜rÓ"ˆÜ�5‹z˜QŠÐˆð  Ÿ™÷
àØ—i‘iò
ð Ü˜"œgÔ&¨2¯7©7°eÑ+;ò ð
Øð
ˆñ 
ñ œÑFÀÔFÔFÐFÐFÜ.¨vÔ6Ð6Ð6Ð6ùó
s   Á8B&c                óV   — t        dd«      }|j                  }|€J ‚t        |«      rJ ‚y )Nzmfrom typing import Any, Generator

def gen(values: Any) -> Generator[int, None, None]:
    yield from values
rD   )r    r   r   )r<   r   rL   s      r   Ú+test_address_taken_register_is_not_promotedz?TestGeneratorSpills.test_address_taken_register_is_not_promoteda   s@   € Üðð ó
ˆð ×%Ñ%ˆØÐ!Ð!Ð!Ü.¨vÔ6Ð6Ð6Ð6rB   c                óŒ   — d}t        |d«      }t        |d«      }t        |«      }t        |«      }|rJ ‚t        |«      dk(  sJ ‚y )Na  from typing import Iterator

def before_yield(values: list[int]) -> Iterator[int]:
    total = 0
    for value in values:
        total += value
    yield total

def inside_loop(values: list[int]) -> Iterator[int]:
    for value in values:
        yield value
Úbefore_yield_genÚinside_loop_genr   )r    r$   r   )r<   r   ÚbeforeÚinsideÚbefore_slotsÚinside_slotss         r   Ú1test_loop_state_only_crossing_a_yield_uses_a_slotzETestGeneratorSpills.test_loop_state_only_crossing_a_yield_uses_a_sloto   sU   € ðˆô ! Ð);Ó<ˆÜ  Ð):Ó;ˆÜ% fÓ-ˆÜ% fÓ-ˆÙÐÐÜ�<Ó  AÒ%Ð%Ñ%rB   c           	     óô   — t        dd«      }|j                  }|€J ‚|j                  D ��cg c];  }|j                  D ]*  }t	        |t
        «      r|j                  j                  s|‘Œ, Œ= }}}|sJ ‚y c c}}w )Nzsfrom typing import Any

async def collect(values: Any) -> Any:
    return {value: value async for value in values}
Úcollect_gen)r    r   rH   rI   r:   r   ÚdestÚis_arg)r<   r   rL   rN   rO   Úresult_copiess         r   Ú0test_return_value_is_copied_before_frame_cleanupzDTestGeneratorSpills.test_return_value_is_copied_before_frame_cleanup„   sŽ   € Üðð ó
ˆð ×%Ñ%ˆØÐ!Ð!Ð!ð  Ÿ™÷
àØ—i‘iò
ð Ü˜"œfÔ%¨b¯g©g¯nªnò ð
Øð
ˆñ 
ñ Ð‰}ùó
s   ¬A A4c                ó”   — t        dd«      }|j                  sJ ‚t        dz   |j                  v sJ ‚t	        |«      }d|v sJ ‚d|vsJ ‚y )Nz‰from typing import Generator

def gen() -> Generator[int, None, int]:
    local = 1
    crossing = 2
    yield local
    return crossing
rD   ÚcrossingÚlocal)r    r&   r   r#   r*   )r<   r   Ú	variabless      r   Ú*test_only_crossing_named_local_is_on_framez>TestGeneratorSpills.test_only_crossing_named_local_is_on_frame˜   sd   € Üðð ó
ˆð × Ò Ð Ð Ü/°*Ñ<ÀÇÁÑMÐMÐMÜ# BÓ'ˆ	Ø˜YÑ&Ð&Ð&Ø˜iÑ'Ð'Ñ'rB   c                ó¼   — t        dj                  «       «      \  }}}}t        d„ |j                  D «       «      }t        d„ |j                  D «       «      }||usJ ‚y )Nz‘from typing import Any, Iterator

def gen() -> Iterator[Any]:
    captured = "value"
    def get() -> str:
        return captured
    yield get
c              3  ó@   K  — | ]  }|j                   d k(  sŒ|–— Œ y­w)rD   N©r   r1   s     r   r3   zOTestGeneratorSpills.test_captured_local_stays_in_environment.<locals>.<genexpr>µ   s   è ø€ ÒI˜B°B·G±G¸yÓ4H”RÑIùó   ‚—c              3  óL   K  — | ]  }t         d z   |j                  v sŒ|–— Œ y­w)ÚcapturedN)r   r#   r1   s     r   r3   zOTestGeneratorSpills.test_captured_local_stays_in_environment.<locals>.<genexpr>¶   s)   è ø€ ò 
ØÔ+EÈ
Ñ+RÐVX×VcÑVcÒ+cŒBñ
ùs   ‚$�$)r   r   r9   r   )r<   r   r   r=   r?   s        r   Ú(test_captured_local_stays_in_environmentz<TestGeneratorSpills.test_captured_local_stays_in_environment«   sc   € Ü3ð 5÷ �JƒLó‰ˆ��1�aô ÑI &§.¡.ÔIÓIˆÜñ 
ØŸ™ô
ó 
ˆð  %Ñ'Ð'Ñ'rB   c                óè  — t        dj                  «       «      \  }}}}t        d„ |j                  D «       «      }|j                  t
           }t        |t        «      sJ ‚|j                  }t        dz   |j                  v sJ ‚dt        |«      v sJ ‚t        dz   |j                  vsJ ‚t        d„ |j                  D «       «      }|j                  }|€J ‚t        d„ |j                  D «       «      rJ ‚y )NzËfrom typing import Iterator

def outer() -> Iterator[int]:
    def captured() -> Iterator[int]:
        yield 1
    def private() -> Iterator[int]:
        yield from captured()
    yield from private()
c              3  ó@   K  — | ]  }|j                   d k(  sŒ|–— Œ y­w©Ú	outer_genNri   r1   s     r   r3   zZTestGeneratorSpills.test_nested_generator_private_helper_stays_on_frame.<locals>.<genexpr>Æ   ó   è ø€ ÒQ "¸"¿'¹'À[Ó:Pœ2ÑQùrj   rl   Úprivatec              3  óX   K  — | ]"  }|j                   j                  d «      sŒ|–— Œ$ y­w)Úprivate_gen___N)r   r   r1   s     r   r3   zZTestGeneratorSpills.test_nested_generator_private_helper_stays_on_frame.<locals>.<genexpr>Ï   s"   è ø€ Òa B¸B¿G¹G×<NÑ<NÐO_Õ<`œRÑaùs   ‚ *£*c              3  óŽ   K  — | ]=  }|j                   D ],  }t        |t        «      xr |j                  t        d z   k(  –— Œ. Œ? y­w)rs   N)rI   r:   r   rJ   r   )r2   rN   rO   s      r   r3   zZTestGeneratorSpills.test_nested_generator_private_helper_stays_on_frame.<locals>.<genexpr>Ò   sP   è ø€ ò 
àØ—i‘iò
ð ô �rœ7Ó#ÒY¨¯©Ô3MÐPYÑ3YÑ(YÓYð
ØYñ
ùs   ‚AA)r   r   r9   r   r#   r   r:   r   r;   r   r*   r   ÚanyrH   )r<   r   r   Úouter_frameÚenvironment_typer?   Úprivate_frameÚprivate_helpers           r   Ú3test_nested_generator_private_helper_stays_on_framezGTestGeneratorSpills.test_nested_generator_private_helper_stays_on_frame»   sÿ   € Ü3ð 	5÷ �JƒLó	‰ˆ��1�aô ÑQ¨¯©ÔQÓQˆØ&×1Ñ1´-Ñ@ÐÜÐ*¬IÔ6Ð6Ð6Ø&×/Ñ/ˆä)¨JÑ6¸+×:PÑ:PÑPÐPÐPØœO¨KÓ8Ñ8Ð8Ð8Ü)¨IÑ5¸[×=SÑ=SÑSÐSÐSäÑa¨&¯.©.ÔaÓaˆØ&×8Ñ8ˆØÐ)Ð)Ð)Üñ 
à'×.Ñ.ô
ô 
ð 	
ð 
ð 
rB   c                ó&  — t        dj                  «       «      \  }}}}t        d„ |j                  D «       «      }|j                  t
           }t        |t        «      sJ ‚|j                  }t        dz   |j                  v sJ ‚dt        |«      vsJ ‚y )Nzøfrom typing import Iterator

def outer(flag: bool) -> Iterator[int]:
    if flag:
        def helper(n: int) -> int:
            return 1
    else:
        def helper(n: int) -> int:
            return helper(1) if not n else 2
    yield helper(0)
c              3  ó@   K  — | ]  }|j                   d k(  sŒ|–— Œ y­wrp   ri   r1   s     r   r3   z_TestGeneratorSpills.test_redefined_nested_helper_has_single_storage_location.<locals>.<genexpr>å   rr   rj   rL   )r   r   r9   r   r#   r   r:   r   r;   r   r*   )r<   r   r   rx   ry   r?   s         r   Ú8test_redefined_nested_helper_has_single_storage_locationzLTestGeneratorSpills.test_redefined_nested_helper_has_single_storage_locationØ   s‘   € Ü3ð 5÷ �JƒLó‰ˆ��1�aô ÑQ¨¯©ÔQÓQˆØ&×1Ñ1´-Ñ@ÐÜÐ*¬IÔ6Ð6Ð6Ø&×/Ñ/ˆä)¨HÑ4¸×8NÑ8NÑNÐNÐNØœ¨{Ó;Ñ;Ð;Ñ;rB   c                óŒ   — t        dd«      }|j                  }|€J ‚t        |«      }t        |«      dk(  sJ ‚t	        |«      rJ ‚y )Nz¥from typing import Iterator

def gen() -> Iterator[int]:
    try:
        try:
            yield 1
        finally:
            yield 2
    finally:
        yield 3
rD   é   )r    r   r$   r   r   )r<   r   rL   rM   s       r   Ú1test_nested_finally_promotes_each_saved_exceptionzETestGeneratorSpills.test_nested_finally_promotes_each_saved_exceptioní   sZ   € Üðð ó
ˆð ×%Ñ%ˆØÐ!Ð!Ð!Ü˜rÓ"ˆÜ�5‹z˜QŠÐˆÜ.¨vÔ6Ð6Ð6Ð6rB   c                ó:   — t        dd«      }dt        |«      vsJ ‚y )NzÞfrom typing import Any, Iterator

def outer(edges: list[list[int]]) -> Any:
    def walk(node: int) -> Iterator[int]:
        for child in edges[node]:
            yield from walk(child)
        yield node
    return walk
Úwalk_genÚwalk)r    r*   )r<   r   s     r   Ú4test_recursive_generator_reference_is_rematerializedzHTestGeneratorSpills.test_recursive_generator_reference_is_rematerialized  s+   € Üð	ð ó
ˆð œ_¨RÓ0Ñ0Ð0Ñ0rB   N)ÚreturnÚNone)Ú__name__Ú
__module__Ú__qualname__rA   rQ   rS   r[   ra   rf   rm   r|   r   r‚   r†   © rB   r   r,   r,   /   s9   „ ó6ó.7ó47ó&ó*ó((ó&(ó 
ó:<ó*7ô,1rB   r,   )r   Ústrr   r�   r‡   r   )r   r   r‡   zset[str])Ú__doc__Ú
__future__r   ÚunittestÚmypyc.commonr   r   r   r   r   r	   r
   Úmypyc.ir.class_irr   Úmypyc.ir.opsr   r   Úmypyc.ir.rtypesr   Úmypyc.test.testutilr   Ú mypyc.transform.generator_spillsr   r    r$   r*   ÚTestCaser,   rŒ   rB   r   ú<module>r˜      sW   ðñõ
 #ã ÷÷ ñ õ &ß (Ý %Ý 9Ý Hó	óVó\ô
d1˜(×+Ñ+õ d1rB   