U
    c,.                     @   s4  d Z ddlmZmZ ddlZddlZddlZddlm	Z	 ddl
m  mZ ddlmZ ddlmZ ddlmZ ddlmZ ddlmZ ejjejd	d
dejdddgejdddgdd Zejjejdd
ddd Zejdd Zejj ej!dejj"de#dde$ G dd dZ%dS )zl
Tests parsers ability to read and parse non-local files
and hence require a network connection to be read.
    )BytesIOStringION)is_ci_environment)	DataFrame)_compression_to_extension)read_featherread_csvVhttps://github.com/pandas-dev/pandas/raw/main/pandas/tests/io/parser/data/salaries.csvT)urlZcheck_before_testmodeexplicitinferenginepythoncc                 C   s@   t | }d}|| }|dkr |}t|d||d}t||  d S )Nr
   r   	)sepcompressionr   )r   r	   tmassert_frame_equal)Zsalaries_tabler   r   Zcompression_only	extensionbase_urlr   Z	url_table r   G/tmp/pip-unpacked-wheel-g7fro6k3/pandas/tests/io/parser/test_network.pytest_compressed_urls   s    r   ghttps://raw.githubusercontent.com/pandas-dev/pandas/main/pandas/tests/io/parser/data/unicode_series.csvc                  C   s(   d} t | ddd}|jd dks$tdS )zN
    read_csv should honor the requested encoding for URLs.

    GH 10424
    r   zlatin-1N)encodingheader)      u$   Á köldum klaka (Cold Fever) (1994))r	   locAssertionError)pathdfr   r   r   test_url_encoding_csv5   s    r%   c                 C   s   t | ddddS )z DataFrame with the tips dataset.iodatacsvtips.csvr   )datapathr   r   r   tips_dfK   s    r+   s3_resourcez$CI race condition GH 45433, GH 44584F)reasonraisesstrictc                   @   s   e Zd Zeddd Zdd Zdd Zdd	 Zd
d Z	dd Z
dd Zdd Zdd Zdd Zejjddddd Zejjdddeddd Zejjdd Zejjejje d d!d"d# Zd$d% Zedd&d' Zd(S ))TestS3s3fsc                 C   sx   dD ]>\}}t d| ||d}t|ts,t|jr6tt|| qt d|d}t|ts^t|jrhtt|| d S )N) N).gzgzip).bz2bz2s3://pandas-test/tips.csv)r   storage_optionszs3://cant_get_it/tips.csvr9   r	   
isinstancer   r"   emptyr   r   selfr+   s3soextcompr$   r   r   r   test_parse_public_s3_bucketZ   s    

z"TestS3.test_parse_public_s3_bucketc                 C   s@   t dd|d}t|tst|jr&tt|jd d | d S )Nzs3n://pandas-test/tips.csv
   nrowsr9   r	   r<   r   r"   r=   r   r   ilocr?   r+   r@   r$   r   r   r   test_parse_public_s3n_bucketo   s    
z#TestS3.test_parse_public_s3n_bucketc                 C   s@   t dd|d}t|tst|jr&tt|jd d | d S )Nzs3a://pandas-test/tips.csvrD   rE   rG   rI   r   r   r   test_parse_public_s3a_bucketw   s    
z#TestS3.test_parse_public_s3a_bucketc                 C   sT   dD ]J\}}t d| d||d}t|ts.t|jr8tt|jd d | qd S )Nr2   r8   rD   )rF   r   r9   rG   r>   r   r   r   !test_parse_public_s3_bucket_nrows~   s    
z(TestS3.test_parse_public_s3_bucket_nrowsc           
   
   C   s   d}dD ]\}}t d| |||dd}|j|ks4tdD ]J}| }t|tsRt|jr\t|j|| ||d   }	t	|	| q8W 5 Q R X qd S )N   r2   r8   )	chunksizer   r9   r   r       r    
r	   rN   r"   Z	get_chunkr<   r   r=   rH   r   r   
r?   r+   r@   rN   rA   rB   Z	df_readerZi_chunkr$   Ztrue_dfr   r   r   #test_parse_public_s3_bucket_chunked   s&    
 
