
    CdY                    N   d dl 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
 d dlmZmZ d dlmZ d dlmZ d dlZd dlZd d	lmZ d dl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 d dl m!Z!  G d d          Z"ej#        $                    d          d             Z%d Z&d Z'e G d d                      Z( ed           G d d                      Z)ej#        *                    de(e)f          d             Z+d Z,d Z-d Z.d Z/d  Z0d! Z1d" Z2d# Z3d$ Z4d% Z5d& Z6d' Z7d( Z8d) Z9d* Z:d+ Z;d, Z<d- Z=d. Z>d/ Z?d0 Z@d1 ZAd2 ZBej#        *                    d3d4d5gd6ed4d5fg d7g          d8             ZCd9 ZDd: ZEej#        $                    d          d;             ZFd< ZGd= ZHd> ZId? ZJd@ ZKdA ZLdB ZMdC ZNdD ZOdE ZPdF ZQdG ZRdH ZSdI ZTdJ ZUdK ZVdL ZWedM             ZX edNO          dP             ZYej#        *                    dQ eeW           ejZ        eXej#        [                    dRS          T           ejZ        eYej#        [                    dRS          T          g          dU             Z\ej#        *                    dQ eeW          eXeYg          dV             Z]dW Z^dX Z_dY Z`dS )Z    )annotationsN)
namedtuple)	dataclassfield)partial)addsetitem)random)
NamedTuple)merge)compute)Delayeddelayedto_task_dask)HighLevelGraph)incc                  V    e Zd Z eej        j                  Zd Zd Z	d Z
d Zd ZdS )Tuplec                "    || _         || _        d S N)_dask_keys)selfdskkeyss      7lib/python3.11/site-packages/dask/tests/test_delayed.py__init__zTuple.__init__   s    



    c                    | j         S r   r   r   s    r   __dask_tokenize__zTuple.__dask_tokenize__   
    zr   c                    | j         S r   )r   r!   s    r   __dask_graph__zTuple.__dask_graph__"   r#   r   c                    | j         S r   r    r!   s    r   __dask_keys__zTuple.__dask_keys__%   r#   r   c                    t           dfS )N tupler!   s    r   __dask_postcompute__zTuple.__dask_postcompute__(   s    byr   N)__name__
__module____qualname__staticmethoddaskthreadedget__dask_scheduler__r   r"   r%   r'   r,   r)   r   r   r   r      st        %dm&788            r   r   z#ignore:The dask.delayed:UserWarningc                    t          dd          } t          dd          }t          | |dg          \  }}|g dk    sJ t          | |df          \  }}|t          g dfk    sJ t          |          t	          | j        |j                  k    sJ 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    sJ t          |          t	          | j        |j                  k    sJ t          dg d	          } || |d          }t          |          \  }}||dddfk    sJ t          |          t	          | j        |j                  k    sJ t          t          | |d                    \  }}|t          dddfk    sJ t          |          t	          | j        |j                  k    sJ  G d
 dt                    }t           |                      \  }}t          |          |u sJ t          |          i k    sJ t          ddt          ddfd	g d	          }t          |          \  }}||v sJ  |j        |          }|t          g d	fk    sJ ||j        k    sJ d S )N   aname   b   )r7   r;   r<   fr7   r;   cc                      e Zd ZdS )"test_to_task_dask.<locals>.MyClassN)r-   r.   r/   r)   r   r   MyClassrA   F   s        r   rB   )r   r   r+   dictr   r1   r   slicetyper   r   popr   )r7   r;   taskr1   r=   xrB   s          r   test_to_task_daskrI   ,   s   AAq!Qi((JD$===    q!Qi((JD$E===)))))::qvqv......q!Ql++JD$DC8c1X.///4DC8cSTXBV;W3W3W3W3W::qvqv......3((A	!Q

AaJD$AsC#####::qvqv......eAq!nn--JD$E3Q'''''::qvqv......    $    ggii((JD$::    :: 	AAS#sO44oooFFAaJD$4<<<<A(((((17??????r   c                    t          t                    }  | dd                                          dk    sJ  | dd          dz                                   dk    sJ  |  | dd          d                                          dk    sJ t          d          }|                                dk    sJ d|j                                        v sJ  |  | |d          d          }|j        |j        v sJ d S )Nr6   r:   r<      )r   r   r   r1   valueskey)add2r7   r;   s      r   test_delayedrO   V   s   3<<D41::1$$$$DAJJN##%%****4Q

