B
    T\7B                 @   s  d dl mZ d dlmZmZ d dlZd dlmZ d dlZd dlZd dl	m
Z
mZmZ d dlZd dlZd dlmZ d dlmZmZ d dlmZmZmZ d d	lmZ yd d
lmZ W n ek
r   dZY nX G dd de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'd#d$ Z(d%d& Z)d'd( Z*d)d* Z+d+d, Z,d-d. Z-d/d0 Z.d1d2 Z/d3d4 Z0d5d6 Z1d7d8 Z2d9d: Z3d;d< Z4d=d> Z5d?d@ Z6dAdB Z7dCdD Z8dEdF Z9dGdH Z:dIdJ Z;dKdL Z<dMdN Z=dOdP Z>dQdR Z?dSdT Z@dUdV ZAdS )W    )
namedtuple)addsetitemN)random)identitypartialmerge)compute)PY2PY3)delayedto_task_daskDelayed)inc)matmulc               @   s@   e Zd ZeejjZdd Zdd Z	dd Z
dd Zd	d
 ZdS )Tuplec             C   s   || _ || _d S )N)_dask_keys)selfdskkeys r   6lib/python3.7/site-packages/dask/tests/test_delayed.py__init__   s    zTuple.__init__c             C   s   | j S )N)r   )r   r   r   r   __dask_tokenize__   s    zTuple.__dask_tokenize__c             C   s   | j S )N)r   )r   r   r   r   __dask_graph__!   s    zTuple.__dask_graph__c             C   s   | j S )N)r   )r   r   r   r   __dask_keys__$   s    zTuple.__dask_keys__c             C   s   t dfS )Nr   )tuple)r   r   r   r   __dask_postcompute__'   s    zTuple.__dask_postcompute__N)__name__
__module____qualname__staticmethoddaskZthreadedgetZ__dask_scheduler__r   r   r   r   r   r   r   r   r   r      s   r   c           	   C   s4  t jdd tddd} tddd}t| |dg\}}|dddgksLtt| |df\}}|tdddgfksttt|t| j|jkstt| d|di\}}|tddgddggfks|tddgddggfkstt|t| j|jkstt	d	d
dg}|dd}t|\}}||ks"tt|i ks4ttt
| |d\}}|t
dddfks^tt|t| j|jksztG dd dt}t| \}}t||kstt|i ksttddtddfddddg}t|\}}||kst||}|tdddgfkst||jks&tW d Q R X d S )NT)record   a)name   b   fxyc               @   s   e Zd ZdS )z"test_to_task_dask.<locals>.MyClassN)r   r    r!   r   r   r   r   MyClassF   s   r/   )r'   r*   cr0   )warningscatch_warningsr   r   AssertionErrorr   dictr   r#   r   slicetyper   r   popr   )r'   r*   taskr#   r,   r-   r/   r   r   r   test_to_task_dask+   s<    

r9   c              C   s   t t} | dd dkst| ddd  dks8t| | ddd dksTtt d}| dksltd|j ks~t| | |dd}|j|jkstd S )Nr&   r)   r+      )r   r   r	   r3   r#   valueskey)Zadd2r'   r*   r   r   r   test_delayedV   s    r=   c              C   sb   t d} | ddtfg}td}td||di}dd }t||}| dks^td S )Ndataclasses
ADataClassr'   r+   )r'   c             S   s
   | d j S )Nr'   )r'   )objr   r   r   return_nestedl   s    z2test_delayed_with_dataclass.<locals>.return_nested)pytestimportorskipZmake_dataclassintr#   r   r	   r3   )r>   r?   literalZ
with_classrA   finalr   r   r   test_delayed_with_dataclassc   s    

rG   c              C   s  t dddg} | d  dks"t| |   ddddddgksBtt d}| d |  ddgksftt d} | d  dkstd|   dkst| d?  dkst| dk st| d  dksttrG d	d
 d
}t | }t | }td dkstd S )Nr&   r)   r+   r   
         d   c               @   s   e Zd Zdd ZdS )ztest_operators.<locals>.dummyc             S   s   dS )N   r   )r   otherr   r   r   