z*TestS3.test_parse_public_s3_bucket_chunkedc           
   
   C   s   d}dD ]\}}t d| ||d|dd}|j|ks6tdD ]J}| }t|tsTt|jr^t|j|| ||d   }	t	|	| q:W 5 Q R X qd S )NrM   r2   r8   r   )rN   r   r   r9   rO   r    rQ   rR   r   r   r   *test_parse_public_s3_bucket_chunked_python   s(    
 
z1TestS3.test_parse_public_s3_bucket_chunked_pythonc                 C   sJ   dD ]@\}}t d| d||d}t|ts.t|jr8tt|| qd S )Nr2   r8   r   r   r   r9   r;   r>   r   r   r   "test_parse_public_s3_bucket_python   s    
z)TestS3.test_parse_public_s3_bucket_pythonc                 C   sF   dD ]<}t d| dd|d}t|ts*t|jr4tt|| qd S )N)r3   r4   r6   r8   r   r   rU   r;   )r?   r+   r@   rA   r$   r   r   r   test_infer_s3_compression   s    
z TestS3.test_infer_s3_compressionc                 C   sV   dD ]L\}}t d| dd||d}t|ts0t|jr:tt|jd d | qd S )Nr2   r8   r   rD   )r   rF   r   r9   rG   r>   r   r   r   (test_parse_public_s3_bucket_nrows_python   s    
z/TestS3.test_parse_public_s3_bucket_nrows_pythonc              	   C   sP   d}t jt|d td|d W 5 Q R X t jt|d td W 5 Q R X d S )N#The specified bucket does not existmatchzs3://nyqpug/asdf.csvr:   zs3://cant_get_it/file.csv)pytestr.   OSErrorr	   )r?   r@   msgr   r   r   test_read_s3_fails   s
    zTestS3.test_read_s3_failszGH#39155 s3fs upgradeF)r-   r/   c              	   C   s@   dd l }t|jjf}tj|dd |jd|d W 5 Q R X d S )Nr   rY   rZ   z/s3://an_s3_bucket_data_doesnt_exit/not_real.csvr:   )botocoreFileNotFoundError
exceptionsClientErrorr\   r.   to_csvr?   r+   r@   r`   errorr   r   r   test_write_s3_csv_fails   s     zTestS3.test_write_s3_csv_failsZpyarrowc              	   C   s@   dd l }t|jjf}tj|dd |jd|d W 5 Q R X d S )Nr   rY   rZ   z3s3://an_s3_bucket_data_doesnt_exit/not_real.parquetr:   )r`   ra   rb   rc   r\   r.   Z
to_parquetre   r   r   r   test_write_s3_parquet_fails   s    z"TestS3.test_write_s3_parquet_failsc              	   C   sj   |j jjddd}t|d  }t|dd}W 5 Q R X t|tsHt|j	rRtt|}t
|| d S )Npandas-testr)   )BucketKeyBodyutf8)r   )metaclientZ
get_objectr   readr	   r<   r   r"   r=   r   r   )r?   r,   Z	tips_fileZ	s3_objectbufferresultexpectedr   r   r   $test_read_csv_handles_boto_s3_object  s     
z+TestS3.test_read_csv_handles_boto_s3_objectzThis test can hang in our CI min_versions build and leads to '##[error]The runner has received a shutdown signal...' in GHA. GH: 45651)r-   c              	   C   s   dd l }ttjddtdd}t }|| t|	 
d}|djd|d	 |j  |jtjd
d, tdd|d ddd |jD kstW 5 Q R X d S )Nr   i    Zabcd)columnszutf-8ri   zlarge-file.csv)rk   rl   r1   )loggerzs3://pandas-test/large-file.csvrM   rE   )r   i  T c                 s   s   | ]}|j d d V  qdS )N)args).0xr   r   r   	<genexpr>4  s     z8TestS3.test_read_csv_chunked_download.<locals>.<genexpr>)r1   r   nprandomZrandnlistr   rd   r   getvalueencoderj   Z
put_objectZS3FileSystemZclear_instance_cacheZat_levelloggingDEBUGr	   recordsr"   )r?   r,   Zcaplogr@   r1   r$   Zstr_bufbufr   r   r   test_read_csv_chunked_download  s    	

z%TestS3.test_read_csv_chunked_downloadc                 C   s   t d|d}t|| d S )Nzs3://pandas-test/tips#1.csvr:   )r	   r   r   )r?   r+   r@   rr   r   r   r   test_read_s3_with_hash_in_key6  s    z$TestS3.test_read_s3_with_hash_in_keyc                 C   s$   t |}t d|d}t|| d S )Nz's3://pandas-test/simple_dataset.featherr:   )r   r   r   )r?   Zfeather_filer@   rs   resr   r   r   test_read_feather_s3_file_path;  s     z%TestS3.test_read_feather_s3_file_pathN)__name__
__module____qualname__tdZ
skip_if_norC   rJ   rK   rL   rS   rT   rV   rW   rX   r_   r\   markxfailrg   rh   
single_cpurt   Zskipifr   r   r   r   r   r   r   r   r0   Q   s6   	



r0   )&__doc__r&   r   r   r   Znumpyr}   r\   Zpandas.compatr   Zpandas.util._test_decoratorsutilZ_test_decoratorsr   Zpandasr   Zpandas._testingZ_testingr   Z pandas.tests.io.test_compressionr   Zpandas.io.feather_formatr   Zpandas.io.parsersr	   r   networkZparametrizer   r%   Zfixturer+   r   Zusefixturesr   ra   Zskip_if_not_us_localer0   r   r   r   r   <module>   sH   