A&&((A----

A99;;!TT!QZZA5AF??????r   c                      G d dt                     } t          j        d          }t          j        d | |          i          }d } t          |          |          }|                                dk    sJ d S )Nc                      e Zd ZU ded<   dS )1test_delayed_with_namedtuple.<locals>.ANamedTupleintr7   Nr-   r.   r/   __annotations__r)   r   r   ANamedTuplerR   d   s         r   rV   r<   r7   r7   c                    | d         j         S )Nr7   rW   objs    r   return_nestedz3test_delayed_with_namedtuple.<locals>.return_nestedj   s    3xzr   )r   r1   r   r   )rV   literal
with_classr[   finals        r   test_delayed_with_namedtupler_   c   s        j    l1ooGsKK'$:$:$:;<<J   #GM"":..E==??ar   c                      e Zd ZU ded<   dS )ANonFrozenDataClassrS   r7   NrT   r)   r   r   ra   ra   r            
FFFFFr   ra   T)frozenc                      e Zd ZU ded<   dS )AFrozenDataClassrS   r7   NrT   r)   r   r   re   re   w   rb   r   re   clsc                    t          d          }t          d | |          i          }d } t          |          |          }|                                dk    sJ d S )Nr<   datarW   c                    | d         j         S Nrh   rW   rY   s    r   r[   z2test_delayed_with_dataclass.<locals>.return_nested       6{}r   r   r   )rf   r\   r]   r[   r^   s        r   test_delayed_with_dataclassrm   |   sp    ajjG&##...122J   #GM"":..E==??ar   c                 f   t                       G d d                      } t          j        d          }t          j        t
                    5 }t          j        d | |          i           d d d            n# 1 swxY w Y   |                    d           |                    d           d S )Nc                  "    e Zd ZU ded<   ddZdS )@test_delayed_with_dataclass_with_custom_init.<locals>.ADataClassrS   r7   r;   c                    || _         d S r   rW   r   r;   s     r   r   zItest_delayed_with_dataclass_with_custom_init.<locals>.ADataClass.__init__       DFFFr   Nr;   rS   r-   r.   r/   rU   r   r)   r   r   
ADataClassrp      3         	 	 	 	 	 	r   rv   r<   rh   r;   z custom __init__ is not supported)r   r1   r   pytestraises	TypeErrormatch)rv   r\   es      r   ,test_delayed_with_dataclass_with_custom_initr~      s    [[       [ l1ooG	y	!	! 6Qfjj733345556 6 6 6 6 6 6 6 6 6 6 6 6 6 6 GGMGG/00000s   !A<<B B c                     t                       G d d                      } t          d | d          i          }d } t          |          |          }|                                dk    sJ d S )Nc                  "    e Zd ZU ded<   ddZdS )Ftest_delayed_with_dataclass_with_eager_custom_init.<locals>.ADataClassrS   r7   r;   c                    || _         d S r   rW   rr   s     r   r   zOtest_delayed_with_dataclass_with_eager_custom_init.<locals>.ADataClass.__init__   rs   r   Nrt   ru   r)   r   r   rv   r      rw   r   rv   rh   r<   rx   c                    | d         j         S rj   rW   rY   s    r   r[   zItest_delayed_with_dataclass_with_eager_custom_init.<locals>.return_nested   rk   r   r   r   r   )rv   r]   r[   r^   s       r   2test_delayed_with_dataclass_with_eager_custom_initr      s    [[       [ &**q///233J   #GM"":..E==??ar   c                     t            G d d                      fd} t          d | d          i          }d } t          |          |          }|                                dk    sJ d S )Nc                  <    e Zd ZU ded<    ed          Zded<   dS )Otest_delayed_with_eager_dataclass_with_set_init_false_field.<locals>.ADataClassrS   r7   Finitr;   Nr-   r.   r/   rU   r   r;   r)   r   r   rv   r      9         E"""""""""r   rv   c                .     |           }d|_         |S NrW      rx   r7   rh   rv   s     r   prep_dataclasszStest_delayed_with_eager_dataclass_with_set_init_false_field.<locals>.prep_dataclass       zAr   rh   r<   c                    | d         j         S rj   rW   rY   s    r   r[   zRtest_delayed_with_eager_dataclass_with_set_init_false_field.<locals>.return_nested   rk   r   r   )r   r]   r[   r^   rv   s       @r   ;test_delayed_with_eager_dataclass_with_set_init_false_fieldr      s    # # # # # # # Y#    
 &.."3"3455J   #GM"":..E==??ar   c                 ^   t            G d d                      t          j        d          } fd}t          j        t
                    5 }t          j         ||                      d d d            n# 1 swxY w Y   |                    d           |                    d           d S )Nc                  <    e Zd ZU ded<    ed          Zded<   dS )Itest_delayed_with_dataclass_with_set_init_false_field.<locals>.ADataClassrS   r7   Fr   r;   Nr   r)   r   r   rv   r      r   r   rv   r<   c                .     |           }d|_         |S r   rx   r   s     r   r   zMtest_delayed_with_dataclass_with_set_init_false_field.<locals>.prep_dataclass   r   r   z`init=False` are not supported)r   r1   r   ry   rz   
ValueErrorr|   )r\   r   r}   rv   s      @r   5test_delayed_with_dataclass_with_set_init_false_fieldr      s   # # # # # # # Y# l1ooG    
 
