a nEb<. @sdZddlZddlZddlmZddlZddlZddlZddl Z ddl Z ddl m Z ddlmmZddlZddlmZe jje ddZGdddejjZGd d d eZGd d d eZGd ddeZGdddeZ GdddeZ!GdddeZ"GdddeZ#GdddeZ$GdddejjZ%ddZ&ddZ'e j(dd Z)e jj*d!eej+dfeej,dfe!ej-d"fe j.e"ej-d#ej/gd$e#ej0dfe$ej1dfeej+dfe ej,dfgd gd%d&d'Z2e jj*d!eej+dfeej,dfe!ej-d"fe j.e"ej-d#ej/gd$e#ej0dfe$ej1dfeej+dfe ej,dfgd gd%d(d)Z3e jj*d*e%ej+fgd gd%d+d,Z4e j*d-d"d#gd.d/Z5dS)0z9 Tests for the pandas custom headers in http(s) requests N)BytesIO)is_ci_environmentzThis test can hang in our CI min_versions build and leads to '##[error]The runner has received a shutdown signal...' in GHA. GH 45651)reasonc@s(eZdZdZddZddZddZdS) BaseUserAgentResponderz Base class for setting up a server that can be set up to respond with a particular file format with accompanying content-type headers. The interfaces on the different io methods are different enough that this seemed logical to do. cCs,|d|jd|_td|jgi}|S)z< shared logic at the start of a GET request User-Agentheader) send_responseheadersZrequested_from_user_agentpd DataFrame)self response_dfr_/home/ja/django-apps/lartica_env/lib/python3.9/site-packages/pandas/tests/io/test_user_agent.pystart_processing_headers$s  z/BaseUserAgentResponder.start_processing_headersc CsjtP}tj|dd}||Wdn1s60Y|}Wdn1s\0Y|S)zQ some web servers will send back gzipped files to save bandwidth w)fileobjmodeN)rgzipGzipFilewritegetvalue)r response_bytesbioZzipperrrr gzip_bytes1s (&z!BaseUserAgentResponder.gzip_bytescCs|j|dS)z: shared logic at the end of a GET request N)wfiler)r rrrrwrite_back_bytes;sz'BaseUserAgentResponder.write_back_bytesN)__name__ __module__ __qualname____doc__rrrrrrrrs  rc@seZdZddZdS)CSVUserAgentRespondercCs<|}|dd||jddd}||dS)N Content-Typetext/csvFindexutf-8)r send_header end_headersto_csvencoderr rrrrrdo_GETCs  zCSVUserAgentResponder.do_GETNrrr r-rrrrr"Bsr"c@seZdZddZdS)GzippedCSVUserAgentRespondercCsR|}|dd|dd||jddd}||}||dS)Nr#r$Content-EncodingrFr%r')rr(r)r*r+rrr,rrrr-Ns   z#GzippedCSVUserAgentResponder.do_GETNr.rrrrr/Msr/c@seZdZddZdS)JSONUserAgentRespondercCs8|}|dd||d}||dS)Nr#application/jsonr')rr(r)to_jsonr+rr,rrrr-[s  zJSONUserAgentResponder.do_GETNr.rrrrr1Zsr1c@seZdZddZdS)GzippedJSONUserAgentRespondercCsN|}|dd|dd||d}||}||dS)Nr#r2r0rr')rr(r)r3r+rrr,rrrr-fs   z$GzippedJSONUserAgentResponder.do_GETNr.rrrrr4esr4c@seZdZddZdS) ParquetPyArrowUserAgentRespondercCs8|}|dd||jddd}||dS)Nr#application/octet-streamFpyarrow)r&engine)rr(r) to_parquetrr,rrrr-ss  z'ParquetPyArrowUserAgentResponder.do_GETNr.rrrrr5rsr5c@seZdZddZdS)$ParquetFastParquetUserAgentRespondercCsx|}|dd|ddl}|jddddd|dd}|}Wdn1s`0Y||dS) Nr#r6rz'memory://fastparquet_user_agent.parquetF fastparquet)r&r8 compressionrb)rr(r)fsspecr9openreadr)r rr>frrrrr-~s &z+ParquetFastParquetUserAgentResponder.do_GETNr.rrrrr:}sr:c@seZdZddZdS)PickleUserAgentRespondercCsB|}|dd|t}|||}||dS)Nr#r6)rr(r)rZ to_picklerrr rrrrrrr-s  zPickleUserAgentResponder.do_GETNr.rrrrrBsrBc@seZdZddZdS)StataUserAgentRespondercCsF|}|dd|t}|j|dd|}||dS)Nr#r6F)Z write_index)rr(r)rZto_statarrrCrrrr-s zStataUserAgentResponder.do_GETNr.rrrrrDsrDc@seZdZdZddZdS)AllHeaderCSVResponderz? Send all request headers back for checking round trip cCsPt|j}|d|dd||jddd}|j |dS)Nrr#r$Fr%r') r r r itemsr r(r)r*r+rrr,rrrr-s   zAllHeaderCSVResponder.do_GETN)rrr r!r-rrrrrEsrEcsfdd}|S)Ncs8z|i|WStjjy0tdYq0qdS)N皙?)urlliberrorURLErrortimesleep)argskwargsfuncrrinnerszwait_until_ready..innerr)rPrMrNrQrrOrwait_until_readys rRcCsFtjd|f|}|Wdn1s00Y|dS)N localhost)httpserverZ HTTPServerhandle_request server_close) responderportrUrrrprocess_servers&rZccst&}|d|d}Wdn1s40Ytjt|j|fd}||V|d| d}d}| r||kr| qqz|d7}t dqz|dS) z Fixture that starts a local http server in a separate process on localhost and returns the port. Running in a separate process instead of a thread to allow termination/killing of http server upon cleanup. )rSrN)targetrM rrG)socketbind getsocknamemultiprocessingProcessrZparamstartjoin terminateis_alivekillrKrLclose)requestsockrYZserver_processZ kill_timeZ wait_timerrrrXs&  *   rXz&responder, read_method, parquet_enginer7r;)Zmarks)ZindirectcCsd|dur$t||dkr$tdt|}|durD|d|}n|d||d}|jr`JdS)Nr;r>http://localhost:)r8)pytest importorskiprRempty)rX read_methodparquet_enginedf_httprrrtest_server_and_default_headerss  rtcCs|dur$t||dkr$tdd}td|gi}t|}|dur`|d|d|id}n|d|d|i|d}t||dS) Nr;r>Super Cool Onerrmrstorage_optionsrwr8)rnror r rRtmassert_frame_equal)rXrqrrcustom_user_agentdf_truersrrrtest_server_and_custom_headerss$  r}zresponder, read_methodcsd}d}||dt|}|d|d}||d}|dg}|ddg}t}t|fdd |Dd }|dg}|jd gd d }t ||dS)NruzSuper Secret OnerZAuthrmrv01csg|] }|qSrr).0krvrr _z6test_server_and_all_custom_headers..)rrr&r[)Zaxis) rRisinkeysZ sort_valuesZ reset_indexlistr r Zdropryrz)rXrqr{Zcustom_auth_tokenrsrr|rrvr"test_server_and_all_custom_headersFs$   rr8cCsjddd}t|tddgi}d}tjt|d |jd||d Wdn1s\0YdS) NZcustomZ other_customr~Z column_nameZ column_valuezxstorage_options passed with file object or non-fsspec file path|storage_options passed with buffer, or non-supported URL)matchz/tmp/junk.parquetrx)rnror r Zraises ValueErrorr9)r8r Ztrue_dfmsgrrr,test_to_parquet_to_disk_with_storage_optionsfs  r)6r!rZ http.serverrTiorrbr_rK urllib.errorrHrnZ pandas.compatrZpandas.util._test_decoratorsutilZ_test_decoratorstdZpandasr Zpandas._testingZ_testingrymarkZskipifZ pytestmarkrUZBaseHTTPRequestHandlerrr"r/r1r4r5r:rBrDrErRrZZfixturerXZ parametrizeZread_csv read_jsonZ read_parquetrdZ&skip_array_manager_not_yet_implementedZ read_pickleZ read_statartr}rrrrrrs   &