__matmul__   s    z(test_operators.<locals>.dummy.__matmul__N)r   r    r!   rN   r   r   r   r   dummy   s   rO   zc @ drL   )r   r	   r3   r   eval)r'   r*   rO   r0   dr   r   r   test_operatorst   s      

rR   c              C   s   t d} | d dddddgks(t|  dd	 d	 d
ksNt| jdddj| jdddjksrt| jddd}|jdkstd S )Nz	a b c d e r'   r*   r0   rQ   eBAr)   T)pureZtest)dask_key_name)r   splitr	   r3   upperreplacecountr<   )r'   or   r   r   test_methods   s     &$r^   c              C   s\   t d} | jj| jjkst| j dks.t| j dks@t| j| j  dksXtd S )Ny       @      ?r)   r&   r+   )r   real_keyr3   r	   imag)r'   r   r   r   test_attributes   s
    rb   c              C   s>   t dddg} | d}ttdd |  D ks:td S )Nr&   r)   r+   c             s   s   | ]}|d  V  qdS )r   Nr   ).0vr   r   r   	<genexpr>   s    z5test_method_getattr_call_same_task.<locals>.<genexpr>)r   indexgetattrsetr   r;   r3   )r'   r]   r   r   r   "test_method_getattr_call_same_task   s    
ri   c           	   C   s^   t d} td}t t | | W d Q R X t| jdgddj | dksZtd S )Nnumpyr&   Zf8)dtype)	rB   rC   r   raises	TypeErrorrk   Zarrayr	   r3   )npr-   r   r   r   test_np_dtype_of_delayed   s
    
ro   c                  s   t dddg tt fdd tt fdd tt fdd tt fdd tt fd	d tt fd
d d S )Nr&   r)   r+   c                  s   t  ddS )Nfoor&   )setattrr   )r'   r   r   <lambda>   s    z%test_delayed_errors.<locals>.<lambda>c                  s   t  ddS )Nr&   r   )r   r   )r'   r   r   rr      s    c                  s   d kS )Nr&   r   r   )r'   r   r   rr      s    c                  s   t  S )N)listr   )r'   r   r   rr      s    c                  s      S )N)Z_hiddenr   )r'   r   r   rr      s    c                  s   t  S )N)boolr   )r'   r   r   rr      s    )r   rB   rl   rm   AttributeErrorr   r   )r'   r   test_delayed_errors   s    rv   c              C   sX   t dddg} | d | d  }| d j|jks2t| j|jksBtt|jdksTtd S )Nr&   r)   r+   r   )r   r<   r#   r3   len)r'   resr   r   r   test_common_subexpressions   s
    ry   c              C   sD   t ddtdftdfd} t| \}t|j ddgks@td S )Nr*   r&   r'   )r'   r*   r0   )r   r   r#   optimizesortedr   r3   )r-   x2r   r   r   test_delayed_optimize   s
    r}   c              C   s4   t d} t d}t t| |g}| dks0td S )Nr&   r)   r+   )r   sumr	   r3   )r'   r*   r0   r   r   r   
test_lists   s    r   c              C   s   t d} | d }| |df}t | dks.t| |dg}t | dddgksRtt| |df}t | tdksxt| d|dddi}t | ddddkstt ||   dkst| |dd}t | ddddkstt |d  dkstd S )	Nr&   r+   )r&   r)   r+   r)   r'   r*   r0   )r'   r*   r0   )r   r	   r3   rh   )r'   r*   litr   r   r   test_literates   s    

r   c              C   sV   t d} | d }| |df}t |jt |jks2tt |ddjt |ddjksRtd S )Nr&   r+   T)rW   )r   r<   r3   )r'   r*   r   r   r   r   test_literates_keys   s
    
