U
    >vf>$                     @   s  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mZmZmZ d dlmZ d dlmZ ejdd Zejdd Zdd	 Zd
d Zdd Zdd Zejdddgdd Zdd Zdd Z dd Z!e"ddd Z#e"ddd  Z$ej%e"dd!d" Z&ejj'e"d#d$d% Z(ejj'ejd&d'd(d)ge"d#d*d+ Z)ejj'ej%e"d#e"dd,d- Z*e+d.d/d0 Z,e"dd1d2 Z-d3d4 Z.d5d6 Z/d7d8 Z0e"d9d:d; Z1e"dd<d= Z2dS )>    N)
	DataFrame
date_rangeread_csv
read_excelread_feather	read_jsonread_parquetread_pickle
read_stata
read_table)_test_decoratorsc                   C   s(   t ddgdtjgddgtdddd	S )
N      g       @tsz
2018-06-18   )Zperiods)intfloatstrdt)r   npnanr    r   r   ?/tmp/pip-unpacked-wheel-vdrwu74i/pandas/tests/io/test_fsspec.pydf1   s    
r   c                  c   s(   t d} | d}|V  |j  d S )NfsspecZmemory)pytestimportorskip
filesystemstoreclear)r   Zmemfsr   r   r   
cleared_fs"   s    

r!   c              	   C   sT   t |jdd }| dd}|| W 5 Q R X tddgd}t|| d S )NFindexztest/test.csvwbmemory://test/test.csvr   )parse_dates)r   to_csvencodeopenwriter   tmassert_frame_equal)r!   r   textwdf2r   r   r   test_read_csv+   s
    r0   c              	   C   sl   ddl m} tjtdd td W 5 Q R X d}| |dd|d	 tjt|d td
 W 5 Q R X d S )Nr   )known_implementationsZnosuchprotocolmatchznosuchprotocol://test/test.csvztest error messageZ
couldexistzunimportable.CouldExist)classerrzcouldexist://test/test.csv)Zfsspec.registryr1   r   raises
ValueErrorr   setitemImportError)Zmonkeypatchr!   r1   err_msgr   r   r   test_reasonable_error4   s    r;   c                 C   s.   |j ddd tddgdd}t|| d S )Nr%   Tr"   r   r   r&   Z	index_col)r'   r   r+   r,   )r!   r   r/   r   r   r   test_to_csvC   s    r=   c                 C   sF   t d d}d| }|j|dd t|dgdd}t|| d S )	Nopenpyxlxlsxzmemory://test/test.Tr"   r   r   r<   )r   r   to_excelr   r+   r,   )r!   r   extpathr/   r   r   r   test_to_excelK   s    

rC   binary_modeFTc              	   C   s   t d}d}|rdnd}|j||d }|j|dd |jrFtW 5 Q R X |dd}|j||d }t|d	gd
d}|jrtW 5 Q R X t	|| d S )Nr   r%   r$   r.   )modeTr"   rr   r   r<   )
r   r   r)   r'   closedAssertionErrorreplacer   r+   r,   )r!   rD   r   r   rB   rE   Zfsspec_objectr/   r   r   r   test_to_csv_fsspec_objectV   s    
rJ   c                 C   sZ   t ddgi}|jdddidd | jd dks4ttdddid	 | jd dksVtd S 
Nar   testmem://test/test.csvtestZ	csv_writeFstorage_optionsr#   Zcsv_readrP   )r   r'   rN   rH   r   
fsspectestdfr   r   r   test_csv_optionsl   s      rU   c                 C   sZ   t ddgi}|jdddidd | jd dks4ttdddid	 | jd dksVtd S rK   )r   r'   rN   rH   r   rR   r   r   r   test_read_table_optionsv   s      rV   c                 C   sr   t d d}tddgi}d| }|j|ddidd	 | jd dksLtt|dd
id | jd d
ksntd S )Nr>   r?   rL   r   ztestmem://test/test.rN   r*   FrO   readrQ   )r   r   r   r@   rN   rH   r   )rS   	extensionrT   rB   r   r   r   test_excel_options   s    

