U
    f/e#P                     @   s0  d dl Z d dlZd dlZd dlmZ d dlZd dlm	Z
 d dlmZ d dlmZmZ ejddgddd	 Zejd
dgddd ZG dd dZejde je jgdd Zdd Zdd Zejdejejgejded d ddgeddddggdd Zejded d d ged ddgfed d d gdded ddgfed d d gdded ddgfed d d gdded ddgfed d d gdded ddgfgejdejejgd d! Z d"d# Z!d$d% Z"ejd&d
dgd'd( Z#ejd&d ej$gejde j%e j&gd)d* Z'dS )+    N)np_version_under1p20)ops)SparseArraySparseDtypeintegerblock)paramsc                 C   s   | j S )z.kind kwarg to pass to SparseArray/SparseSeriesparamrequest r   O/tmp/pip-unpacked-wheel-tiezk1ph/pandas/tests/arrays/sparse/test_arithmetics.pykind   s    r   TFc                 C   s   | j S Nr	   r   r   r   r   mix   s    r   c                   @   s
  e Zd ZejZeZdd Zdd Z	dd Z
dd Zd	d
 Zejddddgejddddgdd Zdd Z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jdd&d'ejgd(d) Zejdd&d'ejgd*d+ Zd,d- Zd.d/ Zd0d1 ZdS )2TestSparseArrayArithmeticsc                 C   s   t || d S r   )tmassert_numpy_array_equal)selfabr   r   r   _assert"   s    z"TestSparseArrayArithmetics._assertc           
   	   C   s   t jddd |r$||| }n||| }|tjtjfkrR||d |}n
|||}|tjtjfkrt 	|}	|	
 rt j||	< | || W 5 Q R X d S )Nignore)invaliddivide      ?)nperrstateto_denseoperatortruedivr   Zrtruedivfloordiv	rfloordivisinfanynanr   )
r   r   r   a_denseb_denser   opresultexpectedmaskr   r   r   _check_numeric_ops%   s    


z-TestSparseArrayArithmetics._check_numeric_opsc                 C   sF   t || jstt |jts t|jjtjks2tt |jt	sBtd S r   )