r   c              C   sH   t d} t d}t t| dg|dggdd dd }| dksDtd S )Nr&   r)   rH      c             S   s   | d S )Nr   r   )r-   r   r   r   rr      s    z)test_lists_are_concrete.<locals>.<lambda>)r<   )r   maxr	   r3   )r'   r*   r0   r   r   r   test_lists_are_concrete   s    $r   c              C   sd   t d} t d}t tt| |g}| dks4tdd }t |t| |g}| dks`td S )Nr&   r)   r+   c             S   s   t | S )N)r~   )seqr   r   r   r,      s    ztest_iterators.<locals>.f)r   r~   iterr	   r3   )r'   r*   r0   r,   r   r   r   test_iterators   s    r   c              C   sV  dd } t |  }|dddg}t |dd }t|dks>t|d	 |ksNt|dd  |dd  ksjt| || |ff}t |dd }t|tst|d	 | kst|d |kstd| |f|g}t |dd }t|tst|d	 dkst|d d	 | kr|d d |kst|d |ks,tt d}t |dd}| dksRtd S )
Nc              W   s   t dd S )Nzshouldn't have computed)
ValueError)argsr   r   r   fail  s    z!test_traverse_false.<locals>.failr&   r)   r+   F)ZtraverserL   r   )r   r	   rw   r3   
isinstancer   rs   )r   r'   r-   rx   r*   r   r   r   test_traverse_false  s*    
(r   c              C   sT   t tdddd} t tdddd}| j|jks4tt t}| j| jksPtd S )NT)rW   r&   r)   )r   r   r<   r3   r   )Zv1Zv2Zmyrandr   r   r   	test_pure&  s
    r   c           	   C   s\  t t} tjjdd" | ddj| ddjks4tW d Q R X tjjdd" | ddj| ddjksjtW d Q R X t tdd} tjjdd" | ddj| ddjkstW d Q R X t djt djksttjjdd t djt djkstW d Q R X tjjdd( t dddjt dddjks2tW d Q R X t dddg}|dj|djkshttjjddJ |dj|djkst|jdddj|jdddjkstW d Q R X tjjdd, |jdddj|jdddjkstW d Q R X tjjdd> |jj|jjks,t|d }|| j|| jksNtW d Q R X d S )	NT)Zdelayed_purer&   r)   F)rW   r+   r   )r   r   r#   Zconfigrh   r<   r3   rf   )funcdataelementr   r   r   test_pure_global_setting/  s0    &&&",00r   c                 s  t dd ddd} | d}t|dks*t|\ }t |dksDt jd ksRt|jd ks`ttt fdd tt fd	d ttd
d  ttdd  t t	d d} | d  jd ksttt fdd tt fdd t dd ddd} | d}t|dkst|\   dks6t jd ksFttt fdd t dd ddd} | d}t|dkst| t
 kstd S )Nc             S   s
   | |  fS )Nr   )r-   r   r   r   rr   X  s    ztest_nout.<locals>.<lambda>r)   T)noutrW   r&   )r&   c                  s   t  S )N)rw   r   )r'   r   r   rr   _  s    c                  s   t  S )N)rs   r   )r'   r   r   rr   `  s    c               S   s   t tddS )Nr   )r   )r   r   r   r   r   r   rr   b  s    c               S   s   t tddS )NT)r   )r   r   r   r   r   r   rr   c  s    )r   c                  s   t  S )N)rs   r   )r'   r   r   rr   h  s    c                  s   t  S )N)rw   r   )r'   r   r   rr   i  s    c             S   s   | fS )Nr   )r-   r   r   r   rr   k  s    c                  s   t  S )N)rw   r   )r'   r   r   rr   q  s    c             S   s   t  S )N)r   )r-   r   r   r   rr   s  s    r   )r   rw   r3   r	   _lengthrB   rl   rm   r   r   r   )r   r-   r*   r   )r'   r   	test_noutW  s6    r   c              C   s  ddd} t | }|ddt ddg|ddd}| d	ks@tt | d