rY   fastparquetc                 C   s   |j ddddd dS )CRegression test for writing to a not-yet-existent GCS Parquet file.r%   TrZ   N)r#   enginecompression)
to_parquet)r!   r   r   r   r   test_to_parquet_new_file   s       r_   pyarrowc                 C   s^   t ddgi}|jdddddid | jd dks6ttdddd	id
 | jd d	ksZtdS )r[   rL   r   rM   r`   NrN   parquet_writer\   r]   rP   parquet_readr\   rP   r   r^   rN   rH   r   rR   r   r   r   test_arrowparquet_options   s    rf   c                 C   s^   t ddgi}|jdddddid | jd dks6ttdddd	id
 | jd d	ksZtdS )r[   rL   r   rM   rZ   NrN   ra   rb   rc   rd   re   rR   r   r   r   test_fastparquet_options   s    rg   Zs3fsc                 C   sL   t td|dt| t td|dt| t td|dt| d S )Nzs3://pandas-test/tips.csvrQ   zs3://pandas-test/tips.csv.gzzs3://pandas-test/tips.csv.bz2r+   assert_equalr   )s3_resource	tips_files3sor   r   r   test_from_s3_csv   s    
 

rm   protocolZs3Zs3aZs3nc                 C   s"   t t| d|dt| d S )Nz://pandas-test/tips.csvrQ   rh   )rj   rk   rn   rl   r   r   r   test_s3_protocols   s    ro   c                 C   s6   d}|j |ddd |d t|d|d}t|| d S )Nzs3://pandas-test/test.parquetFrZ   )r#   r\   r]   rP   rd   )r^   r   r+   ri   )rj   rl   r   fnr/   r   r   r   test_s3_parquet   s        rq   r   c               	   C   s*   d} t jt| d td W 5 Q R X d S )Nz?Missing optional dependency 'fsspec'|fsspec library is requiredr2   r%   )r   r6   r9   r   )msgr   r   r   test_not_present_exception   s    rs   c                 C   sd   t ddgi}|jdddid | jd dks2ttdddid}| jd dksTtt|| d S )NrL   r   testmem://afilerN   Zfeather_writerQ   Zfeather_read)r   Z
to_featherrN   rH   r   r+   r,   rS   rT   outr   r   r   test_feather_options   s    rw   c                 C   sd   t ddgi}|jdddid | jd dks2ttdddid}| jd dksTtt|| d S )NrL   r   rt   rN   Zpickle_writerQ   Zpickle_read)r   Z	to_picklerN   rH   r	   r+   r,   ru   r   r   r   test_pickle_options   s    rx   c                 C   sh   t ddgi}|jd|ddid | jd dks4ttd|ddid}| jd dksXtt|| d S )NrL   r   rt   rN   Z
json_write)r]   rP   Z	json_read)r   Zto_jsonrN   rH   r   r+   r,   )rS   r]   rT   rv   r   r   r   test_json_options   s    ry   c                 C   sl   t ddgi}|jdddidd | jd dks4ttdddid	}| jd dksVtt||d
 d S )NrL   r   rt   rN   Zstata_writeF)rP   Zwrite_indexZ
stata_readrQ   Zint64)r   Zto_statarN   rH   r
   r+   r,   Zastyperu   r   r   r   test_stata_options  s      rz   Ztabulatec                 C   sD   t ddgi}|jdddid | jd dks2t| ds@td S )NrL   r   rt   rN   Zmd_writerQ   )r   Zto_markdownrN   rH   catrR   r   r   r   test_markdown_options  s    r|   c               	   C   s   t jtdd tdddid W 5 Q R X t jtdd tdddid W 5 Q R X t } t jtdd t| ddid W 5 Q R X tddgi}t jtdd |jdddid W 5 Q R X d S )	NrP   r2   	localfilerL   TrQ   r   Znonfsspecpath)	r   r6   r7   r   r   ioBytesIOr   r^   )ZbyrT   r   r   r   test_non_fsspec_options$  s    r   )3r~   Znumpyr   r   Zpandasr   r   r   r   r   r   r   r	   r
   r   Zpandas._testingZ_testingr+   Zpandas.utilr   tdZfixturer   r!   r0   r;   r=   rC   markZparametrizerJ   rU   rV   rY   Z
skip_if_nor_   rf   Z&skip_array_manager_not_yet_implementedrg   Z
single_cpurm   ro   rq   Zskip_if_installedrs   rw   rx   ry   rz   r|   r   r   r   r   r   <module>   s`   0

	



	

		
