U
    ~fh                     @  s   d dl mZ d dlZd dlmZmZmZmZ d dlm	Z	 d dl
mZ dddhZe	je	jhZee d	d
ddZd	d
ddZd	d
ddZddddddZddddddZG dd dZG dd dZG d d! d!ZG d"d# d#Zd$dd$d%d&d'ZdS )(    )annotationsN)AnyIterableOptionalUnion)HelloCompat)_SENSITIVE_COMMANDSsnappyzlibzstdbool)returnc                  C  s*   zdd l } W dS  tk
r$   Y dS X d S Nr   TF)r	   ImportError)r	    r   ?/tmp/pip-unpacked-wheel-36gvocj8/pymongo/compression_support.py_have_snappy   s
    r   c                  C  s*   zdd l } W dS  tk
r$   Y dS X d S r   )r
   r   )r
   r   r   r   
_have_zlib$   s
    r   c                  C  s*   zdd l } W dS  tk
r$   Y dS X d S r   )	zstandardr   )r   r   r   r   
_have_zstd-   s
    r   r   zUnion[str, Iterable[str]]	list[str])dummyvaluer   c                 C  s   z| d}W n tk
r*   t|}Y nX |d d  D ]}|tkrd|| tjd| dd q8|dkrt s|| tjddd q8|dkrt s|| tjddd q8|d	kr8t	 s8|| tjd
dd q8|S )N,zUnsupported compressor:    )
stacklevelr	   zuWire protocol compression with snappy is not available. You must install the python-snappy module for snappy support.r
   zWWire protocol compression with zlib is not available. The zlib module is not available.r   zwWire protocol compression with zstandard is not available. You must install the zstandard module for zstandard support.)
splitAttributeErrorlist_SUPPORTED_COMPRESSORSremovewarningswarnr   r   r   )r   r   compressorsZ
compressorr   r   r   validate_compressors6   s6    



r$   strint)optionr   r   c                 C  s\   zt |}W n* tk
r6   t|  d|dd Y nX |dk sH|dkrXtd| |f |S )Nz must be an integer, not .	   z$%s must be between -1 and 9, not %d.)r&   	Exception	TypeError
ValueError)r'   r   levelr   r   r   validate_zlib_compression_levelZ   s    r/   c                   @  s,   e Zd ZdddddZdddd	d
ZdS )CompressionSettingsr   r&   r#   zlib_compression_levelc                 C  s   || _ || _d S Nr1   )selfr#   r2   r   r   r   __init__e   s    zCompressionSettings.__init__zOptional[list[str]]z4Union[SnappyContext, ZlibContext, ZstdContext, None])r#   r   c                 C  sB   |r>|d }|dkrt  S |dkr,t| jS |dkr:t S d S d S )Nr   r	   r
   r   )SnappyContextZlibContextr2   ZstdContext)r4   r#   Zchosenr   r   r   get_compression_contexti   s    
z+CompressionSettings.get_compression_contextN)__name__
__module____qualname__r5   r9   r   r   r   r   r0   d   s   r0   c                   @  s$   e Zd ZdZedddddZdS )r6      bytesdatar   c                 C  s   dd l }|| S Nr   )r	   compress)r@   r	   r   r   r   rB   {   s    zSnappyContext.compressNr:   r;   r<   compressor_idstaticmethodrB   r   r   r   r   r6   x   s   r6   c                   @  s.   e Zd ZdZddddZddddd	Zd
S )r7   r   r&   r.   c                 C  s
   || _ d S r3   rF   )r4   r.   r   r   r   r5      s    zZlibContext.__init__r>   r?   c                 C  s   dd l }||| jS rA   )r
   rB   r.   )r4   r@   r
   r   r   r   rB      s    zZlibContext.compressN)r:   r;   r<   rD   r5   rB   r   r   r   r   r7      s   r7   c                   @  s$   e Zd ZdZedddddZdS )r8      r>   r?   c                 C  s   dd l }| | S rA   )r   ZZstdCompressorrB   )r@   r   r   r   r   rB      s    zZstdContext.compressNrC   r   r   r   r   r8      s   r8   r>   )r@   rD   r   c                 C  sn   |t jkr dd l}|t| S |tjkr<dd l}|| S |tjkr\dd l	}|
 | S td|f d S )Nr   zUnknown compressorId %d)r6   rD   r	   
uncompressr>   r7   r
   
decompressr8   r   ZZstdDecompressorr-   )r@   rD   r	   r
   r   r   r   r   rI      s    



rI   )
__future__r   r!   typingr   r   r   r   Zpymongo.hellor   Zpymongo.helpers_sharedr   r   ZCMDZ
LEGACY_CMDZ_NO_COMPRESSIONupdater   r   r   r$   r/   r0   r6   r7   r8   rI   r   r   r   r   <module>   s"   

			$