z	"	" .a^^G,,---. . . . . . . . . . . . . . . GGMGG-.....s   A88A<?A<c                    t            G d d                      } t          j        d          }t          d | |          i          }d } t          |          |          }|                                dk    sJ d S )Nc                  <    e Zd ZU ded<    ed          Zded<   dS )Ktest_delayed_with_dataclass_with_unset_init_false_field.<locals>.ADataClassrS   r7   Fr   r;   Nr   r)   r   r   rv   r      r   r   rv   r<   rh   rW   c                    | d         j         S rj   rW   rY   s    r   r[   zNtest_delayed_with_dataclass_with_unset_init_false_field.<locals>.return_nested   rk   r   )r   r1   r   r   )rv   r\   r]   r[   r^   s        r   7test_delayed_with_dataclass_with_unset_init_false_fieldr      s    # # # # # # # Y# l1ooG&**w"7"7"7899J   #GM"":..E==??ar   c                    t          g d          } | d                                         dk    sJ | | z                                   g dk    sJ t          d          }| d |                                         ddgk    sJ t          d          } | dz                                   dk    sJ d| z                                   dk    sJ | dz	                                  dk    sJ | dk                                    sJ | dz                                  d	k    sJ  G d
 d          }t           |                      }t           |                      }||z                                  dk    sJ d S )Nr6   r:   r<   r   r6   )r6   r:   r<   r6   r:   r<   r:   
         d   c                      e Zd Zd ZdS )test_operators.<locals>.dummyc                    dS )Nr   r)   )r   others     r   
__matmul__z(test_operators.<locals>.dummy.__matmul__       1r   N)r-   r.   r/   r   r)   r   r   dummyr      s#        	 	 	 	 	r   r   r   rl   )r7   r;   r   r?   ds        r   test_operatorsr      s   			AQ4<<>>QE?? 2 2 22222

ARaR5==??q!f$$$$AE??""""E??""""F""""E??qD>>s""""        	AAE??!!!!!!r   c                    t          d          } |                     d                                          g dk    sJ |                                                     dd                                                              d                                          dk    sJ |                     dd          j        |                     dd          j        k    sJ |                     dd	
          }|j        d	k    sJ d S )Nz	a b c d e )r7   r;   r?   r   r}   BAr:   Tpuretestdask_key_name)r   splitr   upperreplacecountrM   r7   os     r   test_methodsr      s    A773<<!!%>%>%>>>>>7799S#&&,,..44S99AACCqHHHH773T7""&!''#D'*A*A*EEEEE	6**A5F??????r   c                 &   t          d          } | j        j        | j        j        k    sJ | j                                        dk    sJ | j                                        dk    sJ | j        | j        z                                   dk    sJ d S )N       @      ?r:   r6   r<   )r   real_keyr   imagrW   s    r   test_attributesr     s    A6;!&+%%%%6>>q    6>>q    FQVO$$&&!++++++r   c                     t          g d          } |                     d          }t          d |                                                                D             vsJ d S )Nr   r6   c                    h | ]
}|d          S )r   r)   ).0vs     r   	<setcomp>z5test_method_getattr_call_same_task.<locals>.<setcomp>  s    EEEA1Q4EEEr   )r   indexgetattrr%   rL   r   s     r   "test_method_getattr_call_same_taskr     s`    			A	

AEE)9)9););)B)B)D)DEEEEEEEEEr   c                    t          j        d          } t          d          }t          j        t                    5  |                     |           d d d            n# 1 swxY w Y   t          |                     dgd                    j                                        |                     d          k    sJ d S )Nnumpyr6   f8dtype)ry   importorskipr   rz   r{   r   arrayr   )nprH   s     r   test_np_dtype_of_delayedr     s     
	W	%	%B

A	y	!	!  
              288QCt8,,--3;;==$OOOOOOs   AA#&A#c                    d }  t          j        |           d          }|                                 t          j        t
          d          5  |                    d           d d d            n# 1 swxY w Y   t          j        t
          d          5  |                                 d d d            d S # 1 swxY w Y   d S )Nc                    | dz   S Nr6   r)   rH   s    r   r   z(test_delayed_visualise_warn.<locals>.inc"      1ur   r6   z/dask.delayed objects have no `visualise` methodr|   zdesk_graph.svg)	file_name)r1   r   r   ry   warnsUserWarning	visualise)r   zs     r   test_delayed_visualise_warnr     sT      	S!AIIKKK	L
 
 
 0 0 	
.///0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 
L
 
 
   	
                 s$   A33A7:A7B;;B?B?c                     t          g d           t          j        t           fd           t          j        t           fd           t          j        t           fd           t          j        t           fd           t          j        t           fd           t          j        t           fd           d S )Nr   c                 &    t           dd          S )Nfoor6   )setattrrW   s   r   <lambda>z%test_delayed_errors.<locals>.<lambda>7  s    WQq%9%9 r   c                 &    t           dd          S )Nr6   r   )r	   rW   s   r   r   z%test_delayed_errors.<locals>.<lambda>8  s    WQ1%5%5 r   c                     d v S r   r)   rW   s   r   r   z%test_delayed_errors.<locals>.<lambda>:  s    Q!V r   c                 "    t                     S r   listrW   s   r   r   z%test_delayed_errors.<locals>.<lambda>;      T!WW r   c                 ,                                      S r   )_hiddenrW   s   r   r   z%test_delayed_errors.<locals>.<lambda>=  s    !))++ r   c                 "    t                     S r   )boolrW   s   r   r   z%test_delayed_errors.<locals>.<lambda>?  r   r   )r   ry   rz   r{   AttributeErrorrW   s   @r   test_delayed_errorsr   4  s    			A
M)9999:::
M)5555666
M)^^^^,,,
M)____---
M."5"5"5"5666
M)____-----r   c                     t          g d          } | d         | d         z   }| d         j        |j        v sJ | j        |j        v sJ t          |j                  dk    sJ d S )Nr   r   r<   )r   rM   r1   len)r7   ress     r   test_common_subexpressionsr   B  sq    			A
A$1+CQ48sx5CHsx==Ar   c                     t          ddt          dft          dfd          } t          j        |           \  }t	          |j                                                  ddgk    sJ |j        |j        k    sJ d S )Nr;   r6   r7   r>   )r   r   r1   optimizesortedr   _layerr   )rH   x2s     r   test_delayed_optimizer   J  sv    1C:S#J??@@AM!ER"',,..!!c3Z////9r   c                     t          d          } t          d          } t          t                    | |g          }|                                dk    sJ d S )Nr6   r:   r<   )r   sumr   r>   s      r   
test_listsr   T  sP    

