
    bi                       d dl mZ d dlZd dlZd dlZd dlmZmZmZm	Z	 d dl
mZ  G d de      Z G d de      Zd	 Zd
 Z G d de      Zej"                  j%                         d        Z G d de      Z G d de      Z G d de      Z G d d      Z G d dee      Zd Zej"                  j%                         d        Z G d de      Zd Zy)    )annotationsN)ExprProhibitReuseSingletonExpr_ExprSequenceDataNodec                      e Zd ZddgZy)MyExprfoobarN)__name__
__module____qualname___parameters     O/home/cdr/jupyterlab/.venv/lib/python3.12/site-packages/dask/tests/test_expr.pyr   r      s    %.Kr   r   c                      e Zd Zy)MyExpr2Nr   r   r   r   r   r   r   r      s    r   r   c                     t        dd      } d| _        | j                  dk(  sJ t        j                  t              5  d| _        d d d        y # 1 sw Y   y xY wN      r   r         )r   r   pytestraisesAttributeErrorbazes    r   test_setattrr%      sM    1!AAE55A::	~	&      AAc                     t        dd      } d| _        | j                  dk(  sJ t        j                  t              5  d| _        d d d        y # 1 sw Y   y xY wr   )r   r   r   r    r!   r"   r#   s    r   test_setattr2r(      sM    A1AAE55A::	~	&   r&   c                  L    e Zd ZdZddgZed        Zej                  d        Zy)MyExprCachedPropertyFr   r   c                4    | j                   | j                  z   S Nr   selfs    r   r"   zMyExprCachedProperty.baz)   s    xx$((""r   c                    t         j                  rt        d      dt         _        | j                  | j                  z   S )NzNo!T)r*   called_cached_propertyRuntimeErrorr   r   r-   s    r   cached_propertyz$MyExprCachedProperty.cached_property-   s2    66u%%6:3xx$((""r   N)	r   r   r   r0   r   propertyr"   	functoolsr2   r   r   r   r*   r*   %   s=    "%.K# # # #r   r*   c                 x    t        j                  d       ddlm}  ddlm}  |d| dg      d        } |        y )	Ndistributedr   )Nanny)gen_clusterT) r   )clientWorkernthreadsc                  K   t        dd      }t        d      D ]$  }|j                  dk(  sJ |j                  dk(  r$J  t         j                  du sJ t        j                  t        j                  |            }|j                  dk(  sJ t         j                  du sJ t        dd      }t        j                  t              5  |j                   d d d        d }| j                  ||       d {    y # 1 sw Y   'xY w7 w)Nr   r   r   
   r   Tc                v    t         j                  du sJ | j                  dk(  sJ t         j                  du sJ y )NFr   )r*   r0   r2   )exprs    r   fz6test_pickle_cached_properties.<locals>.test.<locals>.fN   s@     (>>%GGG''1,,,'>>%GGGr   )r*   ranger"   r2   r0   pickleloadsdumpsr   r    r1   submit)csar@   _rtexpr3rA   s           r   testz+test_pickle_cached_properties.<locals>.test;   s      $q1r 	-A88q= =''1,,,	- $::dBBB\\&,,t,-!!Q&&&#::dBBB %2]]<( 	"!!	"	H hhq$	" 	" 	 s/   <D
B
D
	C< D
6D7D
<DD
)r   importorskipr6   r7   distributed.utils_testr8   )r7   r8   rM   s      r   test_pickle_cached_propertiesrP   5   s;    
&!2UgY?  @ 8 	Fr   c                      e Zd Zy)MySingletonNr   r   r   r   rR   rR   [       r   rR   c                      e Zd Zd Zy)MySingletonWithCustomInitc                     y r,   r   r.   argskwargss      r   __init__z"MySingletonWithCustomInit.__init___       r   Nr   r   r   rZ   r   r   r   rU   rU   ^       ,r   rU   c                      e Zd Zy)MySingletonInheritsCustomInitNr   r   r   r   r_   r_   b   rS   r   r_   c                      e Zd Zd Zy)Mixinc                     y r,   r   rW   s      r   rZ   zMixin.__init__f   r[   r   Nr\   r   r   r   ra   ra   e   r]   r   ra   c                      e Zd Zy)$MySingletonInheritsCustomInitAsMixinNr   r   r   r   rd   rd   i   rS   r   rd   c                     t        dd      t        dd      u sJ t        dd      t        dd      usJ t        dd      t        dd      usJ t        dd      t        dd      usJ y )Nr   r   )rR   rU   r_   rd   r   r   r   test_singleton_exprrf   l   s    q!Aq 1111 %Q*2KAq2QQQQ(A.6S	17    0	11!Q78 8 8r   c                    t        j                  d      t        j                  d      } t        j                  d      }|j                  d|j                  d      5 }fd}t	        d      D cg c]  }|j                  ||       }} |d	      d d	 }| j                  ||      }|j                          ~|j                          d d d        y c c}w # 1 sw Y   y xY w)
Npandaszdask.dataframer6   r   z:0)	n_workersworker_classdashboard_addressc                2    j                  d| gi| g      S )NA)index)	DataFrame)ipds    r   genz%test_refcounting_futures.<locals>.gen   s    <<qc
1#<66r   r   r   )r   rN   Clientr;   rB   rF   from_delayedcompute)	ddr6   r:   rr   rp   futuresmetadfrq   s	           @r   test_refcounting_futuresrz   y   s    			X	&B			-	.B%%m4K 
		+"4"4 
 
 		7 38(;Q6==a(;;1vbqz__Wd+





  < s   C2CA CCC"c                      e Zd ZddZy)FooExprc                    dt        dd      iS )Nr   *   r   r-   s    r   _layerzFooExpr._layer   s    xr*++r   N)returndict)r   r   r   r   r   r   r   r|   r|      s    ,r   r|   c                 .   t               } t        j                  j                  t                	 t	        | t        |             j                         j                         }t        |      dk(  sJ  |j                  d             }|j                         \  }}|j                  d      r|dk7  sJ |j                  rJ  |       |u sJ 	 t        j                  j                  t                y # t        j                  j                  t                w xY w)Nr   r   )r|   r   _ALLOWED_TYPESappendr   optimize__dask_graph__lenpoppopitem
startswithdependenciesremove)oncedskfirstkeyvals        r   test_prohibit_reuser      s    9D  ''05D-"56??APPR3x1}} ;;=S~~e$55####u~~$$++G4$$++G4s   BC/ /%D)
__future__r   r4   rC   r   
dask._exprr   r   r   r   dask._task_specr	   r   r   r%   r(   r*   markslowrP   rR   rU   r_   ra   rd   rf   rz   r|   r   r   r   r   <module>r      s    "    H H $!T !	f 	
#4 #  " "J &- %- - D$= C- - D; C
8  2,d ,
5r   