U
    f/ef6                     @  s\  d Z ddlm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mZ ddlmZ dd	lmZmZmZ dd
lmZmZ ddlmZ ddlmZmZmZmZm Z  ddl!m"Z" ddl#m$Z$ ddl%m&Z& ddl'm(Z( edddZ)d*ddddddZ*d+ddddddZ+dddd Z,ddd!d"Z-d#d$ Z.d%d& Z/G d'd( d(e&Z0G d)d de0Z1dS ),z?
Shared methods for Index subclasses backed by ExtensionArray.
    )annotations)HashableTypeVarN)	ArrayLike)function)cache_readonlydoc)rewrite_exception)is_dtype_equalis_object_dtypepandas_dtype)ABCDataFrame	ABCSeries)validate_putmask)CategoricalDatetimeArrayIntervalArrayPeriodArrayTimedeltaArray)NDArrayBackedExtensionArray)deprecate_ndim_indexing)Index)get_op_result_name_TNDArrayBackedExtensionIndex)boundFstrbool)namecachewrapc                   s   t | t ts"t jdkr~|rJfdd}|_ j|_t|}qfdd}fdd}|_ j|_t||}n*t s }n fdd	}|_ j|_|S )
a  
    Make an alias for a method of the underlying ExtensionArray.

    Parameters
    ----------
    name : str
        Name of an attribute the class should inherit from its EA parent.
    delegate : class
    cache : bool, default False
        Whether to convert wrapped properties into cache_readonly
    wrap : bool, default False
        Whether to wrap the inherited result in an Index.

    Returns
    -------
    attribute, method, property, or cache_readonly
    Zgetset_descriptorc                   s   t | j S N)getattr_dataselfr    A/tmp/pip-unpacked-wheel-tiezk1ph/pandas/core/indexes/extension.pycachedG   s    z!inherit_from_data.<locals>.cachedc                   sZ   t | j }rVt|t| jr4t| j|| jdS t|trH|| S t|| jdS |S Nr&   )	r"   r#   
isinstancetype_simple_newr   r   	set_indexr   r%   result)r   r    r'   r(   fgetP   s    

zinherit_from_data.<locals>.fgetc                   s   t | j | d S r!   )setattrr#   r%   valuer&   r'   r(   fsetZ   s    zinherit_from_data.<locals>.fsetc                   s`    | j f||}r\t|t| j r:t| j|| jdS t|trN|| S t|| jdS |S r*   )r#   r+   r,   r-   r   r   r.   r   )r%   argskwargsr0   )attrr    r'   r(   methodh   s    

z!inherit_from_data.<locals>.method)r"   r+   propertyr,   __name____doc__r   callable)r   delegater   r    r)   r9   r1   r5   r'   )r8   r   r    r(   inherit_from_data/   s$    



r?   z	list[str])namesr   r    c                   s    fdd}|S )a  
    Class decorator to pin attributes from an ExtensionArray to a Index subclass.

    Parameters
    ----------
    names : List[str]
    delegate : class
    cache : bool, default False
    wrap : bool, default False
        Whether to wrap the inherited result in an Index.
    c                   s*   D ] }t | d}t| || q| S )N)r   r    )r?   r2   )clsr   methr   r>   r@   r    r'   r(   wrapper   s    zinherit_names.<locals>.wrapperr'   )r@   r>   r   r    rD   r'   rC   r(   inherit_namesw   s    rE   opnamec                   s    fdd} |_ |S )zC
    Create a comparison method that dispatches to ``._data``.
    c                   s,   t |tr|j}t|}t| j }||S r!   )r+   r   Z_values_maybe_unwrap_indexr"   r#   )r%   otheroprF   r'   r(   rD      s
    
