
    bi$                       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 d dlmZmZ 	 d dlZej$                  j'                  edu d	      Zd
 Zd Zd Zd Zd Zd Zed        Zed        Zed        Zed        Zd Zeej$                  jA                  dg d      ej$                  jA                  de!dg      d                      Z"eej$                  jA                  de!dg      d               Z#ej$                  jA                  dg d      d        Z$ed        Z%ed        Z&d Z'd Z(d  Z)d! Z*d" Z+y# e$ r dZY w xY w)#    )annotationsN)array)Version)get_context)sizeof)funcnametmpdirzrequires pandas)reasonc                 J    t        d      t        j                  d      k(  sJ y )N   r   sys	getsizeof     Q/home/cdr/jupyterlab/.venv/lib/python3.12/site-packages/dask/tests/test_sizeof.py	test_baser      s    !9a((((r   c                 ,    t        t              dk(  sJ y )Nr   )r   r   r   r   r   	test_namer      s    Fx'''r   c                     t        dddgg      t        j                  d      dz  t        j                  g       z   kD  sJ y )Nr         r   r   r   r   test_containersr      s9    1a!+#--"2Q"6r9J"JKKKr   c            	     :   dt        t        d            cxk  rdk  sJ  J dt        t        d            cxk  rdk  sJ  J dt        t        t        d                  cxk  rdk  sJ  J dt        t	        dt        d                  cxk  rdk  sJ  J y )N  i  @  d(#  )r   bytes	bytearray
memoryviewr   ranger   r   r   test_bytes_liker#   "   s    6%+&.$.....6)D/*2d222226*U4[12:d:::::6%U4[12:d:::::r   c                     t        j                  d      } dt        | j                  dd            cxk  rdk  sJ  J | j	                  d      }t        |      t        j                  |      k(  sJ y )Nnumpyr   r   f8dtyper   )pytestimportorskipr   emptyr(   r   r   )npdts     r   
test_numpyr.   )   se    			W	%B6"((4t(45======	$B":r****r   c                 r    t        j                  d      } | j                  dd      }t        |      dk  sJ y )Nr%   r   )d   r0   r0      )r)   r*   broadcast_tor   )r,   xs     r   test_numpy_0_stridedr4   0   s2    			W	%B
?+A!9>>r   c                 n   t        j                  g dg ddg d      } t        |       t        | j                        t        | j                        z   t        | j
                        z
  k\  sJ t        | j                        t        | j
                        k\  sJ t        | j                        dk\  sJ t        | j
                        dk\  sJ t        t        |       t              sJ t        t        | j                        t              sJ t        t        | j
                        t              sJ y )Nr   r   r   daaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaadbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbdccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccr3   y
         indexi,  r?   )pd	DataFramer   r3   r<   rB   
