o
    d>$                     @   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&g d'e"d#d(d) Z)ejj'ej%e"d#e"dd*d+ Z*e+d,d-d. Z,e"dd/d0 Z-d1d2 Z.d3d4 Z/d5d6 Z0e"d7d8d9 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   O/app/.heroku/python/lib/python3.10/site-packages/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   sh   t |jdd }| dd}|| W d    n1 s w   Y  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   s   ddl m} tjtdd td W d    n1 sw   Y  d}| |dd|d	 tjt|d td
 W d    d S 1 sDw   Y  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r#J W d    n1 s-w   Y  |dd}|j||d}t|d	gd
d}|jrMJ W d    n1 sWw   Y  t|| d S )Nr   r%   r$   r.   )modeTr"   rr   r   r<   )	r   r   r)   r'   closedreplacer   r+   r,   )r!   rD   r   r   rB   rE   Zfsspec_objectr/   r   r   r   test_to_csv_fsspec_objectV   s"   
rI   c                 C   Z   t ddgi}|jdddidd | jd dksJ tdddid	 | jd dks+J d S 
Nar   testmem://test/test.csvtestZ	csv_writeFstorage_optionsr#   Zcsv_readrP   )r   r'   rN   r   
fsspectestdfr   r   r   test_csv_optionsl   s   
rU   c                 C   rJ   rK   )r   r'   rN   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s&J t|dd
id | jd d
ks7J d S )Nr>   r?   rL   r   ztestmem://test/test.rN   r*   FrO   readrQ   )r   r   r   r@   rN   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   ^   t ddgi}|jdddddid | jd dksJ tdddd	id
 | jd d	ks-J dS )r[   rL   r   rM   r`   NrN   parquet_writer\   r]   rP   parquet_readr\   rP   r   r^   rN   r   rR   r   r   r   test_arrowparquet_options   s   rg   c                 C   ra   )r[   rL   r   rM   rZ   NrN   rb   rc   rd   re   rf   rR   r   r   r   test_fastparquet_options   s   rh   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   

rn   protocol)Zs3Zs3aZs3nc                 C   s"   t t| d|dt| d S )Nz://pandas-test/tips.csvrQ   ri   )rk   rl   ro   rm   r   r   r   test_s3_protocols   s   rp   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   re   )r^   r   r+   rj   )rk   rm   r   fnr/   r   r   r   test_s3_parquet   s   
rr   r   c                  C   s@   d} t jt| d td W d    d S 1 sw   Y  d S )Nz?Missing optional dependency 'fsspec'|fsspec library is requiredr2   r%   )r   r6   r9   r   )msgr   r   r   test_not_present_exception   s   
"rt   c                 C   d   t ddgi}|jdddid | jd dksJ tdddid}| jd dks*J t|| d S )NrL   r   testmem://afilerN   Zfeather_writerQ   Zfeather_read)r   Z
to_featherrN   r   r+   r,   rS   rT   outr   r   r   test_feather_options   s   ry   c                 C   ru   )NrL   r   rv   rN   Zpickle_writerQ   Zpickle_read)r   Z	to_picklerN   r	   r+   r,   rw   r   r   r   test_pickle_options   s   rz   c                 C   sh   t ddgi}|jd|ddid | jd dksJ td|ddid}| jd dks,J t|| d S )NrL   r   rv   rN   Z
json_write)r]   rP   Z	json_read)r   to_jsonrN   r   r+   r,   )rS   r]   rT   rx   r   r   r   test_json_options   s   r|   c                 C   sl   t ddgi}|jdddidd | jd dksJ tdddid	}| jd dks+J t||d
 d S )NrL   r   rv   rN   Zstata_writeF)rP   Zwrite_indexZ
stata_readrQ   Zint64)r   Zto_statarN   r
   r+   r,   Zastyperw   r   r   r   test_stata_options  s   
r}   tabulatec                 C   sD   t ddgi}|jdddid | jd dksJ | ds J d S )NrL   r   rv   rN   Zmd_writerQ   )r   Zto_markdownrN   catrR   r   r   r   test_markdown_options  s   r   c                  C   s  t jtdd tdddid W d    n1 sw   Y  t jtdd tdddid W d    n1 s9w   Y  t } t jtdd t| ddid W d    n1 s\w   Y  tddgi}t jtdd |jdddid W d    d S 1 sw   Y  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   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parametrizerI   rU   rV   rY   Z
skip_if_nor_   rg   Z&skip_array_manager_not_yet_implementedrh   Z
single_cpurn   rp   rr   Zskip_if_installedrt   ry   rz   r|   r}   r   r   r   r   r   r   <module>   sb    0

	



	

		
