3
YdY                 @   s  d Z ddl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 dd	lmZmZ G d
d dejZG dd dejZG dd dejZG dd dejZG dd dejZG dd dejZg ag aG dd dejZG dd dejZG dd dejZdaG dd dejZG dd dejZ G d d! d!ejZ!g a"G d"d# d#ejZ#G d$d% d%ejZ$G d&d' d'ejZ%ej&j'd(d)d*gd+d, Z(d-d. Z)d/d0 Z*d1d2 Z+d3d4 Z,d5d6 Z-d7d8 Z.d9d: Z/d;d< Z0d=d> Z1d?d@ Z2dAdB Z3dCdD Z4dEdF Z5dS )GzTests for join expansion
    N   )config   )engine)base)IdentityInterfaceFunctionMerge)traitsFilec               @   s$   e Zd ZejeddddddZdS )PickFirstSpecT)existsz%s   )Zargstrposition	mandatoryN)__name__
__module____qualname__r
   Listr   in_files r   r   X/var/www/html/virt/lib/python3.6/site-packages/nipype/pipeline/engine/tests/test_join.pyr      s   r   c               @   s   e Zd ZeddZdS )PickFirstOutSpecT)r   N)r   r   r   r   output1r   r   r   r   r      s   r   c               @   s$   e Zd ZeZeZdd Zdd ZdS )	PickFirstc             C   s
   d|_ |S )Nr   )
returncode)selfruntimer   r   r   _run_interface   s    zPickFirst._run_interfacec             C   s    | j  j }| jjd |d< |S )Nr   r   )_outputsgetinputsr   )r   outputsr   r   r   _list_outputs   s    zPickFirst._list_outputsN)	r   r   r   r   
input_specr   output_specr   r#   r   r   r   r   r      s   r   c               @   s.   e Zd ZejjdddZejjddddZdS )IncrementInputSpecTinput)r   desc   	increment)Z
usedefaultdefault_valuer(   N)r   r   r   nibr
   Intinput1incr   r   r   r   r&   %   s   r&   c               @   s   e Zd ZejjddZdS )IncrementOutputSpecouput)r(   N)r   r   r   r,   r
   r-   r   r   r   r   r   r0   *   s   r0   c               @   s   e Zd ZeZeZdd ZdS )IncrementInterfacec             C   s    d|_ | jj| jj | jd< |S )Nr   r   )r   r!   r.   r/   _results)r   r   r   r   r   r   2   s    z!IncrementInterface._run_interfaceN)r   r   r   r&   r$   r0   r%   r   r   r   r   r   r2   .   s   r2   c               @   s"   e Zd ZejjejjdddZdS )SumInputSpecTr'   )r   r(   N)r   r   r   r,   r
   r   r-   r.   r   r   r   r   r4   =   s   r4   c               @   s.   e Zd ZejjddZejjejjddZdS )SumOutputSpecr1   )r(   operandsN)	r   r   r   r,   r
   r-   r   r   r6   r   r   r   r   r5   A   s   r5   c               @   s   e Zd ZeZeZdd ZdS )SumInterfacec             C   sJ   d|_ | jj| jd< t| jj| jd< tj| jj tjt| jj |S )Nr   r6   r   )r   r!   r.   r3   sum_sum_operandsappend_sums)r   r   r   r   r   r   J   s    zSumInterface._run_interfaceN)r   r   r   r4   r$   r5   r%   r   r   r   r   r   r7   F   s   r7   c               @   s"   e Zd ZejjejjdddZdS )SetInputSpecTr'   )r   r(   N)r   r   r   r,   r
   Setr-   r.   r   r   r   r   r<   Y   s   r<   c               @   s   e Zd ZejjddZdS )SetOutputSpecr1   )r(   N)r   r   r   r,   r
   r-   r   r   r   r   r   r>   ]   s   r>   c               @   s$   e Zd ZeZeZdd Zdd ZdS )SetInterfacec             C   s
   d|_ |S )Nr   )r   )r   r   r   r   r   r   e   s    zSetInterface._run_interfacec             C   s$   | j  j }t| jj a|d< |S )Nr   )r   r    lenr!   r.   _set_len)r   r"   r   r   r   r#   i   s    zSetInterface._list_outputsN)	r   r   r   r<   r$   r>   r%   r   r#   r   r   r   r   r?   a   s   r?   c               @   s,   e Zd ZejjdddZejjdddZdS )ProductInputSpecTr.   )r   r(   input2N)r   r   r   r,   r
   r-   r.   rC   r   r   r   r   rB   t   s   rB   c               @   s   e Zd ZejjdddZdS )ProductOutputSpecToutput)r   r(   N)r   r   r   r,   r
   r-   r   r   r   r   r   rD   y   s   rD   c               @   s$   e Zd ZeZeZdd Zdd ZdS )ProductInterfacec             C   s
   d|_ |S )Nr   )r   )r   r   r   r   r   r      s    zProductInterface._run_interfacec             C   s2   | j  j }| jj| jj |d< tj|d  |S )Nr   )r   r    r!   r.   rC   	_productsr:   )r   r"   r   r   r   r#      s    zProductInterface._list_outputsN)	r   r   r   rB   r$   rD   r%   r   r#   r   r   r   r   rF   }   s   rF   needed_outputstruefalsec          
   C   s  | j   g ag ag atjdd}tjdd| tjdd}tj	t