A

AaVA99;;!r   c                    t          d          } | dz   }| |df}t          |                                          dk    sJ | |dg}t          |                                          g dk    sJ | |dh}t          |                                          h dk    sJ | d|dddi}t          |                                          ddddk    sJ t          |          |                                          dk    sJ | |dd}t          |                                          dd	ddk    sJ t          |          d                                         dk    sJ d S )
Nr6   r<   r   >   r6   r:   r<   r7   r;   r?   r>   r:   rl   r7   r;   lits      r   test_literatesr  [  sx   

A	AAa)C3<<!!Y....a)C3<<!!YYY....a)C3<<!!YYY....c1c1c
"C3<<!!%=%=====3<<?""$$++++
"
"C3<<!!111%=%=====3<<$$&&!++++++r   c                     t          d          } | dz   }| |df}t          |          j        t          |          j        k    sJ t          |d          j        t          |d          j        k    sJ d S )Nr6   r<   Tr   )r   rM   r  s      r   test_literates_keysr  l  s{    

A	AAa)C3<<ws||/////3T"""&'#D*A*A*A*EEEEEEEr   c                     t          d          } t          d          } t          t                    | dg|dggd           d         }|                                dk    sJ d S )Nr6   r:   r      c                    | d         S )Nr   r)   r   s    r   r   z)test_lists_are_concrete.<locals>.<lambda>w  s
    qt r   )rM   )r   maxr   r>   s      r   test_lists_are_concreter
  t  se    

A

Aq"g2w'^^<<<Q?A99;;"r   c                 V   t          d          } t          d          } t          t                    t          | |g                    }|                                dk    sJ d } t          |          t          | |g                    }|                                dk    sJ d S )Nr6   r:   r<   c                     t          |           S r   r   )seqs    r   r=   ztest_iterators.<locals>.f  s    3xxr   )r   r   iterr   )r7   r;   r?   r=   s       r   test_iteratorsr  |  s    

A

AT1a&\\""A99;;!   	

4A<<  A99;;!r   c                 &   d }  t          |                       }|dddg}t          |d                                          }t          |          dk    sJ |d         |u sJ |dd          |dd          k    sJ | || |ff}t          |d                                          }t          |t                    sJ |d         | k    sJ |d         |u sJ d| |f|g}t          |d                                          }t          |t
                    sJ |d         dk    sJ |d         d         | k    r|d         d         |u sJ |d         |u sJ t          d          }t          |d          }|                                dk    sJ d S )	Nc                      t          d          )Nzshouldn't have computed)r   )argss    r   failz!test_traverse_false.<locals>.fail  s    2333r   r6   r:   r<   F)traverser   r   )r   r   r   
isinstancer+   r   )r  r7   rH   r   r;   s        r   test_traverse_falser    s   4 4 4 	A 
Aq!A
!e
$
$
$
,
,
.
.Cs88q====q6Q;;;;qrr7ae 
q4)A
!e
$
$
$
,
,
.
.Cc5!!!!!q6T>>>>q6Q;;;; 
T1IqA
!e
$
$
$
,
,
.
.Cc4     q6Q;;;;q6!9Qaq6Q;;;; 	