d}t ddg}|dd||ddd}| d	ks~t|dd||dddj|jkst|dd||dddj|jkst|dd|ddj|jkst|dd|ddj|dd|ddjkstd S )Nr   c             [   s   | | t | t |  S )N)r~   r;   )r'   r*   r0   kwargsr   r   r   mysumz  s    ztest_kwargs.<locals>.mysumr&   r)   r+   r   )r0   ZfourrH   T)rW   rL   )r   )r   r	   r3   r<   )r   ZdmysumZtenr0   r   r   r   test_kwargsy  s    
""r   c              C   s   t ddtddfddddg} ttdd| d	}ttdd| }ttdd| j|jks\t| d
ksltt||| dkstd S )Nr&   r)   r'   r*   )r'   r*   r0   r0   T)rW   )rL   rJ   r:   )r&   r)   r+   rL   rJ   r:   )r+   )r&   r)   r+   rL   rJ   r:   )r&   r)   r+   )r   r   r   rw   r<   r3   r	   )r-   r|   nr   r   r   test_custom_delayed  s    r   c           	   C   s0  t d} t d}| dd}|j|dd}tt||dg}t|tsRt	| 
| || d snt	|  || d  kst	|d  || d d kst	tjd	d
 t|\}}W d Q R X t|j}t|}||st	||}	t|	dkst	t|}
|
 |k s,t	d S )Nrj   z
dask.arrayrK   )rH   rH   )rJ   rJ   )chunksr&   )r   r   T)r%   )rB   rC   arangereshape
from_arrayr   r~   r   r   r3   Zallcloser	   r1   r2   r   rh   r#   issubset
differencerw   all)rn   daZarrZdarrvalr8   r   ZorigrF   ZdiffZdelayed_arrr   r   r   test_array_delayed  s$    

  

r   c        
      C   s   t d} t d}t d}|dd}||j}|j|dd}|j|dd}| dd	d
g}|||||g}tt	dd |D }	|	
 d	|	  d	|	   t	dd	d
g kstd S )Nzdask.bagz
dask.arrayrj   rK   )rH   rH   )rJ   rJ   )r   r&   r)   r+   c             S   s   g | ]}|  qS r   )r~   )rc   ir   r   r   
<listcomp>  s    z*test_array_bag_delayed.<locals>.<listcomp>)rB   rC   r   r   dotTr   Zfrom_sequencer   r~   r	   r3   )
Zdbr   rn   Zarr1Zarr2Zdarr1Zdarr2r*   r   outr   r   r   test_array_bag_delayed  s    


r   c              C   s  t tddddd} tt| }| j|jks4t| j|jksDt| j|jksTtt d} tt| }| j|jks|t| j|jkst| j	|j	kst| j
|j
kst| j} tt| }| jj|jjkst| jj|jjkst| j|jkst| j|jkstd S )Nr)   T)r   rW   r&   y       @      ?)r   divmodpickleloadsdumpsr#   r3   r`   r   Z_noutZ_purer_   Z_objZ_attr)r-   r.   r   r   r   test_delayed_picklable  s"    r   c              C   s   t dd } | jdd d S )Nr&   r)   rH   )Zbogus_keyword)r   r	   )r-   r   r   r   #test_delayed_compute_forward_kwargs  s    r   c               C   s   t tjd d S )N    )r   bytesdecoder   r   r   r   test_delayed_method_descriptor  s    r   c              C   sX   t tdd} | dd}|j|jtddfiks0t| j| jtiksDt|  tksTtd S )NT)rW   r&   r)   )r   r   r#   r<   r3   r	   )r,   rd   r   r   r   test_delayed_callable  s
    
r   c              C   s(   t tdd} | ddddjdks$td S )NT)rW   r&   r)   rp   )rX   )r   r   r`   r3   )r,   r   r   r   test_delayed_name_on_call  s    r   c              C   sX   G dd dt } | d}t|}| |ks0t|j dksBt|  dksTtd S )Nc               @   s   e Zd Zdd Zdd ZdS )ztest_callable_obj.<locals>.Fooc             S   s
   || _ d S )N)r'   )r   r'   r   r   r   r     s    z'test_callable_obj.<locals>.Foo.__init__c             S   s   dS )Nr)   r   )r   r   r   r   __call__  s    z'test_callable_obj.<locals>.Foo.__call__N)r   r    r!   r   r   r   r   r   r   Foo  s   r   r&   r)   )objectr   r	   r3   r'   )r   rp   r,   r   r   r   test_callable_obj  s    r   c              C   s   t tdd} dddddgd}tr4| |jdks4ttrJ| |jd	ksJtdd
d}| |j| |jkslttr| djdksttr| djdkstd S )NT)rW   r&      r)   r+   )r-   r.   zz)identity-6700b857eea9a7d3079762c9a253ffbdz)identity-84c5e2194036c17d1d97c4e3a2b90482r-   )r-   r&   z)identity-91f02358e13dca18cde218a63fee436az)identity-7126728842461bf3d2caecf7b954fa3b)r   r   r
   r`   r3   r   )r   r   r   r   r   %test_name_consistent_across_instances  s    