dgddd}ddd	gfg|_tj	t d
d}tj	t dd}tjt dddd}tj	t dd}tj	t dd}	|j||d gf||d!gf||	d"gf||d#gf||d$gf||	d%gfg |j }
dd |
j D }t|dks*tdt|
j dksDtdttdksZtdtd dksxtdtd  td ddgkstdtd  ttd	kstdtjdd| d S )&NZ	executionZremove_unnecessary_outputstest)namen)fields	inputspecr)   r   	pre_join1	pre_join2r.   join)
joinsource	joinfieldrL   
post_join1
post_join2r   rC   c             S   s   g | ]}|j d kr|qS )rR   )rL   ).0noder   r   r   
<listcomp>   s    z'test_join_expansion.<locals>.<listcomp>z-The number of join result nodes is incorrect.   z*The number of expanded nodes is incorrect.z'The number of join outputs is incorrectr      z+The join Sum output value is incorrect: %s.r   r   z$The join Sum input is incorrect: %s.z5The number of iterated post-join outputs is incorrect)rM   r.   )r   r.   )r   rC   )r   r.   )r   r.   )r   r.   )chdirrG   r9   r;   r   r    setpeWorkflowNoder   	iterablesr2   JoinNoder7   rF   connectrunnodesr@   AssertionError)tmpdirrH   Z
prev_statewfrO   rP   rQ   rR   rU   rV   resultZjoinsr   r   r   test_join_expansion   s@    




rj   c             C   sZ   | j   tjtdgddd}dddgfg|_tjt |ddd	}|j|jksVt	d
dS )z&Test setting the joinsource to a Node.rM   )rN   rO   )rL   r)   r   r.   rR   )rS   rT   rL   z+The joinsource is not set to the node name.N)
r\   r^   r`   r   ra   rb   r?   rS   rL   rf   )rg   rO   rR   r   r   r   test_node_joinsource   s    rk   c             C   s   | j   tjdd}tjtdgddd}ddddddgfg|_tjt d	d}|j|d|d
 tjt	 dd
dd}|j|d|d
 |j
  tdkstdt dS )z%Test collecting join inputs to a set.rK   )rL   rM   )rN   rO   r)   r   r   rP   r.   rR   )rS   rT   rL   r   z+The join Set output value is incorrect: %s.N)r\   r^   r_   r`   r   ra   r2   rc   rb   r?   rd   rA   rf   )rg   rh   rO   rP   rR   r   r   r   test_set_join_node   s    rl   c             C   s   g a | j  tjdd}tjtdgddd}ddddddgfg|_tjt d	d}|j|d|d
 tj	t
 dd