AE"""A99;;!r   c                     t          t          d          dd          }  t          t          d          dd          }| j        |j        k    sJ t          t                    } |            j         |            j        k    sJ d S NTr   r6   r:   )r   r   rM   r
   )v1v2myrands      r   	test_purer    s    	 4	 	 	 A	&	&B	 4	 	 	 A	&	&B6RVV__F688<6688<''''''r   c                 J   t          t                    } t          j                            d          5   | dd          j         | dd          j        k    sJ 	 d d d            n# 1 swxY w Y   t          j                            d          5   | dd          j         | dd          j        k    sJ 	 d d d            n# 1 swxY w Y   t          t          d          } t          j                            d          5   | dd          j         | dd          j        k    sJ 	 d d d            n# 1 swxY w Y   t          d          j        t          d          j        k    sJ t          j                            d          5  t          d          j        t          d          j        k    sJ 	 d d d            n# 1 swxY w Y   t          j                            d          5  t          dd          j        t          dd          j        k    sJ 	 d d d            n# 1 swxY w Y   t          g d          }|                    d          j        |                    d          j        k    sJ t          j                            d          5  |                    d          j        |                    d          j        k    sJ |                    dd          j        |                    dd          j        k    sJ 	 d d d            n# 1 swxY w Y   t          j                            d          5  |                    dd          j        |                    dd          j        k    sJ 	 d d d            n# 1 swxY w Y   t          j                            d          5  |j        j        |j        j        k    sJ |d         }||z   j        ||z   j        k    sJ 	 d d d            d S # 1 swxY w Y   d S )	NT)delayed_purer6   r:   Fr   r   r   )r   r   r1   configsetrM   r   )funcrh   elements      r   test_pure_global_settingr$    s   3<<D	d	+	+ 0 0tAqzz~a/////0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 
e	,	, 0 0tAqzz~a/////0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 3T"""D	e	,	, 0 0tAqzz~a/////0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 1::>WQZZ^++++	d	+	+ 0 0qzz~/////0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 
e	,	, F Fqt$$$(GAD,A,A,A,EEEEEEF F F F F F F F F F F F F F F 999D::a==

1 11111	d	+	+ N Nzz!}} DJJqMM$55555zz!%z((,

15
0I0I0MMMMMMN N N N N N N N N N N N N N N 
e	,	, L Lzz!$z''+tzz!$z/G/G/KKKKKKL L L L L L L L L L L L L L L 
e	,	, B Bz~////q''!&7W+<*AAAAAAB B B B B B B B B B B B B B B B B Bs   'A))A-0A-'CCC	'D==EE-GGG91H77H;>H;+A5L--L14L1=N""N&)N&=PPPc                    t          d dd          }  | d          }t          |          dk    sJ |\  }t          |          dk    sJ j        J |j        J t	          j        t          fd           t	          j        t          fd           t	          j        t          d	            t	          j        t          d
            t          t          d           }  | d          j        J t	          j        t          fd           t	          j        t          fd           t          d dd          }  | d          }t          |          dk    sJ |\                                  dk    sJ j        J t	          j        t          fd           t          d dd          }  | d          }t          |          dk    sJ |                                t                      k    sJ d S )Nc                    | |  fS r   r)   r   s    r   r   ztest_nout.<locals>.<lambda>  s    a!W r   r:   Tnoutr   r6   )r6   c                 "    t                     S r   r   rW   s   r   r   ztest_nout.<locals>.<lambda>      SVV r   c                 "    t                     S r   r   rW   s   r   r   ztest_nout.<locals>.<lambda>  r   r   c                 .    t          t          d          S )Nr)  r(  r   r   r)   r   r   r   ztest_nout.<locals>.<lambda>  s    gc&;&;&; r   c                 .    t          t          d          S )NTr/  r0  r)   r   r   r   ztest_nout.<locals>.<lambda>  s    gc&=&=&= r   r/  c                 "    t                     S r   r   rW   s   r   r   ztest_nout.<locals>.<lambda>  r   r   c                 "    t                     S r   r+  rW   s   r   r   ztest_nout.<locals>.<lambda>  r,  r   c                    | fS r   r)   r   s    r   r   ztest_nout.<locals>.<lambda>  s    aT r   c                 "    t                     S r   r+  rW   s   r   r   ztest_nout.<locals>.<lambda>  r,  r   c                    t                      S r   r*   r   s    r   r   ztest_nout.<locals>.<lambda>  s
    UWW r   r   )
r   r   r   _lengthry   rz   r{   r   r   r+   )r"  rH   r;   r7   s      @r   	test_noutr8    s	   $$14888DQAq66Q;;;;DAq1a==G####99
M)^^^^,,,
M)____---
M*;;<<<
M*==>>>3T"""DQA9
M)____---
M)^^^^,,,>>555DQAq66Q;;;;DQ99;;!9
M)^^^^,,,$$14888DQAq66Q;;;;99;;%''!!!!!!r   rH   r6   r:   )r6   r:   r)   c                    t          |           }t          | |          }t          |          t          t          |                    cxk    r|k    sn J |                                | k    sJ d S )Nr/  )r   r   r   r   )rH   lengthr   s      r   test_nout_with_tasksr;     sw    
 VVFAq66Sa\\++++V++++++99;;!r   c            	        dd} t          |           } |ddt          d          dg |dd                    }|                                dk    sJ t          | d	
          }t          d          dg} |dd| |dd                    }|                                dk    sJ  |dd| |dd                    j        |j        k    sJ  |dd| |dd                    j        |j        k    sJ  |dd|d          j        |j        k    sJ  |dd|d          j         |dd|d          j        k    sJ d S )Nr)   c                p    | |z   t          |          z   t          |                                          z   S r   )r   rL   )r7   r;   r?   kwargss       r   mysumztest_kwargs.<locals>.mysum  s+    1us1vv~FMMOO 4 444r   r6   r:   r<   r   )r?   fourr   Tr   r   )r)   )r   r   rM   )r?  dmysumtenr?   s       r   test_kwargsrC    s   5 5 5 5 U^^F
&A'!**avva||
<
<
<C;;==BU&&&F	QA
&A1
.
.
.C;;==B6!Q!&&A,,///3sw>>>>6!Q!&&A,,///3sw>>>>6!Q!!$$$(CG33336!Q!!$$$(FF1a11,E,E,E,IIIIIIIr   c                    t          ddt          ddfdg d          }  t          t          d          | d          } t          t          d          |           } t          t          d          |           j        |j        k    sJ |                                d	k    sJ t          |||           d
k    sJ d S )Nr6   r:   r7   r;   r>   Tr   )r   r   rK   r6   r:   r<   r   r   rK   )r<   rE  r   )r   r   r   r   rM   r   )rH   r   ns      r   test_custom_delayedrG    s    AAS#sO44oooFFA	 4	 	 	 I	.	.B$""A"73T"""1%%)QU2222::<<-----1b! BBBBBBBr   c                    t          j        d          } t          j        d          }|                     d                              d          }|                    |d          } t          t                    ||dg          }t          |t                    sJ | 	                    |
                                ||z   dz             sJ |                                
                                ||z   dz                                   k    sJ |d         
                                ||z   dz   d         k    sJ t          |          \  }}|j                                        |                                z
  rJ |                                |j                                        z
  }t          |          dk    sJ t          |          }|
                                |k                                    sJ d S )	Nr   
