a
    Pf                     @   s   d dl mZ d dlZd dlZd dlZd dlZd dlZd dlZd dl	m
Z
mZmZmZmZmZmZ d dlmZ d dlmZ ejdZejdd Zejjejdg d	d
d Zeee dddZ!ejdddgdd Z"dd Z#e$ddd Z%dS )    )BytesION)	DataFrameIndex
date_rangeread_csv
read_excel	read_jsonread_parquet)_test_decoratorsz=ignore:Passing a BlockManager to DataFrame:DeprecationWarningc                     sN   t d t d} t  dd  _G  fddd| j}| jd|dd	  S )
z"Emulate GCS using a binary buffer.gcsfsfsspecc                   S   s   dS )NT r   r   r   Q/var/www/ai-form-bot/venv/lib/python3.9/site-packages/pandas/tests/io/test_gcs.py<lambda>"       zgcs_buffer.<locals>.<lambda>c                       s$   e Zd Ze fddZdd ZdS )z%gcs_buffer.<locals>.MockGCSFileSystemc                     s     d  S )Nr   )seek)argskwargs
gcs_bufferr   r   open%   s    
z*gcs_buffer.<locals>.MockGCSFileSystem.openc                 [   s   |ddgS )Nfile)nametyper   )selfpathr   r   r   r   ls*   s    z(gcs_buffer.<locals>.MockGCSFileSystem.lsN)__name__
__module____qualname__staticmethodr   r   r   r   r   r   MockGCSFileSystem$   s   r!   gsT)clobber)pytestimportorskipr   closeAbstractFileSystemZregister_implementation)r   r!   r   r   r   r      s    


r   format)csvjsonparquetexcelmarkdownc           
         sp  t ddgdtjgddgtdddd	}d
| }|dkr\|j|dd t|dgdd}n|dkrd}|| t|dgdd}n|dkr|| t	|dgd}n|dkr>t
d t
d G  fddd j}| 0}| d| || t|}W d   n1 s0    Y  | }	|	jdks`J n"|dkr`t
d || |}t|| dS )zE
    Test that many to/read functions support GCS.

    GH 33987
                 @ts
2018-06-18   Zperiodsintfloatstrdtzgs://test/test.r)   T)indexr:   r   )Zparse_dates	index_colr,   zgs://test/test.xlsxr*   )Zconvert_datesr+   Zpyarrowz
pyarrow.fsc                       s   e Zd Ze fddZdS )z(test_to_read_gcs.<locals>.MockFileSystemc                    s,   t d t| dd  } |S )NzUsing pyarrow filesystemzgs:// )printpathlibPathreplaceabsoluteas_uriZLocalFileSystem)r   Zto_localZpa_fsr   r   from_uriX   s    z1test_to_read_gcs.<locals>.MockFileSystem.from_uriN)r   r   r   r    rE   r   rD   r   r   MockFileSystemW   s   rF   
FileSystemNz2Using pyarrow filesystem
Using pyarrow filesystem
r-   tabulate)r   npnanr   to_csvr   Zto_excelr   to_jsonr   r$   r%   rG   contextsetattr
to_parquetr	   Z
readouterroutZto_markdowntmassert_frame_equal)
r   r(   monkeypatchZcapsysdf1r   Zdf2rF   mZcapturedr   rD   r   test_to_read_gcs5   sB    	
	







(


rV   )resultexpectedcompressionc              	   C   s~  |dkrt t| j}t t|<}t| | D ]\}}|j|jks:J q:W d   n1 sh0    Y  W d   n1 s0    Y  n|dkrltjt| d}tjt|dz}t| | D ]V\}	}
|	|	}|	|
}|du |du ksJ |dur|dur|
 |
 ksJ qW d   n1 s@0    Y  W d   n1 s`0    Y  n| |kszJ dS )a  
    For zip compression, only compare the CRC-32 checksum of the file contents
    to avoid checking the time-dependent last-modified timestamp which
    in some CI builds is off-by-one

    See https://en.wikipedia.org/wiki/ZIP_(file_format)#File_headers
    zipNtar)fileobj)zipfileZipFiler   rZ   infolistCRCtarfiler   
getmembersextractfileread)rW   rX   rY   expresZres_infoZexp_infoZtar_expZtar_resZtar_res_infoZtar_exp_infoZactual_fileZexpected_filer   r   r   assert_equal_zip_safel   s*    P


Xrg   encodingutf-8cp1251c                 C   s&  t dtdd ttdtdtdd tdD tdd	}d
|i}|dkrXd|d< t }|j	|||dd d}|j	|||d | 
 }|
 }	t||	| t|d||d}
t||
 || }d|d
< |d| 7 }|j	|||d | 
 }|
 }	t||	| t|dd|d}
t||
 dS )z
    Compression and encoding should with GCS.

    GH 35677 (to_csv, compression), GH 26124 (to_csv, encoding), and
    GH 32392 (read_csv, encoding)
    g?x   )      ZABCD)Zdtypec                 S   s   g | ]}d | qS )zi-r   ).0ir   r   r   
<listcomp>   r   z8test_to_csv_compression_encoding_gcs.<locals>.<listcomp>rl   )columnsr;   methodgzipr.   mtimewb)rY   rh   modegs://test/test.csv)rY   rh   r   )r<   rY   rh   Zinfer.N)r   rI   ZarangeZreshaper   listobjectranger   rK   getvaluerg   r   rQ   rR   )r   Zcompression_onlyrh   Zcompression_to_extensionZdfrY   bufferZpath_gcsrf   rX   Zread_dfZfile_extr   r   r   $test_to_csv_compression_encoding_gcs   s8    
r~   c                    s~   t d t d ddlm} tddgdtjgdd	gtd
ddd}G  fddd|}| d| |j	ddddd dS )zCRegression test for writing to a not-yet-existent GCS Parquet file.Zfastparquetr   r   )r'   r.   r/   r0   r1   r2   r3   r4   r5   r6   c                       s   e Zd Zd fdd	ZdS )z7test_to_parquet_gcs_new_file.<locals>.MockGCSFileSystemrc                    s$   d|vrt ttj d|ddS )Nwztest.parquetri   )rh   )FileNotFoundErrorr   osr   join)r   r   rv   r   tmpdirr   r   r      s    z<test_to_parquet_gcs_new_file.<locals>.MockGCSFileSystem.openN)r   )r   r   r   r   r   r   r   r   r!      s   r!   zgcsfs.GCSFileSystemrw   TN)r;   ZenginerY   )
r$   r%   r   r'   r   rI   rJ   r   rN   rO   )rS   r   r'   rT   r!   r   r   r   test_to_parquet_gcs_new_file   s    


	r   r   c                   C   s6   t t td W d    n1 s(0    Y  d S )Nrw   )rQ   Zexternal_error_raisedImportErrorr   r   r   r   r   test_gcs_not_present_exception   s    r   )&ior   r   r?   ra   r]   numpyrI   r$   Zpandasr   r   r   r   r   r   r	   Zpandas._testingZ_testingrQ   Zpandas.utilr
   tdmarkfilterwarningsZ
pytestmarkZfixturer   Z
single_cpuZparametrizerV   bytesr9   rg   r~   r   Zskip_if_installedr   r   r   r   r   <module>   s.   $	
5
0