o
    g                     @   s  d Z ddlZddlmZ ddlZddlZddlmZ ddlm	Z
 ejddgddd	 Zejd
d Zejdg ddd Zejdddg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d d! Zejd"g d#d$d% Zd&d' ZdS )(zZ
Tests compressed data parsing functionality for all
of the parsers defined in parsers.py
    N)Path)	DataFrameTF)paramsc                 C   s   | j S )N)param)request r   e/home/ubuntu/cloudmapper/venv/lib/python3.10/site-packages/pandas/tests/io/parser/test_compression.pybuffer   s   r	   c                 C   sJ   | }t |d}| }W d    n1 sw   Y  ||}|||fS )Nrb)openreadread_csv)all_parserscsv1parserfdataexpectedr   r   r   parser_and_data   s   


r   compression)zipinferzip2c           	   	   C   s   | \}}}t dT}tj|dd}|d| W d    n1 s#w   Y  |dkrIt|d}|j|dd}W d    n1 sCw   Y  n|j||d}t || W d    d S 1 saw   Y  d S )	Ntest_file.zipwmode	test_filer   r
   r   r   )tmensure_cleanzipfileZipFilewritestrr   r   assert_frame_equal)	r   r   r   r   r   pathtmpr   resultr   r   r   test_zip    s   
"r(   r   r   c           	   	   C   s   | \}}}t dU}ddg}tj|dd}|D ]}||| qW d    n1 s,w   Y  tjtdd |j||d W d    n1 sJw   Y  W d    d S W d    d S 1 sbw   Y  d S )	Nzcombined_zip.zipr   second_filer   r   zMultiple filesmatchr   )	r   r    r!   r"   r#   pytestraises
ValueErrorr   )	r   r   r   r   r   r%   inner_file_namesr&   	file_namer   r   r   test_zip_error_multiple_files1   s   
"r1   c              	   C   s   | \}}}t  F}tj|dd W d    n1 sw   Y  tjtdd |j|dd W d    n1 s:w   Y  W d    d S W d    d S 1 sRw   Y  d S )Nr   r   z
Zero filesr*   r   r   )r   r    r!   r"   r,   r-   r.   r   )r   r   _r%   r   r   r   test_zip_error_no_files@   s   

"r3   c              
   C   s   | \}}}t  E}t|d'}tjtjdd |j|dd W d    n1 s*w   Y  W d    n1 s9w   Y  W d    d S W d    d S 1 sQw   Y  d S )Nr
   zFile is not a zip filer*   r   r   )r   r    r   r,   r-   r!   
BadZipfiler   )r   r   r2   r%   r   r   r   r   test_zip_error_invalid_zipK   s   

"r5   filenamez
test.{ext}c              	   C   s   | \}}}|}|dkrdn|}|d u r|n|j |d}|r$|r$td tj|dB}	t||	| |r6dn|}
|rWt|	d}|j||
d}W d    n1 sQw   Y  n|j|	|
d}t|| W d    d S 1 sow   Y  d S )	Ngzipgz)extz9Cannot deduce compression from buffer of compressed data.)r6   r   r
   r   )	formatr,   skipr   r    write_to_compressedr   r   r$   )r   compression_onlyr	   r6   r   r   r   compress_typer9   r%   r   r   r'   r   r   r   test_compressionT   s"   

"r?   r9   )Nr8   bz2c           	      C   s   | }ddd}|j |fi |}d|d< |r4t|}|j |fi |}W d    n1 s.w   Y  n|r:d| nd}|j || fi |}t|| d S )Nr   T)	index_colparse_datesr   r   . )r   r   r   r$   )	r   r   r	   r9   r   kwargsr   r   r'   r   r   r   test_infer_compressionl   s   

rF   c           	      C   s\   | }| |}tj|d| d}|j||ddd}tddgddgd	}t|| d S )
Nutfz_ex_small.zipr   	)encodingr   sep	Venezuelau   Hugo Chávez FríaszHenrique Capriles R.)CountryTwitter)r:   osr%   joinr   r   r   r$   )	r   csv_dir_path	utf_valueencoding_fmtr   rI   r%   r'   r   r   r   r   test_compression_utf_encoding   s   
rS   invalid_compression)sfarkbz3zipperc                 C   sZ   | }d|i}d| }t jt|d |jdi | W d    d S 1 s&w   Y  d S )Nr   zUnrecognized compression type: r*   r   )r   )r,   r-   r.   r   )r   rT   r   compress_kwargsmsgr   r   r   test_invalid_compression   s   
"rZ   c              	   C   s   | }t dddgi}tdC}td%}|j|dd t|t|  t|j|d d| W d    n1 s=w   Y  W d    d S W d    d S 1 sUw   Y  d S )	Nar      ztest.csvztest.csv.zipF)indexr   )	r   r   r    to_csvr   
write_text	read_textr$   r   )r   r   dfpath_csvpath_zipr   r   r   !test_ignore_compression_extension   s   "rd   )__doc__rN   pathlibr   r!   r,   pandasr   pandas._testing_testingr   fixturer	   r   markparametrizer(   r1   r3   r5   r?   rF   rS   rZ   rd   r   r   r   r   <module>   s2    




	