isinstance_klassAssertionErrordtyper   subtyper   bool_
fill_valuebool)r   resr   r   r   _check_bool_result:   s    z-TestSparseArrayArithmetics._check_bool_resultc              	   C   s  t jdd | ||k | ||k ||k | ||k | ||k ||k | ||k | ||k ||k | ||k | ||k ||k | ||k | ||k ||k | ||k  | ||k  ||k  | ||k | ||k ||k | ||k | ||k ||k | ||k | ||k ||k | ||k | ||k ||k | ||k | ||k ||k | ||k  | ||k  ||k  W 5 Q R X d S )Nr   )r   )r   r   r7   r   r   r   r   r   r'   r(   r   r   r   _check_comparison_ops@   s2    z0TestSparseArrayArithmetics._check_comparison_opsc                 C   s   |  ||@  | ||@  ||@  |  ||B  | ||B  ||B  |  ||@  | ||@  ||@  |  ||B  | ||B  ||B  d S r   )r7   r   r   r8   r   r   r   _check_logical_opsk   s    z-TestSparseArrayArithmetics._check_logical_opsscalarr         r4   N   c                 C   s   |}t s\|tjtjfkr\|tjkr,|dkr,n0|tjkr@|dkr@ntjjtdd}|j	
| | tjdddtjddddtjg
}	| j|	||d}
| |
||	||| d S )Nr   GH#38172Zraisesreasonr<   r>   r   r4   )r   r    r"   r   r#   pytestmarkxfailr0   node
add_marker_baser   r&   r/   r-   )r   r   r   all_arithmetic_functionsr4   r;   r   r)   rD   valuesr   r   r   r   test_float_scalary   s    $z,TestSparseArrayArithmetics.test_float_scalarc                 C   s   |  tjdddtjddddtjg
}| j||d}| |d|d | |d|d | |d|d | j||dd}| |d|d | |d|d | |d|d | j||dd}| |d|d | |d|d | |d|d d S )Nr<   r>   r   r   r=   rB   rH   r   r&   r/   r9   )r   r   rJ   r   r   r   r   test_float_scalar_comparison   s    $z7TestSparseArrayArithmetics.test_float_scalar_comparisonc           
      C   sx   |}|  ddddddddddg
}|  ddddddddddg
}| j||dd}| j||dd}	| ||	|||| d S )	N        r          @      @      @      @r   rB   )rH   r/   r-   )
r   r   r   rI   r   r)   rJ   rvaluesr   r   r   r   r   "test_float_same_index_without_nans   s    z=TestSparseArrayArithmetics.test_float_same_index_without_nansc                 C   s   |}t s:|tjkr:|r|dks:tjjtdd}|j| | 	t
jdddt
jddddt
jg
}| 	t
jdddt
jddddt
jg
}| j||d	}	| j||d	}
| |	|
|||| d S )
Nr   r?   r@   r<   r>   r   r=      rL   )r   r   r#   rC   rD   rE   r0   rF   rG   rH   r   r&   r/   r-   )r   r   r   rI   r   r)   rD   rJ   rT   r   r   r   r   r   test_float_same_index_with_nans   s     $$z:TestSparseArrayArithmetics.test_float_same_index_with_nansc                 C   s   |  tjdddtjddddtjg
}|  tjdddtjddddtjg
}| j||d}| j||d}| |||| |  ddd	d
dddd	ddg
}|  dd	ddddddd	dg
}| j||dd}| j||dd}| |||| d S )Nr<   r>   r   r=   rV   rL   rO   r   rP   rQ   rR   rS   rB   rM   r   r   rJ   rT   r   r   r   r   r    test_float_same_index_comparison   s    $$z;TestSparseArrayArithmetics.test_float_same_index_comparisonc           	      C   s6  |}|  tjdddtjddddtjg
}|  dtjddtjddddtjg
}| j||d}| j||d}| |||||| | ||d ||d || | j||dd}| j||d}| |||||| | j||dd}| j||dd}| |||||| | j||dd}| j||dd}| |||||| d S Nr<   r>   r   r=      rL   rB   rH   r   r&   r/   r-   )	r   r   r   rI   r)   rJ   rT   r   r   r   r   r   test_float_array   s     $$z+TestSparseArrayArithmetics.test_float_arrayc                 C   s6  |}|  tjdddtjddddtjg
}|  dtjddtjddddtjg
}| j|dd}| j|dd}| |||||| | ||d ||d || | j|ddd	}| j|dd}| |||||| | j|ddd	}| j|ddd	}| |||||| | j|ddd	}| j|ddd	}| |||||| d S )
Nr<   r>   r   r=   r[   r   rL   r   rB   r\   )r   r   rI   r)   rJ   rT   r   r   r   r   r   test_float_array_different_kind   s     $$z:TestSparseArrayArithmetics.test_float_array_different_kindc                 C   s  |  tjdddtjddddtjg
}|  dtjddtjddddtjg
}| j||d}| j||d}| |||| | ||d ||d  | j||dd}| j||d}| |||| | j||dd}| j||dd}| |||| | j||dd}| j||dd}| |||| d S rZ   rM   rX   r   r   r   test_float_array_comparison  s    $$z6TestSparseArrayArithmetics.test_float_array_comparisonc           
      C   s  |}t j}| jddddddddddg
|d}| jddddddddddg
|d}| j|||d}|jt|kspt| j|||d}	|	jt|kst| ||	|||| | ||	d ||d || | j|d||d}|jt|kst| j|||d}	|	jt|ks
t| ||	|||| | j|d||d}|jt|ksDt| j|d||d}	|	jt|ksjt| ||	|||| | j|d||d}|jt|dd	kst| j|d||d}	|	jt|dd	kst| ||	|||| d S )
Nr   r<   r>   r1   r=   r[   r1   r   )r4   r1   r   r4   )r   int64rH   r/   r1   r   r0   r-   )
r   r   r   rI   r)   r1   rJ   rT   r   r   r   r   r   test_int_array  s2    ""z)TestSparseArrayArithmetics.test_int_arrayc                 C   s.  d}| j ddddddddddg
|d}| j ddddddddddg
|d}| j|||d}| j|||d}| |||| | ||d ||d  | j|||dd	}| j|||d}| |||| | j|||dd	}| j|||dd	}| |||| | j|||dd	}| j|||dd	}| |||| d S )
Nrc   r   r<   r>   r`   r=   r[   ra   )r1   r   r4   )rH   r/   r9   )r   r   r1   rJ   rT   r   r   r   r   r   test_int_array_comparison:  s     ""z4TestSparseArrayArithmetics.test_int_array_comparisonTFc                 C   sl   | j ddddgtjd}| j ddddgtjd}| j||tj|d}| j||tj|d}| |||| d S NTFr`   )r   r1   r4   rH   r   r3   r/   r:   r   r   r4   rJ   rT   r   r   r   r   r   test_bool_same_indexR  s
    z/TestSparseArrayArithmetics.test_bool_same_indexc                 C   st   | j ddddddgtjd}| j ddddddgtjd}| j||tj|d}| j||tj|d}| |||| d S rf   rg   rh   r   r   r   test_bool_array_logical]  s
    z2TestSparseArrayArithmetics.test_bool_array_logicalc                 C   s  |}t s8|tjtjfkr8|r8tjjtdd}|j	
