o
    Eh*                     @  s   d dl mZ d dlm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 G dd dZG dd	 d	eZG d
d deZG dd deZdS )    )annotations)finalN)is_string_dtype)opsc                   @  s   e Zd ZU eZded< eZded< eZded< eZded< dd
dZ	dddZ
dddZedddZedd ZeefdddZedddZdS )BaseOpsUtiltype[Exception] | Noneseries_scalar_excframe_scalar_excseries_array_exc
divmod_excop_namestrreturn4type[Exception] | tuple[type[Exception], ...] | Nonec                 C  sT   |dv r	| j }|S t|tjrt|tjr| j}|S t|tjr%| j}|S | j}|S )N)
__divmod____rdivmod__)r   
isinstancepdSeriesr
   r   r	   )selfr   objotherresult r   ?lib/python3.10/site-packages/pandas/tests/extension/base/ops.py_get_expected_exception   s   z#BaseOpsUtil._get_expected_exceptionc                 C  s   |S Nr   )r   r   r   r   Zpointwise_resultr   r   r   _cast_pointwise_result(   s   z"BaseOpsUtil._cast_pointwise_resultc                 C  s
   t |S r   )tmget_op_from_name)r   r   r   r   r   r   1   s   
zBaseOpsUtil.get_op_from_nameser	pd.Seriesc                 C  s.   |  |||}| |}| ||||| d S r   )r   r   	_check_op)r   r    r   r   excopr   r   r   check_opname:   s   
zBaseOpsUtil.check_opnamec                 C  sP   t |tjr t|jdkrt|jd d df || }|S |||}|S )N   r   )	r   r   	DataFramelencolumnsNotImplementedErroriloccombineZto_frame)r   r   r   r$   expectedr   r   r   _combineB   s   zBaseOpsUtil._combinec                 C  s   |d u r)|||}|  |||}| ||||}t|t|s!J t|| d S t| ||| W d    d S 1 s?w   Y  d S r   )r.   r   r   typer   Zassert_equalpytestraises)r   r    r$   r   r   r#   r   r-   r   r   r   r"   M   s   
"zBaseOpsUtil._check_opc           	      C  s   |t u r| d||}n| d||}|d u rC|||\}}|t u r,|| || }}n	|| || }}t|| t|| d S t| t || W d    d S 1 sYw   Y  d S )Nr   r   )divmodr   r   assert_series_equalr0   r1   )	r   r    r$   r   r#   Z
result_divZ
result_modZexpected_divZexpected_modr   r   r   _check_divmod_op_   s   "zBaseOpsUtil._check_divmod_opN)r   r   r   r   )r   r   )r    r!   r   r   r    r!   )__name__
__module____qualname__	TypeErrorr   __annotations__r	   r
   r   r   r   r   r   r%   r.   r*   r"   r4   r   r   r   r   r      s"   
 


		

r   c                   @  s   e Zd ZU dZeZded< eZded< eZded< eZ	ded< dd Z
d	d
 Zdd Zdd Zdd Zdd Zejdejejejgejddd ejej D dd ZdS )BaseArithmeticOpsTestsa?  
    Various Series and DataFrame arithmetic ops methods.

    Subclasses supporting various ops should set the class variables
    to indicate that they support ops of that kind

    * series_scalar_exc = TypeError
    * frame_scalar_exc = TypeError
    * series_array_exc = TypeError
    * divmod_exc = TypeError
    r   r   r	   r
   r   c                 C  sB   |dkrt |jrtd |}t|}| |||jd  d S )N__rmod__%Skip testing Python string formattingr   )r   dtyper0   skipr   r   r%   r+   r   dataall_arithmetic_operatorsr   r    r   r   r   test_arith_series_with_scalar   s
   

z4BaseArithmeticOpsTests.test_arith_series_with_scalarc                 C  sD   |dkrt |jrtd |}td|i}| |||d  d S )Nr<   r=   Ar   )r   r>   r0   r?   r   r'   r%   )r   rA   rB   r   Zdfr   r   r   test_arith_frame_with_scalar   s
   
z3BaseArithmeticOpsTests.test_arith_frame_with_scalarc              	   C  s6   |}t |}| ||t |jd gt|  d S Nr   )r   r   r%   r+   r(   r@   r   r   r   test_arith_series_with_array   s   
(z3BaseArithmeticOpsTests.test_arith_series_with_arrayc                 C  s,   t |}| |td | dtj| d S )Nr&   r   r   r4   r2   r   Zrdivmod)r   rA   r    r   r   r   test_divmod   s   
z"BaseArithmeticOpsTests.test_divmodc                 C  sJ   t |}| |t| |}| |tj| t |}| |tj| d S r   rH   )r   rA   Zdata_for_twosr    r   r   r   r   test_divmod_series_array   s   