dask.arrayr   r   r   r   r   chunksr6   )r   r   )ry   r   arangereshape
from_arrayr   r   r  r   allcloser   r   r1   r   r   all)	r   daarrdarrvalrG   r   diffdelayed_arrs	            r   test_array_delayedrY  %  s   		W	%	%B		\	*	*B
))C..
 
 
*
*C==V=,,D
'#,,T1~
&
&Cc7#####;;s{{}}cCi!m44444779939q="5"5"7"77777t939q=$"77777T""ID#y~~#((**,,,,88::	(((Dt99>>>>$--K!!S(--///////r   c                 p   t          j        d          } t          j        d          }|                    d                              d          }|                    |j                  }|                     |d          }|                     |d          }t          j        g d          }|||||g} t          t                    d |D                       }|                                d	|
                                z  d	|
                                z  z   t          g d          z   k    sJ d S )
NrI  r   r   rJ  rK  rL  r   c                6    g | ]}|                                 S r)   r  )r   is     r   
<listcomp>z*test_array_bag_delayed.<locals>.<listcomp>E  s     ---A---r   r:   )ry   r   rN  rO  dotTrP  dbfrom_sequencer   r   r   )	rS  r   arr1arr2darr1darr2r;   r  outs	            r   test_array_bag_delayedrg  ;  s   		\	*	*B		W	%	%B99S>>!!(++D88DFDMM$vM..EMM$vM..E
##AueQ
'C
'#,,-----
.
.C;;==A

NQ^;c)))nnLLLLLLLr   c                     t          t          dd          dd          } t          j        t          j        |                     }| j        |j        k    sJ | j        |j        k    sJ | j        |j        k    sJ t          d          } t          j        t          j        |                     }| j        |j        k    sJ | j        |j        k    sJ | j        |j        k    sJ | j	        |j	        k    sJ | j
        } t          j        t          j        |                     }| j        j        |j        j        k    sJ | j        j        |j        j        k    sJ | j        |j        k    sJ | j        |j        k    sJ d S )Nr:   Tr'  r6   r   )r   divmodpickleloadsdumpsr1   r   r7  _nout_purer   _obj_attrrH   ys     r   test_delayed_picklablers  I  sn   *QT***1a00AV\!__%%A6QV6QV9	!!!!AV\!__%%A6QV6QV7ag7ag	AV\!__%%A6;!&+%%%%6;!&+%%%%7ag6QVr   c                 V    t          d          dz   } |                     d           d S )Nr6   r:   r   )bogus_keywordrl   r   s    r   #test_delayed_compute_forward_kwargsrv  `  s+    