z,_make_wrapped_comparison_op.<locals>.wrapperr;   )rG   rD   r'   rF   r(   _make_wrapped_comparison_op   s    rL   c                   s    fdd} |_ |S )Nc                   sH   t |tr$t|jr$t|tk	r$tS t| j }|t|}t	| ||S r!   )
r+   r   r   dtyper,   NotImplementedr"   r#   rH   _wrap_arithmetic_op)r%   rI   rB   r0   rF   r'   r(   r9      s    
z%make_wrapped_arith_op.<locals>.methodrK   )rG   r9   r'   rF   r(   make_wrapped_arith_op   s    rP   c                 C  sl   |t krt S t|trFt|dks&tt| ||d t| ||d fS t|tsXt|}t| |}||_|S )N   r      )	rN   r+   tuplelenAssertionErrorrO   r   r   r   )r%   rI   r0   Zres_namer'   r'   r(   rO      s    


rO   c                 C  s   t | tr| jS | S )aF  
    If operating against another Index object, we need to unwrap the underlying
    data before deferring to the DatetimeArray/TimedeltaArray/PeriodArray
    implementation, otherwise we will incorrectly return NotImplemented.

    Parameters
    ----------
    obj : object

    Returns
    -------
    unwrapped object
    )r+   r   r#   )objr'   r'   r(   rH      s    
rH   c                   @  sJ  e Zd ZU dZded< ded< ed:dddd	d
ZedZedZ	edZ
edZedZedZeddddZdd Zd;ddddZddddZddddZdd d!d"d#Zd$d% Zd<d&d'Zd(dd)d*d+Zd,d- Zd.d/ Zeejd=d0d1Zeejd>ddd3d4d5Zeddd6d7Zeej ddd8d9Z dS )?ExtensionIndexz>
    Index subclass for indexes backed by ExtensionArray.
    z+IntervalArray | NDArrayBackedExtensionArrayr#   zhtype[Categorical] | type[DatetimeArray] | type[TimedeltaArray] | type[PeriodArray] | type[IntervalArray]	_data_clsNr   )arrayr   c                 C  s@   t || jstt|t| }||_||_i |_|	  |S )z
        Construct from an ExtensionArray of the appropriate type.

        Parameters
        ----------
        array : ExtensionArray
        name : Label, default None
            Attached as result.name
        )
r+   rX   rU   r,   object__new__r#   _name_cacheZ_reset_identity)rA   rY   r   r0   r'   r'   r(   r-      s    
zExtensionIndex._simple_new__eq____ne____lt____gt____le____ge__r   returnc                 C  s   dS )NTr'   r$   r'   r'   r(   _has_complex_internals  s    z%ExtensionIndex._has_complex_internalsc                 C  sH   | j | }t|t| j r<|jdkr6t| || jdS |j}t| |S )NrR   r&   )r#   r+   r,   ndimr\   _ndarrayr   )r%   keyr0   r'   r'   r(   __getitem__  s    

zExtensionIndex.__getitem__left
np.ndarrayc                 C  s   | j j|||dS )N)sidesorter)r#   searchsorted)r%   r4   rm   rn   r'   r'   r(   ro   #  s    zExtensionIndex.searchsortedr   c              	   C  s   t | j|\}}|r|  S z| | W n4 ttfk
r^   | |}| ||| Y S X | j }||| t	| j
|| jdS r*   )r   r#   copy_validate_fill_value
ValueError	TypeError_find_common_type_compatastypeputmaskr,   r-   r   )r%   maskr4   ZnooprM   arrr'   r'   r(   rv   '  s    

zExtensionIndex.putmaskc                 C  s   t | jS r!   )npZasarrayr#   r$   r'   r'   r(   _get_engine_target8  s    z!ExtensionIndex._get_engine_targetr   r0   re   c                 C  s   t | jj|| jdS )N)rM   )r,   r#   Z_from_sequencerM   r/   r'   r'   r(   _from_join_target;  s    z ExtensionIndex._from_join_targetc                 C  s    | j |}t| j|| jdS )z|
        Make new Index with passed location(-s) deleted

        Returns
        -------
        new_index : Index
        r&   )r#   deleter,   r-   r   )r%   locrx   r'   r'   r(   r}   @  s    zExtensionIndex.deletec                 C  s4   t dd|i | jj||d}t| j|| jdS )Nr'   axis)r   r&   )nvZvalidate_repeatr#   repeatr,   r-   r   )r%   Zrepeatsr   r0   r'   r'   r(   r   K  s    zExtensionIndex.repeatint)r~   re   c              	   C  s`   z| j ||}W n4 ttfk