ddd}|j|d|d
 |j  t d dddgkstdt d  dS )z/Test join with the ``unique`` flag set to True.rK   )rL   rM   )rN   rO   r   r)   r   rP   r.   TrR   )rS   rT   uniquerL   r   r   r   z.The unique join output value is incorrect: %s.N)r9   r\   r^   r_   r`   r   ra   r2   rc   rb   r7   rd   rf   )rg   rh   rO   rP   rR   r   r   r   test_unique_join_node   s"    rn   c       
      C   sp  g a | j  tjdd}tjtdgddd}ddddgfg|_tjt d	d}|j|d|d
 tj	tdgddddd}|j|d|d tjt
 dd}|j|d|d
 tj	tddgddddd}|j|d|d |j|d|d tjt
 dd}|j|d|d
 tjt dd}|j|d|d
 |j|d|d |j }	t|	j dksXtdt dgksltddS )z1Test two join nodes, one downstream of the other.rK   )rL   rM   )rN   rO   r)   r   r   rP   r.   vectorjoin1)rS   rT   rL   r   rU   Zscalarjoin2rV   
post_join3rC   rZ   z*The number of expanded nodes is incorrect.Q   z"The post-join product is incorrectN)rG   r\   r^   r_   r`   r   ra   r2   rc   rb   r7   rF   rd   r@   re   rf   )
rg   rh   rO   rP   rp   rU   rq   rV   rr   ri   r   r   r   test_multiple_join_nodes  s<    
rt   c             C   s   g a | j  tjdd}tjtdgddd}ddddgfg|_tjt d	d}|j|d|d
 tj	tdgddddd}|j|d|d tjt
 dd}|j|d|d
 |j }t|j dkstdt d dddgkstdt d  dS )zTest an IdentityInterface join.rK   )rL   rM   )rN   rO   r)   r   r   rP   r.   ro   rR   )rS   rT   rL   r   rU      z*The number of expanded nodes is incorrect.r   r   z$The join Sum input is incorrect: %s.N)r9   r\   r^   r_   r`   r   ra   r2   rc   rb   r7   rd   r@   re   rf   )rg   rh   rO   rP   rR   rU   ri   r   r   r   test_identity_join_nodeP  s&    
rv   c             C   sF  g a | j  tjdd}tjtddgddd}dddgfdd	d
gfg|_tjt dd}|j|d|d tjt dd}|j|d|d tj	tddgdddd}|j|d|d |j|d|d tj
t dddgd}|j|d|d |j|d|d |j }t|j dkstdtt tddddgksBtdt  dS )zTest join on several fields.rK   )rL   mrM   )rN   rO   r)   r   r   r   inc1r.   inc2vector1vector2rR   )rS   rL   r   prodrC   )rL   	iterfield
   z*The number of expanded nodes is incorrect.rZ         z(The post-join products is incorrect: %s.N)rG   r\   r^   r_   r`   r   ra   r2   rc   rb   MapNoderF   rd   r@   re   rf   r]   )rg   rh   rO   rx   ry   rR   r|   ri   r   r   r   test_multifield_join_nodeu  s,    r   c             C   s@  g a | j  tjdd}tjtddgddd}dddgfdd	d
gfg|_d|_tjt dd}|j	|d|d tjt dd}|j	|d|d tj
tddgdddd}|j	|d|d |j	|d|d tjt dddgd}|j	|d|d |j	|d|d |j }t|j dks"tdt ddgks<tdt  dS )zJTest join on an input node which has the ``synchronize`` flag set to True.rK   )rL   rw   rM   )rN   rO   r)   r   r   r   Trx   r.   ry   rz   r{   rR   )rS   rL   r   r|   rC   )rL   r}      z*The number of expanded nodes is incorrect.rZ   r   z(The post-join products is incorrect: %s.N)rG   r\   r^   r_   r`   r   ra   synchronizer2   rc   rb   r   rF   rd   r@   re   rf   )rg   rh   rO   rx   ry   rR   r|   ri   r   r   r   test_synchronize_join_node  s,    r   c       	      C   s\  | j   tjdd}tjtdgddd}dddgfg|_tjt dd}|j|d|d	 tjt d
d}d|_	d	ddgddgdf|_|j|d|d tjt dd}|j|d|d	 tj
tdgdd
ddd}|j|d|d tjt dd}|j|d|d	 |j }t|j dks$tdddgtks>tdt ddgtksXtdt dS ) z7Test join on an input node which has an ``itersource``.rK   )rL   rM   )rN   rO   r)   r   rP   r.   rQ   r   r   ru   r   )r)   r   r   rC   	pre_join3ro   rR   )rS   rT   rL   rU      z*The number of expanded nodes is incorrect.      z$The join Sum input is incorrect: %s.r[   	   N)rO   rM   )r\   r^   r_   r`   r   ra   r2   rc   rF   
itersourcerb   r7   rd   r@   re   rf   r9   )	rg   rh   rO   rP   rQ   r   rR   rU   ri   r   r   r    test_itersource_join_source_node  s4    

r   c       
      C   sR  | j   tjdd}tjtdgddd}dddgfg|_tjt dd}|j|d|d	 tjt d
d}d|_	d	ddgddgdf|_|j|d|d tjt dd}|j|d|d	 tj
tdgdd
ddd}|j|d|d tjt dd}|j|d|d	 tj
tdgddddd}|j|d|d |j }	t|	j dksNtddS )zGTest join with a midstream ``itersource`` and an upstream
    iterable.rK   )rL   rM   )rN   rO   r)   r   rP   r.   rQ   r   r   ru   r   )r)   r   r   rC   r   ro   rp   )rS   rT   rL   rU   rq   r   z*The number of expanded nodes is incorrect.N)rO   rM   )r\   r^   r_   r`   r   ra   r2   rc   rF   r   rb   r7   rd   r@   re   rf   )