QAIIBIr   c                 J     t          t          j                  d           d S )Nr   )r   bytesdecoder)   r   r   test_delayed_method_descriptorrz  e  s"    GEL#r   c                     t          t          d          }  | dd          }|j        |j        t          ddfik    sJ | j        | j        t          ik    sJ |                                 t          k    sJ d S r  )r   r   r1   rM   r   )r=   r   s     r   test_delayed_callabler|  i  s{    $A	!QA6aec1a[)))))6aeS\!!!!99;;#r   c                 d    t          t          d          }  | ddd          j        dk    sJ d S )NTr   r6   r:   r   r   )r   r   r   )r=   s    r   test_delayed_name_on_callr~  r  s?    $A1Q''',555555r   c                     G d d          }  | d          }t          |          }|                                |u sJ |j                                        dk    sJ  |                                            dk    sJ d S )Nc                      e Zd Zd Zd ZdS )test_callable_obj.<locals>.Fooc                    || _         d S r   rW   )r   r7   s     r   r   z'test_callable_obj.<locals>.Foo.__init__y  rs   r   c                    dS )Nr:   r)   r!   s    r   __call__z'test_callable_obj.<locals>.Foo.__call__|  r   r   N)r-   r.   r/   r   r  r)   r   r   Foor  x  s2        	 	 		 	 	 	 	r   r  r6   r:   )r   r   r7   )r  r   r=   s      r   test_callable_objr  w  s            #a&&CA99;;#3;;==A133;;==Ar   c                    | S r   r)   r   s    r   identityr    s    Hr   c                     t          t          d          } ddg dd} | |          j        dk    sJ ddd	} | |          j         | |          j        k    sJ  | d          j        d
k    sJ d S )NTr   r6      r   )rH   rr  r   z)identity-4f318f3c27b869239e97c3ac07f7201arH   )rH   r6   z)identity-7258833899272585e16d0ec36b21a3de)r   r  r   )r"  rh   s     r   %test_name_consistent_across_instancesr    s    8$'''D))),,D4::?IIIIIsD4::?dd4jjo----477<FFFFFFFr   c                      t          t          t          d          d          d          j         t          t          t          d          d          d          j        k    sJ d S )Nr   Tr   r:   r  )r   r   r   r   r)   r   r   test_sensitive_to_partialsr    sn    ,R  t,,,Q//4/773##$///227	8 	8 	8 	8 	8 	8r   c                 P   t          d          j                            d          sJ t          dd          j                            d          sJ t          dd          j        dk    sJ d }  t          |           d          j                            d          sJ d S )	Nr6   zint-Tr   Xr8   c                    | dz   S r   r)   r   s    r   myfuncz!test_delayed_name.<locals>.myfunc  r   r   r  )r   r   
startswithrM   )r  s    r   test_delayed_namer    s    1::?%%f-----14   %008888813$++++   76??1!,,X6666666r   c                 $   t          j        d          } |                     dd          }t          |g          }t	          |j                                      |j                  sJ d t          fd|j        D                       sJ d S )NrI  r   r   rL  c                    t          | t                    r| d         } |                     d          d                             dd          S )Nr   -_ )r  r+   r   r   )ss    r   rM   ztest_finalize_name.<locals>.key  sA    a 	!Awws||A&&sB///r   c              3  R   K   | ]!} |                                           V  "d S r   )isalpha)r   krM   s     r   	<genexpr>z%test_finalize_name.<locals>.<genexpr>  s5      00Ass1vv~~000000r   )ry   r   onesr   r!  r1   issubsetrR  )rS  rH   r   rM   s      @r   test_finalize_namer    s    		\	*	*B
1AAqv;;'''''0 0 0 00000000000000r   c                     t          j        d          } ddlm} |                     dd                                                                          }d |D             } ||d         j                   d S )NrI  r   )
_check_dskrJ  r   rL  c                H    g | ]} t          t                    |           S r)   )r   r   )r   rH   s     r   r]  z(test_keys_from_array.<locals>.<listcomp>  s&    	%	%	%a,'#,,q//	%	%	%r   )ry   r   dask.array.utilsr  r  
to_delayedflattenr1   )rS  r  r  xss       r   test_keys_from_arrayr    s    		\	*	*B++++++
##..0088::A	%	%1	%	%	%BJr!uzr   c                     G d d          }  | d          }|                     dd                                          dk    sJ |                      |dd                                          dk    sJ |                    dd                                          dk    sJ |                     dd                                          dk    sJ |                    dd                                          dk    sJ |                     dd                                          dk    sJ t	          |j         t
          j                  sJ t	          | j        t
          j                  sJ t	          | j        t                    sJ d S )	Nc                  z    e Zd ZdZd Zed             Zeed                         Ze	ed                         Z
dS )+test_delayed_decorator_on_method.<locals>.Ar   c                    || _         d S r   BASE)r   bases     r   r   z4test_delayed_decorator_on_method.<locals>.A.__init__  s    DIIIr   c                    | j         |z   |z   S r   r  )r   rH   rr  s      r   	addmethodz5test_delayed_decorator_on_method.<locals>.A.addmethod  s    9q=1$$r   c                    | j         |z   |z   S r   r  )rf   rH   rr  s      r   addclassz4test_delayed_decorator_on_method.<locals>.A.addclass  s     8a<!##r   c                    | |z   S r   r)   rq  s     r   	addstaticz5test_delayed_decorator_on_method.<locals>.A.addstatic  s     q5Lr   N)r-   r.   r/   r  r   r   r  classmethodr  r0   r  r)   r   r   r   r    s        	 	 	 
	% 	% 
	% 
		$ 	$ 
 
	$ 
		 	 
 
	 	 	r   r   r   r<   r   k         )r  r   r  r  r  types
MethodTyper   )r   r7   s     r    test_delayed_decorator_on_methodr    s          ( 	
#A;;q!$$&&#----;;q!Q''))S0000::a##%%++++::a##%%++++;;q!$$&&!++++;;q!$$&&!++++ ak5#344444aj%"233333 ak7+++++++r   c                     t          d          } t          | j        t                    sJ t          | j        j        t                    sJ t          | j        j        j        t                    sJ d S )N{   )r   r  r7   r   r;   r?   r   s    r   test_attribute_of_attributer    sa    Aac7#####aceW%%%%%acegw'''''''r   c                 B   t          j        d          } ddlm}  |g dd          } |g dd          } t	          d           |          } t	          d	           |          }|                     ||gd
          }| j                            ||           d S )Nzdask.dataframer   )Series)r7   r;   r7   categoryr   )r7   r?   r7   c                    | S r   r)   r   s    r   r   z&test_check_meta_flag.<locals>.<lambda>      1 r   c                    | S r   r)   r   s    r   r   z&test_check_meta_flag.<locals>.<lambda>  r  r   F)verify_meta)ry   r   pandasr  r   from_delayedutils	assert_eq)ddr  r7   r;   rS  r`  r?   s          r   test_check_meta_flagr    s    		-	.	.Bj111Aj111A			a	 	 B			a	 	 B
Re44AHq!r   c                    | dz   S r   r)   r   s    r   modlevel_eagerr    s    q5Lr   c                    | dz   S r   r)   r   s    r   modlevel_delayed1r        q5Lr   Fr   c                    | dz   S r   r)   r   s    r   modlevel_delayed2r  
  r  r   r=   z#3369)reason)marksc                     | d          }t          j        t          j        |t           j                            }|                                dk    sJ d S Nr:   )protocolr<   )rj  rk  rl  HIGHEST_PROTOCOLr   r=   r   s     r   test_pickler    sQ     	
!AV\!f.EFFFGGA99;;!r   c                     | d          }t          j        t          j        |t          j                            }|                                dk    sJ d S r  )cloudpicklerk  rl  rj  r  r   r  s     r   test_cloudpickler    sS     	
!A++A8OPPPQQA99;;!r   c                 r   t          d          } | j        j                                        | j        hk    sJ | j        j        | j        t                      ik    sJ |                                 | j        fk    sJ t          |           }|j        j                                        | j        |j        hk    sJ |j        j        | j        t                      |j        | j        hik    sJ |                                |j        fk    sJ t          j
        dd|j        i|g          }t          j        t          d          5  t          d|           d d d            n# 1 swxY w Y   t          d|d          }|                                dk    sJ |j                                         d S )	Nr6   r   alias)dependenciesznot inr   )layerr   )r   r1   layersr   rM   r  r!  __dask_layers__r  r   from_collectionsry   rz   r   r   validate)d1d2hlgexplicits       r   test_dask_layersr  &  s   	B7>  RVH,,,,7BFCEE?2222BF9,,,,	2		B7>  RVRV$444447BFCEE26BF8#DDDDDBF9,,,,

)%'261BRTQU
V
V
VC	z	2	2	2                 w5111H##%%1111Ms   
E''E+.E+c                    t          j        d          5  t          j        ddt          dft          dfdg           } t          d|           }d d d            n# 1 swxY w Y   t          |j                   t          u sJ t          |j         j                  dk    sJ t          |j         j        d                   dk    sJ |j         j        d         j	        ddik    sJ t          j
        |          \  }t          |j                   t          u sJ t          |j         j                  dk    sJ t          |j         j        d                   d	k    sJ |j         j        d         j	        ddik    sJ d S )
Nbarr  r;   r6   r7   r>   r<   r   r:   )r1   annotater   r  r   r   rE   r   r  r   r   )graphr   d_opts      r   %test_annotations_survive_optimizationr  9  s   	5	!	!	!    /3*C:66
 

 C                              <<>))))qv}""""qv}S!""a''''6=)eU^;;;; }QHU
