U
    /eU                     @   sx  d Z ddlmZ ddlZddlZddlmZmZm	Z	 ddl
mZ ddlmZ ddlmZmZmZ d	d
 Zdd ZG dd deZzddlmZ W n: ek
r   zddlZW n ek
r   dZY nX Y nX dd ZdddZeZeZdd ZeZee8 ddlZed ej Zej!Ze	ej ddZej!ZW 5 Q R X ee ddl"m Z ddl"m!Z W 5 Q R X dd Z dd Z!dS )z Store arrays

We put arrays on disk as raw bytes, extending along the first dimension.
Alongside each array x we ensure the value x.dtype which stores the string
description of the array's dtype.
    )suppressN)valmapidentitypartial   )	Interface)File)frame
framesplitsuffixc                 C   s
   | j  S )z Serialize dtype to bytes

    >>> serialize_dtype(np.dtype('i4'))
    b'<i4'
    >>> serialize_dtype(np.dtype('M8[us]'))
    b'<M8[us]'
    )strencode)dt r   //tmp/pip-unpacked-wheel-ziavv1t7/partd/numpy.pyserialize_dtype   s    r   c                 C   s&   |  drtt| S t| S dS )z Parse text as numpy dtype

    >>> parse_dtype(b'i4')
    dtype('int32')

    >>> parse_dtype(b"[('a', 'i4')]")
    dtype([('a', '<i4')])
       [N)
startswithnpdtypeeval)sr   r   r   parse_dtype   s    	
r   c                   @   sb   e Zd ZdddZdd Zdd Zdd	 Zd
d Zdd Zdd Z	dd Z
edd Zdd ZdS )NumpyNc                 C   s*   |rt |trt|}|| _t|  d S N)
isinstancer   r   partdr   __init__)selfr   r   r   r   r   ,   s    zNumpy.__init__c                 C   s
   d| j iS )Nr   )r   r   r   r   r   __getstate__2   s    zNumpy.__getstate__c                 K   sF   |  D ]"\}}| jt|dt|j q| jjtt|f| d S )N.dtype)	itemsr   Zisetr   r   r   appendr   	serialize)r   datakwargskvr   r   r   r#   5   s    zNumpy.appendc                 K   sD   | j j|f|}| j jdd |D dd}tt|}ttt||S )Nc                 S   s   g | ]}t |d qS r!   r   .0keyr   r   r   
<listcomp><   s     zNumpy._get.<locals>.<listcomp>F)lock)r   _getmapr   listdeserialize)r   keysr&   bytesZdtypesr   r   r   r0   :   s    
z
Numpy._getc                 K   s"   dd |D }| j j|f| d S )Nc                 S   s   g | ]}t |d qS r)   r*   r+   r   r   r   r.   B   s     z Numpy.delete.<locals>.<listcomp>)r   delete)r   r4   r&   keys2r   r   r   r6   A   s    zNumpy.deletec                 C   s   | j ||S r   )r   _iset)r   r-   valuer   r   r   r8   E   s    zNumpy._isetc                 C   s
   | j  S r   )r   dropr   r   r   r   r:   H   s    z
Numpy.dropc                 C   s   | j   d S r   )r   __del__r   r   r   r   r;   K   s    zNumpy.__del__c                 C   s   | j jS r   )r   r/   r   r   r   r   r/   N   s    z
Numpy.lockc                 G   s   |    | jj| f|  d S r   )r:   r   __exit__)r   argsr   r   r   r<   R   s    zNumpy.__exit__)N)__name__
__module____qualname__r   r    r#   r0   r6   r8   r:   r;   propertyr/   r<   r   r   r   r   r   +   s   

r   )msgpackFc              
   C   sd   | j dkrX|   }tt  ttj|ddW  5 Q R  S Q R X ttj	|tj
dS |  S d S )NOT)use_bin_type)protocol)r   flattentolistr   	Exceptionr	   rB   packbpickledumpsHIGHEST_PROTOCOLtobytes)xlr   r   r   r$   _   s    

$r$   c                    s   |dkrz6t jdkrddi nddi  fddt| D }W n& tk
rd   d	d t| D }Y nX tjttt|dd
}d}|D ]$}||||t| < |t|7 }q|S t	| |}|r|
 }|S d S )NrC   )r         rawFencodingzutf-8c                    s   g | ]}t j|f qS r   )rB   unpackbr,   fZunpack_kwargsr   r   r.   q   s   zdeserialize.<locals>.<listcomp>c                 S   s   g | ]}t |qS r   )rJ   loadsrU   r   r   r   r.   t   s     )r   r   )rB   versionr
   rH   r   emptysumr1   lenZ
frombuffercopy)r5   r   r]   blocksresultiblockr   rW   r   r3   i   s(    



r3   c                 C   s   | S r   r   )r5   itemsizer   r   r   <lambda>       rc   )Ztypesize)compress)
decompressc                 C   s    |dkrt | S t| |jS d S NrC   )compress_textcompress_bytesrb   r5   r   r   r   r   re      s    re   c                 C   s   |dkrt | S t| S d S rg   )decompress_textdecompress_bytesrj   r   r   r   rf      s    rf   )F)#__doc__
contextlibr   rJ   Znumpyr   Ztoolzr   r   r   corer   filer   utilsr	   r
   r   r   r   r   ZpandasrB   ImportErrorr$   r3   rh   rk   ri   rl   ZbloscZset_nthreadsre   rf   Zsnappyr   r   r   r   <module>   sF   +




