o
    2g                     @  st   d dl mZ d dlmZmZmZmZmZmZm	Z	m
Z
mZmZmZmZmZ d dlmZ d dlmZ G dd deZdS )    )annotations)Block	Decorator
ExpressionFuncDefFuncItemImport
LambdaExpr
MemberExprMypyFileNameExprNode
SymbolNodeVar)ExtendedTraverserVisitor)Errorsc                      s   e Zd ZdZd: fd
dZd;ddZd< fddZd= fddZd>ddZd?dd Z	d@ fd#d$Z
dA fd'd(ZdBd*d+ZdCd.d/ZdDd2d3ZdEd6d7ZdDd8d9Z  ZS )FPreBuildVisitora  Mypy file AST visitor run before building the IR.

    This collects various things, including:

    * Determine relationships between nested functions and functions that
      contain nested functions
    * Find non-local variables (free variables)
    * Find property setters
    * Find decorators of functions
    * Find module import groups

    The main IR build pass uses this information.
    errorsr   current_filer   decorators_to_removedict[FuncDef, list[int]]returnNonec                   sX   t    i | _i | _g | _t | _i | _i | _i | _	|| _
i | _d | _|| _|| _d S N)super__init__free_variablessymbols_to_funcsfuncssetprop_settersencapsulating_funcsnested_funcsfuncs_to_decoratorsr   module_import_groups_current_import_groupr   r   )selfr   r   r   	__class__ U/home/garg/my-data/venv/lib/python3.10/site-packages/mypyc/irbuild/prebuildvisitor.pyr   %   s   

zPreBuildVisitor.__init__or   boolc                 C  s   t |tsd | _dS )NT)
isinstancer   r%   )r&   r+   r)   r)   r*   visitU   s   
zPreBuildVisitor.visitblockr   c                   s   d | _ t | d | _ d S r   )r%   r   visit_block)r&   r/   r'   r)   r*   r0   Z   s   
zPreBuildVisitor.visit_blockdecr   c                   s   |j r@t|j d tr|j d jdkr| j|j n%|j  }|j| jv r:| j|j }t	|D ]}||= q0|s:d S || j
|j< t | d S )Nr   setter)
decoratorsr-   r
   namer    addfunccopyr   reversedr#   r   visit_decorator)r&   r1   decorators_to_store	to_removeir'   r)   r*   r9   _   s    
zPreBuildVisitor.visit_decoratorfdefr   c                 C  s   |  | | | d S r   )
visit_funcvisit_symbol_node)r&   r=   r)   r)   r*   visit_func_defz   s   
zPreBuildVisitor.visit_func_defexprr	   c                 C     |  | d S r   )r>   r&   rA   r)   r)   r*   visit_lambda_expr      z!PreBuildVisitor.visit_lambda_exprr6   r   c                   sV   | j r| j| j d g | | j d | j|< | j | t | | j   d S N)r   r!   
setdefaultappendr"   r   r>   pop)r&   r6   r'   r)   r*   r>      s   zPreBuildVisitor.visit_funcimpr   c                   s@   | j d ur| j| j  | n	|g| j|< || _ t | d S r   )r%   r$   rI   r   visit_import)r&   rK   r'   r)   r*   rL      s
   
zPreBuildVisitor.visit_importr   c                 C  s$   t |jttfr| |j d S d S r   )r-   noder   r   r?   rC   r)   r)   r*   visit_name_expr   s   zPreBuildVisitor.visit_name_exprvarr   c                 C  rB   r   )r?   )r&   rO   r)   r)   r*   	visit_var   rE   zPreBuildVisitor.visit_varsymbolr   c                 C  s   | j sd S || jv rB| j| }| | j d |r0| j d | j|< | j| j d t | d S | || j d r@| | d S d S | j d | j|< d S rF   )r   r   	is_parentr   rH   r   r5   add_free_variable)r&   rQ   	orig_funcr)   r)   r*   r?      s   

 
z!PreBuildVisitor.visit_symbol_nodefitemchildc                 C  s,   || j vrdS | j | }||kp| ||S )NF)r"   rR   )r&   rU   rV   parentr)   r)   r*   rR      s   

zPreBuildVisitor.is_parentc                 C  s$   | j | }| j|t | d S r   )r   r   rH   r   r5   )r&   rQ   r6   r)   r)   r*   rS      s   
z!PreBuildVisitor.add_free_variable)r   r   r   r   r   r   r   r   )r+   r   r   r,   )r/   r   r   r   )r1   r   r   r   )r=   r   r   r   )rA   r	   r   r   )r6   r   r   r   )rK   r   r   r   )rA   r   r   r   )rO   r   r   r   )rQ   r   r   r   )rU   r   rV   r   r   r,   )__name__
__module____qualname____doc__r   r.   r0   r9   r@   rD   r>   rL   rN   rP   r?   rR   rS   __classcell__r)   r)   r'   r*   r      s    
0





r   N)
__future__r   
mypy.nodesr   r   r   r   r   r   r	   r
   r   r   r   r   r   mypy.traverserr   mypyc.errorsr   r   r)   r)   r)   r*   <module>   s
    <