o
    NiN                     @   s   d dl Z d dlmZmZ d dlZd dlmZ d dlmZ d dl	m
Z
 ddlmZ dZG d	d
 d
ejjZddeeeef  dejfddZG dd deZdS )    N)DictOptional)HTTPAdapter)MaxRetryError)Retry   )SchemeClient)i  i  i  c                   @   s   e Zd ZdZdS )RecoverableServerErrorzB
    Server returned one of `RECOVERABLE_SERVER_ERROR_CODES`.
    N)__name__
__module____qualname____doc__ r   r   L/home/ubuntu/.local/lib/python3.10/site-packages/cached_path/schemes/http.pyr	      s    r	   headersreturnc                 C   sN   t  }tddtd}|dt|d |dt|d | r%|j|  |S )a  
    We ran into an issue where http requests to s3 were timing out,
    possibly because we were making too many requests too quickly.
    This helper function returns a requests session that has retry-with-backoff
    built in. See
    <https://stackoverflow.com/questions/23267409/how-to-implement-retry-mechanism-into-python-requests-library>.

    Parameters
    ----------
    headers : Optional[Dict[str, str]], optional
        Custom headers to add to all requests, by default None
       r   )totalbackoff_factorstatus_forcelistzhttp://)max_retrieszhttps://)requestsSessionr   RECOVERABLE_SERVER_ERROR_CODESmountr   r   update)r   sessionretriesr   r   r   session_with_backoff   s   r   c                       s   e Zd ZdZejef Zddedee	eef  ddf fddZ
edd	 Zdee fd
dZdee fddZdejddfddZdd Z  ZS )
HttpClient)httphttpsNresourcer   r   c                    s    t  | d| _|pi | _dS )ar  
        Initialize an HTTP client for the given resource.

        Parameters
        ----------
        resource : str
            The URL to the resource.
        headers : Optional[Dict[str, str]], optional
            Custom headers to add to all requests, by default None.
            Example: {"Authorization": "Bearer YOUR_TOKEN"} for private resources.
        N)super__init___head_responser   )selfr"   r   	__class__r   r   r$   1   s   zHttpClient.__init__c              
   C   s   | j d u rAzt| j}|j| jdd}W d    n1 sw   Y  W n ty5 } zt|jd }~ww | | || _ | j S | j S )NT)allow_redirects)	r%   r   r   headr"   r   r	   reasonvalidate_response)r&   r   responseer   r   r   head_responseA   s   


zHttpClient.head_responsec                 C   s   | j jdS )NETag)r/   r   get)r&   r   r   r   get_etagO   s   zHttpClient.get_etagc                 C   s"   | j jd}|d u rd S t|S )NzContent-Length)r/   r   r1   int)r&   content_lengthr   r   r   get_sizeR   s   zHttpClient.get_size	temp_filec                 C   s   t | j:}z
|j| jdd}W n ty! } zt|jd }~ww | | |jddD ]	}|r6|	| q-W d    d S 1 sBw   Y  d S )NT)streami   )
chunk_size)
r   r   r1   r"   r   r	   r+   r,   iter_contentwrite)r&   r6   r   r-   r.   chunkr   r   r   get_resourceV   s   


"zHttpClient.get_resourcec                 C   s4   |j dkr
t| j|j tv rt|d|  d S )Ni  )r-   )status_codeFileNotFoundErrorr"   r   r	   raise_for_status)r&   r-   r   r   r   r,   s   s
   



zHttpClient.validate_responseN)r
   r   r   schemer   recoverable_errorsr	   strr   r   r$   propertyr/   r2   r3   r5   ioBufferedWriterr<   r,   __classcell__r   r   r'   r   r   -   s    (
r   r@   )rE   typingr   r   r   requests.adaptersr   urllib3.exceptionsr   urllib3.util.retryr   scheme_clientr   r   
exceptions	HTTPErrorr	   rC   r   r   r   r   r   r   r   <module>   s    "