o
    ҷh                     @   sr  d dl Zd dlZd dlm  mZ d dlmZ d dl	m
Z
mZmZmZ d dlmZ ejdddddggdd Zejdg d	g d
g dg dejg dejjdgdd Zejdddddggejdddgdd Zdd Zdd Zdd Zejdddgdd Zejdg d d!d" Zejdeejd#ed$dejd%ed$dgd&d' ZdS )(    N)is_integer_dtype)	DataFrameIndexPeriodIndexSeriesbyABc                 C   s:   | j |d}| }|D ]\}}|| t|ksJ qd S )Nr   )groupbysizelen)dfr   groupedresultkeygroup r   Q/var/www/html/venv/lib/python3.10/site-packages/pandas/tests/groupby/test_size.py	test_size   s
   r   )r   r   r   r   )r      r   r   )r   r   r   r   )r   NNN)NNNN)marksc           
         s    fddt  D }|rdd | D }t|dd}|r#| }t|jjr9tdd  D s9|j	t
|_d}tjt|d	 | j |||d
}W d    n1 sVw   Y  | }	t|	| d S )Nc                    s$   i | ]  t  fd dD qS )c                 3   s    | ]}| kV  qd S Nr   ).0valuer   r   r   	<genexpr>%       z.test_size_axis_1.<locals>.<dictcomp>.<genexpr>)sum)r   r
   r   r   
<dictcomp>%   s   $ z$test_size_axis_1.<locals>.<dictcomp>c                 S   s   i | ]\}}|d ur||qS r   r   )r   r   r   r   r   r   r   '   s    int64dtypec                 s   s    | ]}|d u V  qd S r   r   )r   xr   r   r   r   +   r   z#test_size_axis_1.<locals>.<genexpr>z+DataFrame.groupby with axis=1 is deprecated)match)r   axissortdropna)dictfromkeysitemsr   
sort_indexr   indexr"   anyastypeinttmassert_produces_warningFutureWarningr   r   assert_series_equal)
r   axis_1r   r&   r'   countsexpectedmsgr   r   r   r
   r   test_size_axis_1   s   r8   r&   TFc                 C   sb   t tjdddtdd}|j|| d }|j|| dd dd	 }t	j
||d
d d S )N      )i     ABCcolumns)r   r&   Cc                 S   s
   | j d S )Nr   )shape)ar   r   r   <lambda>:   s   
 z test_size_sort.<locals>.<lambda>F)check_names)r   nprandomdefault_rngchoicelistr   r   applyr0   r3   )r&   r   r   leftrightr   r   r   test_size_sort5   s    rL   c                  C   s<   t ddgd} tdtg ddd}t| d | d S )Nr   r	   r=   r    namer"   r,   )r   r   r   r0   r3   r   r   )r   outr   r   r   test_size_series_dataframe>   s   rQ   c                  C   sB   t dd d gi} | d }tdtg ddd}t|| d S )Nr   r    rM   rO   r   r   r   r   r   r0   r3   )r   r   r6   r   r   r   test_size_groupby_all_nullE   s   rS   c                  C   s>   t dgtdgdddd} | jdd}| }t||  d S )Nr   2000r   D)rN   freq)r,   level)r   r   r   r   r0   r3   )sergrpr   r   r   r   test_size_period_indexN   s   r[   as_indexc                 C   s   t ddgddggddgd}|d d|d< |jddg| dd }t g d	g d
g dg dgg dd}|d d|d< | rN|ddgd d }t|| d S )Nr   r9   r   r	   r=   categoryF)r\   observedr   r   r   )r   r9   r   )r9   r   r   )r9   r9   r   )r   r	   r   r   )r   r.   r   r   	set_indexrenamer0   assert_equal)r\   r   r   r6   r   r   r   test_size_on_categoricalV   s    rc   r"   )Int64Float64booleanc                 C   sL   t g dg d| d}|jdd }t ddgdd	d
gd}t|| d S )Nr_   rA   rA   b)r,   r"   r   rW   r9   r   rd   rA   rh   rO   )r   r   r   r0   r3   )r"   rY   r   r6   r   r   r   *test_size_series_masked_type_returns_Int64f   s   ri   zstring[pyarrow_numpy]pyarrowstring[pyarrow]c                 C   sj   t g ddd| d}|dd  }| dkrdnd}td	d
gtddgd| dd|d}t|| d S )Nrg   rA   )rA   rh   r!   rh   rk   rd   r    r9   r   )rN   r"   )r,   rN   r"   rR   )r"   r   r   	exp_dtyper6   r   r   r   test_size_stringso   s   
rm   ) numpyrD   pytestpandas.util._test_decoratorsutil_test_decoratorstdpandas.core.dtypes.commonr   pandasr   r   r   r   pandas._testing_testingr0   markparametrizer   paramxfailr8   rL   rQ   rS   r[   rc   ri   object
skip_if_norm   r   r   r   r   <module>   sH    


	