rF   | |}| ||| Y S X t| j|| jdS dS )a  
        Make new Index inserting new item at location. Follows
        Python list.append semantics for negative values.

        Parameters
        ----------
        loc : int
        item : object

        Returns
        -------
        new_index : Index
        r&   N)	r#   insertrr   rs   rt   ru   r,   r-   r   )r%   r~   itemr0   rM   r'   r'   r(   r   P  s    
zExtensionIndex.insertc                 C  s   | j |S )zE
        Convert value to be insertable to underlying array.
        )r#   Z_validate_setitem_valuer3   r'   r'   r(   rq   i  s    z#ExtensionIndex._validate_fill_valuec                 C  s(   | j r
| S | j }t| j|| jdS r*   )Z	is_uniquer#   uniquer,   r-   r   r/   r'   r'   r(   _get_unique_indexo  s    
z ExtensionIndex._get_unique_indexc                 C  s^   z4|| }t |tjrt|}t |ts0td|W S  tk
rX   | t| Y S X d S )Nz,The map function must return an Index object)	r+   ry   Zndarrayr   rs   	Exceptionru   rZ   map)r%   ZmapperZ	na_actionr0   r'   r'   r(   r   v  s    
zExtensionIndex.mapT)rp   re   c              	   C  s   t |}t| j|r$|s| S |  S t| jtjrft|tjrf|jdkrf|dkrftdt| j	 dt
t| jj	t| j	 | jj||d}W 5 Q R X t||j| jddS )NMzM8[ns]zCannot cast z	 to dtype)rp   F)rM   r   rp   )r   r
   rM   rp   r+   ry   kindrs   r,   r;   r	   r#   ru   r   r   )r%   rM   rp   Z
new_valuesr'   r'   r(   ru     s"    
zExtensionIndex.astypec                 C  s
   | j  S r!   )r#   Zisnar$   r'   r'   r(   _isnan  s    zExtensionIndex._isnanc                 C  s.   |  |rdS t|t| s dS | j|jS )NTF)is_r+   r,   r#   equals)r%   rI   r'   r'   r(   r     s
    
zExtensionIndex.equals)N)rk   N)N)N)T)!r;   
__module____qualname__r<   __annotations__classmethodr-   rL   r^   r_   r`   ra   rb   rc   r:   rf   rj   ro   rv   rz   r|   r}   r   r   rq   r   r   r   r   ru   r   r   r   r'   r'   r'   r(   rW      s@   
 
rW   c                      sV   e Zd ZU dZded< edddd fddZd	d
ddZd	ddddZ  Z	S )r   zK
    Index subclass for indexes backed by NDArrayBackedExtensionArray.
    r   r#   Nr   )valuesr   c                   s   t  ||}|j|_|S r!   )superr-   rh   Z_index_data)rA   r   r   r0   	__class__r'   r(   r-     s    z'NDArrayBackedExtensionIndex._simple_newrl   rd   c                 C  s   | j jS r!   )r#   rh   r$   r'   r'   r(   rz     s    z.NDArrayBackedExtensionIndex._get_engine_targetr   r{   c                 C  s    |j | jjj kst| j|S r!   )rM   r#   rh   rU   Z_from_backing_datar/   r'   r'   r(   r|     s    z-NDArrayBackedExtensionIndex._from_join_target)N)
r;   r   r   r<   r   r   r-   rz   r|   __classcell__r'   r'   r   r(   r     s   
 )FF)FF)2r<   
__future__r   typingr   r   Znumpyry   Zpandas._typingr   Zpandas.compat.numpyr   r   Zpandas.util._decoratorsr   r   Zpandas.util._exceptionsr	   Zpandas.core.dtypes.commonr
   r   r   Zpandas.core.dtypes.genericr   r   Zpandas.core.array_algos.putmaskr   Zpandas.core.arraysr   r   r   r   r   Zpandas.core.arrays._mixinsr   Zpandas.core.indexersr   Zpandas.core.indexes.baser   Zpandas.core.opsr   r   r?   rE   rL   rP   rO   rH   rW   r   r'   r'   r'   r(   <module>   s2   H X