o
    í@ËaN  ã                   @   sP   d Z ddlmZ ddlmZ ddlmZ ddlZddl	m
Z
 G dd„ deƒZdS )zC RemoveComprehension turns list comprehension into function calls. é    )ÚImportedIds)ÚTransformationN©Úreducec                   @   sL   e Zd ZdZdd„ Zedd„ ƒZdd„ Zdd	„ Zd
d„ Z	dd„ Z
dd„ ZdS )ÚRemoveComprehensiona  
    Turns all list comprehension from a node into new function calls.

    >>> import gast as ast
    >>> from pythran import passmanager, backend
    >>> node = ast.parse("[x*x for x in (1,2,3)]")
    >>> pm = passmanager.PassManager("test")
    >>> _, node = pm.apply(RemoveComprehension, node)
    >>> print(pm.dump(backend.Python, node))
    list_comprehension0()
    def list_comprehension0():
        __target = builtins.list()
        for x in (1, 2, 3):
            builtins.list.append(__target, (x * x))
        return __target
    c                 C   s   d| _ t | ¡ d S )Nr   )Úcountr   Ú__init__)Úself© r
   úN/usr/lib/python3/dist-packages/pythran/transformations/remove_comprehension.pyr      s   zRemoveComprehension.__init__c                 C   s(   dd„ }t  |j|j|| |jƒgg d¡S )aË  
        Create a ast.For node from a comprehension and another node.

        g is an ast.comprehension.
        x is the code that have to be executed.

        Examples
        --------
        >> [i for i in range(2)]

        Becomes

        >> for i in range(2):
        >>    ... x code with if clauses ...

        It is a reducer as it can be call recursively for mutli generator.

        Ex : >> [i, j for i in range(2) for j in range(4)]
        c                 S   s   t dd„ || ƒS )aˆ  
            Wrap comprehension content in all possibles if clauses.

            Examples
            --------
            >> [i for i in range(2) if i < 3 if 0 < i]

            Becomes

            >> for i in range(2):
            >>    if i < 3:
            >>        if 0 < i:
            >>            ... the code from `node` ...

            Note the nested ifs clauses.
            c                 S   s   t  || gg ¡S ©N)ÚastZIf)ÚnZif_r
   r
   r   Ú<lambda>H   s    zGRemoveComprehension.nest_reducer.<locals>.wrap_in_ifs.<locals>.<lambda>r   )ÚnodeÚifsr
   r
   r   Úwrap_in_ifs7   s   z5RemoveComprehension.nest_reducer.<locals>.wrap_in_ifsN)r   ZForÚtargetÚiterr   )ÚxÚgr   r
   r
   r   Únest_reducer"   s    z RemoveComprehension.nest_reducerc                 G   sž  d| _ |  |j¡|_d || j¡}|  jd7  _|  t|¡}d| _d}t| j	t
|jƒt t tdd„ |dd … t |d t ¡ d d ¡ƒt |t ¡ d d ¡|jgg ¡¡ƒ}t |t |¡¡ t t |t ¡ d d ¡gt t t dt ¡ d d ¡|t ¡ ¡g g ¡d ¡}t t |t ¡ d d ¡¡}	d	d
„ |D ƒ}
t |t |
g d g g d g ¡|||	gg d d ¡}t |t ¡ ¡ | jjj |¡ t t |t ¡ d d ¡dd
„ |
D ƒg ¡S )NTz{0}_comprehension{1}é   r   Z__targetc                 S   s   t  | |t  ¡ ¡S r   )r   Ú	AttributeÚLoad)r   Úyr
   r
   r   r   X   s    ÿz3RemoveComprehension.visit_AnyComp.<locals>.<lambda>Úbuiltinsc                 S   ó    g | ]}t  |t  ¡ d d ¡‘qS r   ©r   ÚNameZParam©Ú.0Úargr
   r
   r   Ú
<listcomp>p   ó     z5RemoveComprehension.visit_AnyComp.<locals>.<listcomp>c                 S   ó"   g | ]}t  |jt  ¡ d d ¡‘qS r   ©r   r   Úidr   r    r
   r
   r   r#   y   ó   " )ÚupdateÚvisitÚeltÚformatr   Úgatherr   Ú
count_iterr   r   ÚreversedÚ
generatorsr   ÚExprÚCallr   r   ÚmetadataÚaddZComprehensionZAssignZStorer   ZReturnÚFunctionDefÚ	argumentsÚLocalÚctxÚmoduleÚbodyÚappend)r	   r   Z	comp_typeÚpathÚnameÚargsZstargetr:   ZinitÚresultÚsargsÚfdr
   r
   r   Úvisit_AnyCompK   sf   
ÿýÿøÿþýú÷
ýýz!RemoveComprehension.visit_AnyCompc                 C   ó   |   |dddd¡S )NÚlistr   r;   ©rB   ©r	   r   r
   r
   r   Úvisit_ListComp}   s   ÿz"RemoveComprehension.visit_ListCompc                 C   rC   )NÚsetr   r4   rE   rF   r
   r
   r   Úvisit_SetComp   s   z!RemoveComprehension.visit_SetCompc                 C   s8   t  t  |j|jgt  ¡ ¡gt  ¡ ¡|_|  |ddd¡S )NÚdictZ__dispatch__r)   )r   ZListZTupleÚkeyÚvaluer   r+   rB   rF   r
   r
   r   Úvisit_DictComp„   s
   þz"RemoveComprehension.visit_DictCompc                 C   sà   d| _ |  |j¡|_d | j¡}|  jd7  _|  t|¡}d| _t| j	t
|jƒt t |j¡¡ƒ}dd„ |D ƒ}t |t |g d g g d g ¡|gg d d ¡}t |t ¡ ¡ | jjj |¡ t t |t ¡ d d ¡dd„ |D ƒg ¡S )NTzgenerator_expression{0}r   r   c                 S   r   r   r   r    r
   r
   r   r#   š   r$   z:RemoveComprehension.visit_GeneratorExp.<locals>.<listcomp>c                 S   r%   r   r&   r    r
   r
   r   r#   ¢   r(   )r)   r*   r+   r,   r   r-   r   r.   r   r   r/   r0   r   r1   ZYieldr5   r6   r3   r4   r7   r8   r9   r:   r;   r2   r   r   )r	   r   r=   r>   r:   r@   rA   r
   r
   r   Úvisit_GeneratorExp   s,   þ
þýz&RemoveComprehension.visit_GeneratorExpN)Ú__name__Ú
__module__Ú__qualname__Ú__doc__r   Ústaticmethodr   rB   rG   rI   rM   rN   r
   r
   r
   r   r      s    
(2	r   )rR   Zpythran.analysesr   Zpythran.passmanagerr   Zpythran.metadatar3   Zgastr   Ú	functoolsr   r   r
   r
   r
   r   Ú<module>   s    