r   c               C   s8   t ttddddjt ttddddjks4td S )NrH   T)rW   r)   r   )r   r   r   r`   r3   r   r   r   r   test_sensitive_to_partials  s    r   c              C   sf   t djdstt dddjds,tt dddjdksBtdd } t | djd	sbtd S )
Nr&   zint-T)rW   X)r(   c             S   s   | d S )Nr&   r   )r-   r   r   r   myfunc  s    z!test_delayed_name.<locals>.myfuncr   )r   r`   
startswithr3   r<   )r   r   r   r   test_delayed_name  s
    r   c                 s`   t d} | jddd}t|g}t|j|js8tdd  t fdd|jD s\td S )	Nz
dask.arrayrH   rJ   )r   c             S   s    t | tr| d } | dd S )Nr   -)r   r   rY   )sr   r   r   r<     s    
ztest_finalize_name.<locals>.keyc             3   s   | ]} |  V  qd S )N)isalpha)rc   k)r<   r   r   re   $  s    z%test_finalize_name.<locals>.<genexpr>)	rB   rC   onesr   rh   r#   r   r3   r   )r   r-   rd   r   )r<   r   test_finalize_name  s    

r   c              C   sL   t d} ddlm} | jddd  }dd |D }||d j d S )	Nz
dask.arrayr   )
_check_dsk)rH   rH   rJ   )r   c             S   s   g | ]}t t|qS r   )r   r   )rc   r-   r   r   r   r   ,  s    z(test_keys_from_array.<locals>.<listcomp>)rB   rC   Zdask.array.utilsr   r   Z
to_delayedZflattenr#   )r   r   r   Zxsr   r   r   test_keys_from_array'  s
    
r   c              C   s   G dd dt } | d}|dd dks0t| |dd dksJt|dd dksbt| dd dkszt|dd dkst| dd dkstt|jtjstt| jtjstt| jt	std S )	Nc               @   sD   e Zd ZdZdd Zedd Zeedd Ze	edd	 Z
d
S )z+test_delayed_decorator_on_method.<locals>.ArH   c             S   s
   || _ d S )N)BASE)r   baser   r   r   r   6  s    z4test_delayed_decorator_on_method.<locals>.A.__init__c             S   s   | j | | S )N)r   )r   r-   r.   r   r   r   	addmethod9  s    z5test_delayed_decorator_on_method.<locals>.A.addmethodc             S   s   | j | | S )N)r   )clsr-   r.   r   r   r   addclass=  s    z4test_delayed_decorator_on_method.<locals>.A.addclassc             S   s   | | S )Nr   )r-   r.   r   r   r   	addstaticB  s    z5test_delayed_decorator_on_method.<locals>.A.addstaticN)r   r    r!   r   r   r   r   classmethodr   r"   r   r   r   r   r   rV   3  s   rV   rK   r+   rL   k         )
r   r   r	   r3   r   r   r   types
MethodTyper   )rV   r'   r   r   r    test_delayed_decorator_on_method2  s    r   c              C   sB   t d} t| jtstt| jjts*tt| jjjts>td S )N{   )r   r   r'   r   r3   r*   r0   )r-   r   r   r   test_attribute_of_attribute]  s    r   )Bcollectionsr   operatorr   r   r   r   r   r1   Ztoolzr   r   r   rB   r#   r	   Zdask.compatibilityr
   r   Zdask.delayedr   r   r   Zdask.utils_testr   r   ImportErrorr   r   r9   r=   rG   rR   r^   rb   ri   ro   rv   ry   r}   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   <module>   sl   
+	
	#	("		+