B
    d                 @   sZ   d Z ddlmZ ddlmZ ddlmZ er8ddlmZ G dd de	Z
G dd	 d	Zd
S )zc
The SCEngine class uses the ``sortedcontainers`` package to implement an
Index engine for Tables.
    )OrderedDict)starmap)HAS_SORTEDCONTAINERS)
SortedListc               @   sT   e Zd ZdZdd Zdd Zdd Zdd	 Zd
d Zdd Z	dd Z
dZdd ZdS )Node)keyvaluec             C   s   || _ || _d S )N)r   r   )selfr   r    r
   _/work/yifan.wang/ringdown/master-ringdown-env/lib/python3.7/site-packages/astropy/table/soco.py__init__   s    zNode.__init__c             C   s,   |j tkr"| j| jf|j|jfk S | j|k S )N)	__class__r   r   r   )r	   otherr
   r
   r   __lt__   s    
zNode.__lt__c             C   s,   |j tkr"| j| jf|j|jfkS | j|kS )N)r   r   r   r   )r	   r   r
   r
   r   __le__   s    
zNode.__le__c             C   s,   |j tkr"| j| jf|j|jfkS | j|kS )N)r   r   r   r   )r	   r   r
   r
   r   __eq__!   s    
zNode.__eq__c             C   s,   |j tkr"| j| jf|j|jfkS | j|kS )N)r   r   r   r   )r	   r   r
   r
   r   __ne__&   s    
zNode.__ne__c             C   s,   |j tkr"| j| jf|j|jfkS | j|kS )N)r   r   r   r   )r	   r   r
   r
   r   __gt__+   s    
zNode.__gt__c             C   s,   |j tkr"| j| jf|j|jfkS | j|kS )N)r   r   r   r   )r	   r   r
   r
   r   __ge__0   s    
zNode.__ge__Nc             C   s   d| j d| jdS )NzNode(z, ))r   r   )r	   r
   r
   r   __repr__7   s    zNode.__repr__)__name__
__module____qualname__	__slots__r   r   r   r   r   r   r   __hash__r   r
   r
   r
   r   r      s   r   c               @   sv   e Zd ZdZdddZdd Zdd Zdd
dZdd Zdd Z	dd Z
dd Zdd ZdddZdd Zdd Zd	S ) SCEngineat  
    Fast tree-based implementation for indexing, using the
    ``sortedcontainers`` package.

    Parameters
    ----------
    data : Table
        Sorted columns of the original table
    row_index : Column object
        Row numbers corresponding to data columns
    unique : bool
        Whether the values of the index must be unique.
        Defaults to False.
    Fc             C   s*   t t|}tttt||| _|| _d S )N)maptupler   r   r   zip_nodes_unique)r	   dataZ	row_indexuniqueZ	node_keysr
   r
   r   r   K   s    
zSCEngine.__init__c             C   s<   | j r&|| jkr&d|dd}t|| jt|| dS )z(
        Add a key, value pair.
        z
duplicate z!rz in unique indexN)r!   r    
ValueErroraddr   )r	   r   r   messager
   r
   r   r%   P   s    zSCEngine.addc             C   s   dd | j ||D S )z;
        Find rows corresponding to the given key.
        c             S   s   g | ]
}|j qS r
   )r   ).0noder
   r
   r   
<listcomp>]   s    z!SCEngine.find.<locals>.<listcomp>)r    irange)r	   r   r
   r
   r   findY   s    zSCEngine.findNc             C   sp   |dk	r<t ||}y| j| W n tk
r6   dS X dS t| j||}x|D ]}| j| qTW t|S )z1
        Remove data from the given key.
        NFT)r   r    remover$   listr*   bool)r	   r   r"   itemitemsr
   r
   r   r,   _   s    

zSCEngine.removec             C   s,   x&| j D ]}|j|kr| jd8  _qW dS )z;
        Decrement rows larger than the given row.
           N)r    r   )r	   rowr(   r
   r
   r   
shift_lefto   s    
zSCEngine.shift_leftc             C   s,   x&| j D ]}|j|kr| jd7  _qW dS )zH
        Increment rows greater than or equal to the given row.
        r1   N)r    r   )r	   r2   r(   r
   r
   r   shift_rightw   s    
zSCEngine.shift_rightc             C   sJ   t  }x:| jD ]0}|j|kr0||j |j q|jg||j< qW | S )z4
        Return a list of key, data tuples.
        )r   r    r   appendr   r0   )r	   resultr(   r
   r
   r   r0      s    
zSCEngine.itemsc             C   s"   xt | jD ]\}}||_qW dS )z6
        Make row order align with key order.
        N)	enumerater    r   )r	   indexr(   r
   r
   r   sort   s    zSCEngine.sortc             C   s   dd | j D S )z?
        Return a list of rows in order sorted by key.
        c             S   s   g | ]
}|j qS r
   )r   )r'   r(   r
   r
   r   r)      s    z(SCEngine.sorted_data.<locals>.<listcomp>)r    )r	   r
   r
   r   sorted_data   s    zSCEngine.sorted_dataTTc             C   s   | j |||}dd |D S )z7
        Return row values in the given range.
        c             S   s   g | ]
}|j qS r
   )r   )r'   r(   r
   r
   r   r)      s    z"SCEngine.range.<locals>.<listcomp>)r    r*   )r	   lowerupperZboundsiteratorr
   r
   r   range   s    zSCEngine.rangec                sH    fdd| j D }x|D ]} |j |_qW | j   | j | dS )z:
        Replace rows with the values in row_map.
        c                s   g | ]}|j  kr|qS r
   )r   )r'   r(   )row_mapr
   r   r)      s    z)SCEngine.replace_rows.<locals>.<listcomp>N)r    r   clearupdate)r	   r@   nodesr(   r
   )r@   r   replace_rows   s
    

zSCEngine.replace_rowsc             C   sj   t | jdkr:t| jd d dg t| jdd   }n| j}ddd |D }d| jj d	| d
S )N      z...z, c             s   s   | ]}t |V  qd S )N)str)r'   r(   r
   r
   r   	<genexpr>   s    z$SCEngine.__repr__.<locals>.<genexpr><z nodes=>)lenr    r-   joinr   r   )r	   rC   Z	nodes_strr
   r
   r   r      s
    ,zSCEngine.__repr__)F)N)r;   )r   r   r   __doc__r   r%   r+   r,   r3   r4   r0   r9   r:   r?   rD   r   r
   r
   r
   r   r   ;   s   
	


r   N)rN   collectionsr   	itertoolsr   Z"astropy.utils.compat.optional_depsr   Zsortedcontainersr   objectr   r   r
   r
   r
   r   <module>   s   +