isinstanceint)dfs    r   test_pandasrH   6   s    	?@
B ":rtt4vbhh7GGGGG"$$<6"((++++"$$<7""""((r!!!fRj#&&&fRTTlC(((fRXX&,,,r   c                     t        j                  ddgddgg      } t        j                  ddgddgg      }t        |      t        |       k  sJ y)zv2+ contiguous columns of the same dtype in the same DataFrame share the same
    surface thus have lower overhead
    r   g@r   g@g?gffffff
@NrC   rD   r   )df1df2s     r   test_pandas_contiguous_dtypesrM   F   sP    
 ,,C1c(+
,C
,,c
S#J/
0C#;$$$r   c                    t         j                  j                  t        d      t	        d      g      } t        j                  |       }d|z  t        |       cxk  r	d|z  k  sJ  J t        t        |       t              sJ y )N2   abcdefghilmnopqrstuvwxyzg      ?r   )
rC   
MultiIndexfrom_productr"   listr   r   r   rE   rF   )rB   actual_sizes     r   test_pandas_multiindexrU   P   sm    MM&&b	48R3S'TUE--&Kve}>q;>>>>>fUmS)))r   c                 r   t        j                  dt        t        d            i      } | g d   }t        j                  t        t        d            t        t        d            d      }dt	        |       cxk  rdk  sJ  J dt	        |      cxk  rdk  sJ  J dt	        |      cxk  rdk  sJ  J y )	Nr3   '  )r3   r3   r3   r;   i8 iL i q i )rC   rD   rS   r"   r   )rG   rL   df3s      r   test_pandas_repeated_columnrY   Y   s    	sDv/0	1B
_
C
,,T%-0tE&M7JK
LCF2J''''''F3K(&(((((VC[*7*****r   c                    t        j                  d      } | j                  d      }t        |j	                               dk\  sJ t        |j                               dk\  sJ t        |j                               dk\  sJ t        |j                               dk\  sJ t        |j                               dk\  sJ t        |j                               dk\  sJ t        |j                               dk\  sJ y )Nzscipy.sparser>                  )r)   r*   eyer   todiatobsrtocootocsctocsrtodoktolil)sparsesps     r   test_sparse_matrixrj   c   s      0F	BB"((*$$$"((*$$$"((*$$$"((*$$$"((*$$$"((*$$$"((*$$$r   cls_name)SeriesrD   Indexr(   zstring[python]c                   t        t        |      } |t        d      D cg c]  }d|d
 c}|       }t        d      dz  t        |      cxk  rdt        d      z  dz  k  sJ  J ddz  }ddz  }d	dz  }d
dz  } |||||gdz  |       }	dt        |	      cxk  rdk  sJ  J  |||||g|       }
 |||||g|       } |||||g|       }t        |      t        |      cxk  rt        |
      k  sJ  J y c c}w )Nr   r3   3dr'   x000r   順 r<   zw i  )getattrrC   r"   r   )r(   rk   clsis1r3   r<   rr   rs   s2s3s4s5s                r   test_pandas_object_dtyper}   p   s+    "h
C	%+.Q!Bj.e	<B&>D 6":IF6N0BT0IIIIIIgAgAgAgA 
aAq\D 	.BVBZ)'))))) 
aAq\	'B	aAq\	'B	aAq\	'B":r
/VBZ/////! /s   C>c                   ddz  }ddz  }ddz  }ddz  }||||g}t        j                  |dz  gdz  |       }d	t        |      cxk  rd
k  sJ  J t        j                  ||g||gg|       }t        j                  ||g||gg|       }t        j                  ||g||gg|       }	t        |	      t        |      cxk  rt        |      k  sJ  J y )Nr3   rq   r<   rr   rs   r   r   r'   rt   ipd rJ   )
r(   r3   r<   rr   rs   objsrK   rL   rX   df4s
             r   test_dataframe_object_dtyper      s     	gAgAgAgA q!Q<D
,,qzD(
6CVC[*7***** ,,AA'u
5C
,,AA'u
5C
,,AA'u
5C#;2vc{22222r   c                    t        j                  d       t        t        |       } |ddz  ddz  gd      }dt	        |      cxk  rd	k  sJ  J y )
Npyarrowr3   rq   r<   iP  zstring[pyarrow]r'   iI ix] )r)   r*   ru   rC   r   )rk   rv   ss      r   test_pandas_string_arrow_dtyper      sS    
	"
"h
CS7]C&L)1BCAVAY((((((r   c                 "   t        j                  g dg ddg d      } | j                  d      }t        |      dkD  sJ t        |j                        dkD  sJ t        |j
                        dkD  sJ t        |j                        dkD  sJ y )Nr6   r7   r;   r=   rA   r   )rC   rD   headr   r3   r<   rB   )rG   r+   s     r   test_pandas_emptyr      s    	?@
B GGAJE%=1%''?Q%''?Q%++"""r   c                    t        j                  d      } t        j                  g dg ddg d      }| j                  j                  |      }t        |      t        |j                  j                        kD  sJ t        t        |      t              sJ t        t        |j                  d         t              sJ t        t        |j                  d         t              sJ t        t        |j                  d	         t              sJ | j                  j                  |j                  d            }t        |      t        |j                  j                        kD  sJ t        |j                  d         dkD  sJ t        |j                  d         dkD  sJ t        |j                  d	         dkD  sJ y )
Nr   r6   r7   r;   r=   rA   r   r   r   )r)   r*   rC   rD   Tablefrom_pandasr   schemametadatarE   rF   columnsr   )parG   tabler+   s       r   test_pyarrow_tabler      sa   			Y	'B	?@
B HH  $E%=6%,,"7"78888fUmS)))fU]]1-.444fU]]1-.444fU]]1-.444HH  ,E%=6%,,"7"78888%--"#a'''%--"#a'''%--"#a'''r   c                    t        j                  d      } | j                  d      }t        d|i      |j                  kD  sJ t        d|gi      |j                  kD  sJ t        dd|igi      |j                  kD  sJ t
        j                  t        d      |      }t        |      |j                  dz  kD  sJ t        t        |      t              sJ y )Nr%   rW   r3   r<   r0   )