| d}| tjdddtjddddtjg
}| jddddddddddg
|d	}	| j||d
}
| j|	|d
}|jt|kst| |
|||	|| | |
|d ||	d || | j||dd}
| j|	|d
}|jt|kst| |
|||	|| | j||dd}
| j|	|dd}|jt|ksZt| |
|||	|| | j||dd}
| j|	|dd}|jt|ddkst| |
|||	|| d S )Nr?   r@   rc   r<   r>   r   r=   r[   r`   rL   rB   rb   )r   r    r"   r   r#   rC   rD   rE   r0   rF   rG   rH   r   r&   r/   r1   r   r-   )r   r   r   rI   r   r)   rD   rdtyperJ   rT   r   r   r   r   r   test_mixed_array_float_inth  s:    $"z5TestSparseArrayArithmetics.test_mixed_array_float_intc                 C   sp  d}|  tjdddtjddddtjg
}| j ddddddddddg
|d}| j||d}| j||d}|jt|ksxt| |||| | ||d ||d  | j||dd	}| j||d}|jt|kst| |||| | j||dd	}| j||dd	}|jt|kst| |||| | j||dd	}| j||dd	}|jt|dd
ks\t| |||| d S )Nrc   r<   r>   r   r=   r[   r`   rL   rB   rb   )rH   r   r&   r/   r1   r   r0   r9   )r   r   rk   rJ   rT   r   r   r   r   r   test_mixed_array_comparison  s(    $"z6TestSparseArrayArithmetics.test_mixed_array_comparisonc                 C   sl   t ddddg}t ddddg}||A }tjjjdtjdddgdd}t dddg|d	}t	|| d S )
NTFrV   r   r<   r>   Zint32r`   )Zsparse_index)
r   pdcoreZarrayssparseZIntIndexr   arrayr   assert_sp_array_equal)r   str*   Zsp_indexr+   r   r   r   test_xor  s    "z#TestSparseArrayArithmetics.test_xor) __name__
__module____qualname__r   rq   rH   r   r/   r   r-   r7   r9   r:   rC   rD   parametrizerK   rN   rU   rW   rY   r]   r^   r_   rd   re   r&   ri   rj   rl   rm   ru   r   r   r   r   r      s4   +$



&r   r)   c                 C   s@   t ddgdd}| |ddg}| |t ddg}t|| d S )Nr   r<   rb   )r   r   rr   )r)   arrr*   r+   r   r   r   test_with_list  s    r{   c                  C   s@   t ddgdd} tddgddgg}| |}|tks<td S )Nr   r<   rb   r>   r=   rV   )r   rn   	DataFrame__add__NotImplementedr0   )rz   Zdfr*   r   r   r   test_with_dataframe  s    
r   c                  C   s6   t ddgdd} | td }| d }t|| d S )Nr   r<   rb   r>   r   r   rq   r   rr   )rz   r*   r+   r   r   r   test_with_zerodim_ndarray  s    r   ufuncrz   r<   c                 C   s8   | |}| |j }t| t||d}t|| d S )Nrb   )r4   r   r   asarrayr   rr   )r   rz   r*   r4   r+   r   r   r   test_ufuncs  s    
r   za, br>   rb   c                 C   sD   | ||}| t |t |}t|ts.ttt || d S r   )r   r   r.   r   r0   r   r   )r   r   r   r*   r+   r   r   r   test_binary_ufuncs  s    
r   c                  C   sL   t ddddg} tddddg}|| 7 }tddddg}t|| d S )Nr   r>   r<   r=   )r   r   rq   r   r   sparrayZndarrayr+   r   r   r   test_ndarray_inplace  s
    r   c                  C   sN   t ddddg} tddddg}| |7 } t ddddgdd}t| | d S )Nr   r>   r<   r=   rb   r   r   r   r   r   test_sparray_inplace  s
    r   r4   c                 C   s   t ddddg}t|| d}| }t| |  d}t|| t| }t|}t|| td|i }td|i}t	|| d S )NTFrb   A)
r   rq   r   r   rr   rn   ZSeriesZassert_series_equalr|   Zassert_frame_equal)r4   rz   r   r*   r+   r   r   r   test_invert  s    
r   c                 C   sL   t ddt jdg}t||d}| |}t| || |d}t|| d S )Nr   r<   r>   rb   )r   rq   r&   r   r   rr   )r)   r4   rz   r   r*   r+   r   r   r   test_unary_op  s
    r   )(r    Znumpyr   rC   Zpandas.compatr   Zpandasrn   Zpandas._testingZ_testingr   Zpandas.corer   Zpandas.core.arrays.sparser   r   Zfixturer   r   r   rD   ry   eqaddr{   r   r   absexpr   rq   Zgreaterr   r   r   r   r&   posnegr   r   r   r   r   <module>   sV   

   
	     

