U
    +ifT                     @   s   d dl mZ d dlZd dlm  mZ d dlZd dl	m
Z d dlm   mZ dd Zdd Zejejjdd Zd	d
 Zdd Zdd Zdd Zejdddgejdddgdd Zdd ZdS )    )BytesIONc              
   C   s   t jdddgdddggddgd	d
dgd}t b}|j|| d t|t j|| d t|| }| 	d}W 5 Q R X t|t | W 5 Q R X d S )Ng~Ϛ?g9DܜJ?g>D)?g.4i(@g33333@g̤AABXYZ)indexcolumnscompressionutf8)
pd	DataFrametmensure_cleanto_jsonassert_frame_equal	read_jsonZdecompress_filereaddecode)r   dfpathfhresult r   X/home/mars/bis/venv/lib/python3.8/site-packages/pandas/tests/io/json/test_compression.pytest_compression_roundtrip   s    
r   c                 C   sD   | dddd}t |}| dddd}t j|dd}t|| d S )Niojsondataztsframe_v012.jsonztsframe_v012.json.zipzipr
   )r   r   r   r   )datapathZuncompressed_pathZuncompressed_dfZcompressed_pathZcompressed_dfr   r   r   test_read_zipped_json   s
    
r"   c              
   C   sv   t d}t >}|j|| d t|d}|djd|d W 5 Q R X W 5 Q R X t jd| |d}t|| d S )	N {"a": [1, 2, 3], "b": [4, 5, 6]}r
   rbzpandas-testztest-1)KeyZBodyzs3://pandas-test/test-1)r   Zstorage_options)	r   r   r   r   r   openZBucketZ
put_objectr   )r   Zs3_resourceZs3sor   r   froundtripped_dfr   r   r   test_with_s3_url(   s    

(  r)   c              	   C   sP   t  >}td}|j|dd| d tj|d| d}t || W 5 Q R X d S )Nr#   recordsTZorientlinesr   )r,   r   )r   r   r   r   r   r   )r   r   r   r(   r   r   r   test_lines_with_compression:   s
    

r-   c              
   C   sh   t  V}td}|j|dd| d tj|dd| d}t|}W 5 Q R X t || W 5 Q R X d S )Nz,{"a": ["foo", "bar", "baz"], "b": [4, 5, 6]}r*   Tr+      )r,   	chunksizer   )r   r   r   r   r   concatr   )r   r   r   resr(   r   r   r   test_chunksize_with_compressionC   s    

   r2   c               
   C   sN   t d} t 2}d}tjt|d | j|dd W 5 Q R X W 5 Q R X d S )Nr#   *Unrecognized compression type: unsupportedmatchunsupportedr
   )r   r   r   r   pytestraises
ValueErrorr   )r   r   msgr   r   r   'test_write_unsupported_compression_typeP   s
    

r;   c               
   C   sD   t  2} d}tjt|d tj| dd W 5 Q R X W 5 Q R X d S )Nr3   r4   r6   r
   )r   r   r7   r8   r9   r   r   )r   r:   r   r   r   &test_read_unsupported_compression_typeX   s    
r<   to_inferTF
read_inferc           
   	   C   s   | }d}|t j| 7 }tddgi}|r.dn|}|r:dn|}t|.}|j||d tj||d}	t|	| W 5 Q R X d S )Nztest.r   r.   Zinferr
   )	icomZ_compression_to_extensionr   r   r   r   r   r   r   )
Zcompression_onlyr>   r=   r   filenamer   Zto_compressionZread_compressionr   r   r   r   r   test_to_json_compression_   s    rA   c              	   C   s4   t ddgi}t }|j|| d W 5 Q R X d S )Nr   r.   r
   )r   r   r   r   )r   expectedbufferr   r   r   test_to_json_compression_modet   s    rD   )r   r   r7   Zpandas.util._test_decoratorsutilZ_test_decoratorstdZpandasr   Zpandas._testingZ_testingr   Zpandas.io.commoncommonr?   r   r"   Zskip_if_not_us_localemarkZ
single_cpur)   r-   r2   r;   r<   ZparametrizerA   rD   r   r   r   r   <module>   s$   
	