~----uz !!Q&&&&uz %&&!++++:S!-%??????s   9AA"Ac                 Z    t           d             } | j        dk    sJ | j        dk    sJ d S )Nc                    | |z   S )This is a docstringr)   rq  s     r   r   z7test_delayed_function_attributes_forwarded.<locals>.addQ  s     1ur   r   r  )r   r-   __doc__)r   s    r   *test_delayed_function_attributes_forwardedr  P  sJ      W <5    ;///////r   )a
__future__r   rj  r  collectionsr   dataclassesr   r   	functoolsr   operatorr   r	   r
   typingr   r  ry   tlzr   r1   dask.bagbagr`  r   dask.delayedr   r   r   dask.highlevelgraphr   dask.utils_testr   r   markfilterwarningsrI   rO   r_   ra   re   parametrizerm   r~   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r  r  r
  r  r  r  r$  r8  r;  rC  rG  rY  rg  rs  rv  rz  r|  r~  r  r  r  r  r  r  r  r  r  r  r  r  r  paramxfailr  r  r  r  r  r)   r   r   <module>r	     s   " " " " " "   " " " " " " ( ( ( ( ( ( ( (       ! ! ! ! ! ! ! !                                     7 7 7 7 7 7 7 7 7 7 . . . . . .             ( ABB& & CB&R
 
 
              $        !46F GHH	  	  IH	 1 1 1"     $     */ / /(     "" " ".  , , ,F F FP P P  *. . .         , , ,"F F F         F( ( (%B %B %BP" " "D VVc1a["b)  	 J J J"C C C ABB0 0 CB0*M M M  .     
    6 6 6
    G G G  7 7 71 1 1   (, (, (,V( ( (
 
 
   	  	 	e   &fk.?.?w.?.O.OPPP&fk.?.?w.?.O.OPPP    ''.
!
!#46G	H     &@ @ @.0 0 0 0 0r   