r)   r*   onesr   nbytesdictfromkeysr"   rE   rF   )r,   r3   r   s      r   	test_dictr      s    			W	%B
A3(ahh&&&3*(((3#q
#$qxx///eCj!$A!9qxx#~%%%fQi%%%r   c                    t         j                  j                  t        j                  |              dd l}|j                  j                          dd l}|j                  |      }t        |      S )Nr   )
r   pathappendosfsdecodedask.sizeofr   _register_entry_point_plugins
class_implImpl)r   sizedaskr   rv   s        r   _get_sizeof_on_pathr      sK    HHOOBKK%& KK--/
//$
C#;r   c                `   | dz  j                  d       | dz  j                  d       | dz  }|j                          |dz  j                  d       t               j                  d      5 }|j	                  t
        | d	f
      d	k(  sJ 	 d d d        j                          y # 1 sw Y   xY w)Nzimpl_sizeof.pys   def sizeof_plugin(sizeof):
    print("REG")
    @sizeof.register_lazy("class_impl")
    def register_impl():
        import class_impl
        @sizeof.register(class_impl.Impl)
        def sizeof_impl(obj):
            return obj.size 
zclass_impl.pysB   class Impl:
    def __init__(self, size):
        self.size = sizezimpl_sizeof-0.0.0.dist-infozentry_points.txts/   [dask.sizeof]
impl = impl_sizeof:sizeof_plugin
r   i)args)write_bytesmkdirr   Poolapplyr   join)tmp_path	dist_infopools      r    test_register_backend_entrypointr      s      --	*	 ,,O 88IOO##00< 
		A	 
$JJ*(J1GJHJV	
V
 	IIK	
 
s   +B$$B-c                 L   t        j                  d      } t        j                  d      }|j                  ddd      j                  t              }|j
                  j                  t        |      f      }| j                  |dgd|i      j                  d      j                         }t        |      t        |      kD  sJ t        |j                        t        |      k\  sJ t        |d         t        |      k\  sJ t        |j                        t        |      k\  sJ y )	Nxarrayr%   C   r   coorddimscoordsfoo)r)   r*   arangeastypefloatrandomlen	DataArrayrename
to_datasetr   r   indexes)xrr,   indarrdatasets        r   test_xarrayr      s    			X	&B			W	%B
))CQ

&
&u
-C
))

CH;
'C 	S> 	 	

 
	  '?VC[((('++&+---''"#vc{222'//"fSk111r   c                    t        j                  d      } t        j                  d      }t        j                  d      }t        j                  d       |j                  ddd      j                  t              }|j
                  j                  t        |      f      }t        |j                        t        d      kD  r6t        |j                        t        d      k  rt        j                  d	       t               5 }| j                  |d
gd
|i      j                  d      j                         j                  |       | j                  |ddi      }|j                   j"                  rJ t%        |      t%        |      cxk  rt%        |      t%        |      z   k  sJ  J t%        |j                         t%        |      k  sJ t%        |d
         t%        |      k\  sJ t%        |j&                        t%        |      k\  sJ |j                   j"                  rJ |j)                          |j                   j"                  sJ t%        |      t%        |      t%        |      z   k\  sJ 	 d d d        y # 1 sw Y   y xY w)Nr   zarrr%   r   r   r   z3.0.0.a0z3.0.0z8consolidated metadata and xarray support is not completer   r   r   r>   )chunks)r)   r*   r   r   r   r   r   r   __version__xfailr	   r   r   r   to_zarr	open_zarrr   
_in_memoryr   r   load)r   r   r,   r   r   r   r   s          r   test_xarray_not_in_memoryr     s   			X	&Bv&D			W	%B

))CQ

&
&u
-C
))

CH;
'C t 7:#667<< 	OP	 <T
S> 	 	
 &

*,wwt},,tUBK,8;;))))c{VG_Hvc{VC[/HHHHHHgkk"VC[000gg&'6#;666goo&&+555;;)))){{%%%%g&+s";;;;%< < <s   E/I::J),
__future__r   r   r   r   r)   packaging.versionr   dask.multiprocessingr   r   r   
dask.utilsr   r	   pandasrC   ImportErrormarkskipifrequires_pandasr   r   r   r#   r.   r4   rH   rM   rU   rY   rj   parametrizeobjectr}   r   r   r   r   r   r   r   r   r   r   r   r   <module>r      s   " 	 
   % ,  ' ++$$R4Z8I$J)(L;+ - - % % * * + +
% %EF6+;"<=0 > G 0* 6+;"<=3 > 3$ %EF) G) 	# 	# ( (*	&:2,!<M  	Bs   E EE