rg   rh   rO   rP   rQ   r   rp   rU   rq   ri   r   r   r   test_itersource_two_join_nodes   s8    

r   c             C   s   | j   tddj  tddj  tjdd}tjtdgddd}d| jdj| jdjgfg|_	tjtdgdd	d}|j
|d|d tjt dd
dd}|j
|d|d
 |j  dS )z%Test collecting join inputs to a set.ztest.niizw+z	test2.niirK   )rL   rM   )rN   rO   rP   r   rR   )rS   rT   rL   N)r\   opencloser^   r_   r`   r   rR   strpathra   rc   rb   r   rd   )rg   rh   rO   rP   rR   r   r   r   test_set_join_node_file_input0  s     r   c             C   sr   | j   ddd}tjddd}x,ddgd	d
ggD ]}||}|j|g q2W |j }t|j dksntddS )z4Test collecting join inputs within a nested workflowsmallwfc             S   s   t jtdgddd}d| fg|_t jt dd}t jtdgddddd}t jd| d	  d}|j|d|d
 |j|d|d |S )NrM   )rN   rO   )rL   pre_joinrR   )rS   rT   rL   zwf_%dr   r.   r   )r^   r`   r   ra   r2   rb   r_   rc   )irL   rO   r   rR   rh   r   r   r   	nested_wfN  s    
z,test_nested_workflow_join.<locals>.nested_wfmeta.)rL   base_dirr)   r   r   r   r   z*The number of expanded nodes is incorrect.N)r   )r\   r^   r_   Z	add_nodesrd   r@   re   rf   )rg   r   Zmeta_wfr   Zmini_wfri   r   r   r   test_nested_workflow_joinI  s    
r   c             C   s|   | j   dd }tjd| jd}tjt|ddd}dd	d
gfg|_tjtd	dddddgd}|j	|d|d |j
  d S )Nc             S   s   | d S )Nr   r   )xr   r   r   sqq  s    z!test_name_prefix_join.<locals>.sqrh   )r   )functionsquare)rL   r   r)   r   T)ravel_inputssquare_joinin1)rL   rS   rT   out)r\   r^   r_   r   r`   r   ra   rb   r	   rc   rd   )rg   r   rh   r   r   r   r   r   test_name_prefix_joinn  s    

r   c             C   s   | j   dd }tjd| jd}tjtdgdddgfgdd	}tjtd
gd
ddgfgdd	}tjt|ddd}tjtddddddgd}tjtddddddgd}|j	||dgf||dgf||dgf||dgfg |j
  d S )Nc             S   s   | | S )Nr   )r   pr   r   r   exponent  s    z'test_join_nestediters.<locals>.exponentrh   )r   r   r)   r   xs)ra   rL   r   r   r   ps)r   exp)rL   T)r   	exp_joinxr   )rL   rS   rT   	exp_joinpr   )r   r   )r   r   )r   r   )r   r   )r\   r^   r_   r   r`   r   r   rb   r	   rc   rd   )rg   r   rh   r   r   r   r   r   r   r   r   test_join_nestediters  s,      






r   )6__doc__Zpytest r   r   r^   Z
interfacesr   r,   Zinterfaces.utilityr   r   r	   Zinterfaces.baser
   r   ZTraitedSpecr   r   ZBaseInterfacer   r&   r0   ZSimpleInterfacer2   r;   r9   r4   r5   r7   rA   r<   r>   r?   rG   rB   rD   rF   markZparametrizerj   rk   rl   rn   rt   rv   r   r   r   r   r   r   r   r   r   r   r   r   <module>   sN   
@:%)(:0%