z/BaseArithmeticOpsTests.test_divmod_series_arrayc                 C  s~   t |}| d||}|d ur,t| ||  W d    d S 1 s%w   Y  d S || }t || }t|| d S )N__add__)r   r   r   r0   r1   r   r3   )r   rA   r    r#   r   r-   r   r   r   $test_add_series_with_extension_array   s   


z;BaseArithmeticOpsTests.test_add_series_with_extension_arrayboxr   c                 C  s   g | ]	}| d s|qS )Z__r)
startswith.0xr   r   r   
<listcomp>   s    z!BaseArithmeticOpsTests.<listcomp>c                 C  s4   ||}t ||rt|||}|tu sJ d S d S r   )hasattrgetattrNotImplemented)r   rA   rM   r   r   r   r   r   r   6test_direct_arith_with_ndframe_returns_not_implemented   s   
zMBaseArithmeticOpsTests.test_direct_arith_with_ndframe_returns_not_implementedN)r6   r7   r8   __doc__r9   r   r:   r	   r
   r   rC   rE   rG   rI   rJ   rL   r0   markparametrizer   r   r'   ZIndexr   Zarithmetic_dunder_methodsZcomparison_dunder_methodsrV   r   r   r   r   r;   s   s(   
 		

r;   c                   @  s*   e Zd ZdZdddZdd Zdd	 Zd
S )BaseComparisonOpsTestsz4Various Series and DataFrame comparison ops methods.r    r!   c           	   
   C  s   |j dv r!|||}|||}| |j |||}t|| d S d }z|||}W n ty> } z|}W Y d }~nd }~ww |d u rZ|||}| |j |||}t|| d S tt| ||| W d    d S 1 ssw   Y  d S )N)eqne)	r6   r,   r   r   r3   	Exceptionr0   r1   r/   )	r   r    rA   r$   r   r   r-   r#   errr   r   r   _compare_other   s*   


"z%BaseComparisonOpsTests._compare_otherc                 C  s   t |}| |||d d S rF   )r   r   r_   )r   rA   comparison_opr    r   r   r   test_compare_scalar   s   
z*BaseComparisonOpsTests.test_compare_scalarc                 C  s<   t |}t j|d gt| |jd}| |||| d S )Nr   )r>   )r   r   r(   r>   r_   )r   rA   r`   r    r   r   r   r   test_compare_array   s   
z)BaseComparisonOpsTests.test_compare_arrayNr5   )r6   r7   r8   rW   r_   ra   rb   r   r   r   r   rZ      s
    
rZ   c                   @  s6   e Zd Zdd Zejdejej	ej
gdd ZdS )BaseUnaryOpsTestsc                 C  s   t j|dd}zdd |d d D  W n< tyP   tt |  W d    n1 s.w   Y  tt |  W d    Y d S 1 sHw   Y  Y d S w | }t j| dd}t|| d S )Nname)rd   c                 S  s   g | ]}| qS r   r   rO   r   r   r   rR      s    z1BaseUnaryOpsTests.test_invert.<locals>.<listcomp>
   )r   r   r9   r0   r1   r   r3   )r   rA   r    r   r-   r   r   r   test_invert   s   &
zBaseUnaryOpsTests.test_invertufuncc                 C  s   t jdt jdt jdi| }d }zt|| }W n9 tyP } z-|}tt|t	f || W d    n1 s9w   Y  W Y d }~d S W Y d }~d S d }~ww ||}t
|| d S )N__pos____neg____abs__)nppositivenegativeabsrT   r]   r0   r1   r/   r9   r   Zassert_extension_array_equal)r   rA   rg   attrr#   r   r^   Zaltr   r   r   #test_unary_ufunc_dunder_equivalence  s   
8z5BaseUnaryOpsTests.test_unary_ufunc_dunder_equivalenceN)r6   r7   r8   rf   r0   rX   rY   rk   rl   rm   rn   rp   r   r   r   r   rc      s    rc   )Z
__future__r   typingr   Znumpyrk   r0   Zpandas.core.dtypes.commonr   Zpandasr   Zpandas._testingZ_testingr   Zpandas.corer   r   r;   rZ   rc   r   r   r   r   <module>   s    d_'