U
    A`                     @   s   d Z ddlmZ ddlmZ ddlmZ ddlmZmZm	Z	m
Z
mZmZ dd ZG dd	 d	ejZd
d Zdd Zg dfddZdd ZdS )a:  Fixer for function definitions with tuple parameters.

def func(((a, b), c), d):
    ...

    ->

def func(x, d):
    ((a, b), c) = x
    ...

It will also support lambdas:

    lambda (x, y): x + y -> lambda t: t[0] + t[1]

    # The parens are a syntax error in Python 3
    lambda (x): x + y -> lambda x: x + y
   )pytree)token)
fixer_base)AssignNameNewlineNumber	Subscriptsymsc                 C   s   t | tjo| jd jtjkS )N    )
isinstancer   Nodechildrentyper   STRING)stmt r   4/usr/lib/python3.8/lib2to3/fixes/fix_tuple_params.pyis_docstring   s    r   c                   @   s(   e Zd ZdZdZdZdd Zdd ZdS )	FixTupleParams   Ta  
              funcdef< 'def' any parameters< '(' args=any ')' >
                       ['->' any] ':' suite=any+ >
              |
              lambda=
              lambdef< 'lambda' args=vfpdef< '(' inner=any ')' >
                       ':' body=any
              >
              c                    s  d|kr ||S g |d }|d }|d jd jtjkrZd}|d jd j}t  nd}d}ttjd d fd
d	}|jt	j
kr|| n<|jt	jkrt|jD ]$\}}	|	jt	j
kr||	|dkd qsd S D ]}
|d |
_q|}|dkr
dd _n&t|d j| r0|d _|d }D ]}
|d |
_q4|d j||< t|d |t d D ]}||d j| _qr|d   d S )Nlambdasuiteargsr      r   z;  Fc                    s\   t  }|  }d|_t|| }|r2d|_| | tt	j
|  g d S )Nr    )r   new_namecloneprefixr   replaceappendr   r   r
   Zsimple_stmt)Z	tuple_arg
add_prefixnargr   endZ	new_linesselfr   r   handle_tupleC   s    

z.FixTupleParams.transform.<locals>.handle_tuple)r"   r   )F)transform_lambdar   r   r   INDENTvaluer   r   ZLeafr
   ZtfpdefZtypedargslist	enumerateparentr   r   rangelenZchanged)r'   noderesultsr   r   startindentr(   ir$   lineZafterr   r%   r   	transform.   sF    


zFixTupleParams.transformc                 C   s   |d }|d }t |d }|jtjkrD| }d|_|| d S t|}t|}| 	t
|}t|dd}	||	  | D ]X}
|
jtjkr|
j|krdd ||
j D }ttj|	 g| }|
j|_|
| qd S )Nr   bodyinnerr   )r   c                 S   s   g | ]}|  qS r   )r   .0cr   r   r   
<listcomp>   s     z3FixTupleParams.transform_lambda.<locals>.<listcomp>)simplify_argsr   r   NAMEr   r   r    find_paramsmap_to_indexr   
tuple_namer   Z
post_orderr+   r   r   r
   Zpower)r'   r0   r1   r   r7   r8   ZparamsZto_indexZtup_nameZ	new_paramr#   Z
subscriptsnewr   r   r   r)   n   s*    
zFixTupleParams.transform_lambdaN)__name__
__module____qualname__Z	run_orderZBM_compatibleZPATTERNr6   r)   r   r   r   r   r      s
   
@r   c                 C   sN   | j tjtjfkr| S | j tjkr>| j tjkr:| jd } q"| S td|  d S )Nr   zReceived unexpected node %s)r   r
   Zvfplistr   r>   vfpdefr   RuntimeErrorr0   r   r   r   r=      s    r=   c                 C   s<   | j tjkrt| jd S | j tjkr,| jS dd | jD S )Nr   c                 S   s    g | ]}|j tjkrt|qS r   )r   r   COMMAr?   r9   r   r   r   r<      s      zfind_params.<locals>.<listcomp>)r   r
   rF   r?   r   r   r>   r+   rH   r   r   r   r?      s
    r?   Nc                 C   sZ   |d kri }t | D ]@\}}ttt|g}t|trHt|||d q|| ||< q|S )N)d)r,   r	   r   strr   listr@   )
param_listr   rJ   r4   objZtrailerr   r   r   r@      s    
r@   c                 C   s<   g }| D ](}t |tr&|t| q|| qd|S )N_)r   rL   r!   rA   join)rM   lrN   r   r   r   rA      s    
rA   )__doc__r   r   Zpgen2r   r   Z
fixer_utilr   r   r   r   r	   r
   r   ZBaseFixr   r=   r?   r@   rA   r   r   r   r   <module>   s    l