o
    LhZ                     @  s  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
mZmZmZ d dlmZmZ d dlmZmZ d dlmZmZ er~d dlmZ d d	lmZmZ d d
lmZmZ d dlm Z 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+m,Z,m-Z- edZ.dQddZ/dRddZ0dSddZ1dTd d!Z2dUd)d*Z3dVd1d2Z4G d3d4 d4eZ5dWd6d7Z6G d8d9 d9eZ7G d:d; d;Z8dXd@dAZ9dYdFdGZ:dZdIdJZ;d[dNdOZ<dPS )\    )annotations)Enumauto)chain)TYPE_CHECKINGAnyCallableLiteralTypeVarcast)is_compliant_expr
zip_strict)is_narwhals_seriesis_numpy_array)InvalidOperationErrorMultiOutputExpressionError)Sequence)NeverTypeIs)CompliantExprCompliantFrameT)
AliasNamesCompliantExprAnyCompliantFrameAnyCompliantNamespaceAnyEagerNamespaceAny	EvalNamesExprSeries)IntoExprNonNestedLiteral_1DArrayTobjr   returnTypeIs[Expr]c                 C     ddl m} t| |S )'Check whether `obj` is a Narwhals Expr.r   r   )narwhals.exprr   
isinstance)r%   r    r,   ]/var/www/html/Persson_Maskin/env/lib/python3.10/site-packages/narwhals/_expression_parsing.pyis_expr$      
r.   TypeIs[Series[Any]]c                 C  r(   )r)   r   r   )narwhals.seriesr    r+   )r%   r    r,   r,   r-   	is_series+   r/   r2   exprs#CompliantExpr[CompliantFrameT, Any]EvalNames[CompliantFrameT]c                    s:   t  d sdt d  d}t|d
 fdd	}|S )Nr   z3Safety assertion failed, expected expression, got: z. Please report a bug.dfr   r&   Sequence[str]c                       d  | d d S Nr      )_evaluate_output_names)r6   r3   r,   r-   evaluate_output_names;      z<combine_evaluate_output_names.<locals>.evaluate_output_names)r6   r   r&   r7   )r   typeAssertionError)r3   msgr=   r,   r<   r-   combine_evaluate_output_names2   s
   rB   r   AliasNames | Nonec                    s$    d j d u r	d S d fdd}|S )Nr   namesr7   r&   c                   r8   r9   _alias_output_names)rD   r<   r,   r-   alias_output_namesG   r>   z6combine_alias_output_names.<locals>.alias_output_names)rD   r7   r&   r7   rE   )r3   rG   r,   r<   r-   combine_alias_output_namesA   s   rH   plxr   other&IntoExpr | NonNestedLiteral | _1DArray
str_as_litbool#CompliantExprAny | NonNestedLiteralc                C  sf   t |r	|| S t|tr|s| |S t|r|j S t|r1t	d| }|j
j||d S |S )Nr   )context)r.   _to_compliant_exprr+   strcolr   _compliant_series_to_exprr   r   _series
from_numpy)rI   rJ   rL   nsr,   r,   r-   extract_compliantM   s   



rX   exprr6   r   excluder7   #tuple[Sequence[str], Sequence[str]]c                   sh   |  |}| jd u r|n| |} r0| jd usJ | jj r0t fddt||D  \}}||fS )Nc                   s    g | ]\}}| vr||fqS r,   r,   ).0xaliasrZ   r,   r-   
<listcomp>l   s
    z5evaluate_output_names_and_aliases.<locals>.<listcomp>)r;   rF   	_metadataexpansion_kindis_multi_unnamedr   )rY   r6   rZ   output_namesaliasesr,   r_   r-   !evaluate_output_names_and_aliases_   s   


rf   c                   @  s   e Zd ZdZe Z	 e Z	 e Z	 e Z	 e Z		 e Z
	 e Z	 e Z	 e Z	 e Z	 e Z	 edddZedddZedd
dZedddZdS )ExprKindz6Describe which kind of expression we are dealing with.r&   rM   c                 C     | t jt jhv S N)rg   LITERALAGGREGATIONselfr,   r,   r-   is_scalar_like      zExprKind.is_scalar_likec                 C  rh   ri   )rg   ORDERABLE_WINDOWORDERABLE_AGGREGATIONrl   r,   r,   r-   is_orderable_window   ro   zExprKind.is_orderable_windowr%   r   c                 C  s0   |j }|jr	tjS |jrtjS |jrtjS tjS ri   )	ra   
is_literalrg   rj   rn   rk   is_elementwiseELEMENTWISEUNKNOWN)clsr%   metar,   r,   r-   	from_expr   s   zExprKind.from_exprrK   rL   c                C  s<   t |r	| |S t|st|st|tr|stjS tjS ri   )	r.   ry   r   r   r+   rQ   rg   ru   rj   )rw   r%   rL   r,   r,   r-   from_into_expr   s   
zExprKind.from_into_exprNr&   rM   )r%   r   r&   rg   )r%   rK   rL   rM   r&   rg   )__name__
__module____qualname____doc__r   rj   rk   rq   ru   rp   WINDOW
FILTRATIONORDERABLE_FILTRATIONNARYOVERrv   propertyrn   rr   classmethodry   rz   r,   r,   r,   r-   rg   u   s@    
rg   7TypeIs[Literal[ExprKind.LITERAL, ExprKind.AGGREGATION]]c                 C  s   | j S ri   rn   )r%   r,   r,   r-   rn      s   rn   c                   @  sF   e Zd ZdZe Z	 e Z	 e Z	 dddZdddZ	dd
dZ
dS )ExpansionKindz8Describe what kind of expansion the expression performs.r&   rM   c                 C  s
   | t ju S ri   )r   MULTI_UNNAMEDrl   r,   r,   r-   rc      s   
zExpansionKind.is_multi_unnamedc                 C  rh   ri   )r   MULTI_NAMEDr   rl   r,   r,   r-   is_multi_output   s   zExpansionKind.is_multi_outputrJ   $Literal[ExpansionKind.MULTI_UNNAMED]c                 C  s4   | t ju r|t ju rt jS d|  d| d}t|)Nz+Unsupported ExpansionKind combination, got z and z, please report a bug.)r   r   r@   )rm   rJ   rA   r,   r,   r-   __and__   s   zExpansionKind.__and__Nr{   )rJ   r   r&   r   )r|   r}   r~   r   r   SINGLEr   r   rc   r   r   r,   r,   r,   r-   r      s    

r   c                   @  s  e Zd ZdZdZddddddddHddZdIddZdJddZedKd d!Z	dLd"d#Z
dLd$d%ZdLd&d'ZdLd(d)ZdLd*d+ZdLd,d-ZdLd.d/ZdLd0d1ZdLd2d3ZedLd4d5ZedLd6d7ZedLd8d9ZedLd:d;ZedLd<d=ZedMdBdCZedNdEdFZdGS )OExprMetadataa  Expression metadata.

    Parameters:
        expansion_kind: What kind of expansion the expression performs.
        has_windows: Whether it already contains window functions.
        is_elementwise: Whether it can operate row-by-row without context
            of the other rows around it.
        is_literal: Whether it is just a literal wrapped in an expression.
        is_scalar_like: Whether it is a literal or an aggregation.
        last_node: The ExprKind of the last node.
        n_orderable_ops: The number of order-dependent operations. In the
            lazy case, this number must be `0` by the time the expression
            is evaluated.
        preserves_length: Whether the expression preserves the input length.
    )rb   has_windowsrt   rs   rn   	last_noden_orderable_opspreserves_lengthFr   Tr   r   r   rt   rn   rs   rb   r   r   rg   r   rM   r   intr   rt   rn   rs   r&   Nonec          	      C  sL   |r|sJ |r|sJ || _ || _|| _|| _|| _|| _|| _|| _d S ri   rb   r   r   r   rt   r   rn   rs   )	rm   rb   r   r   r   r   rt   rn   rs   r,   r,   r-   __init__   s   
zExprMetadata.__init__argsr   kwdsr   c                O  s   d| j }t|)NzCannot subclass )r|   	TypeError)rw   r   r   rA   r,   r,   r-   __init_subclass__  s   zExprMetadata.__init_subclass__rQ   c                 C  sF   d| j  d| j d| j d| j d| j d| j d| j d| j d	S )
Nz ExprMetadata(
  expansion_kind: z,
  last_node: z,
  has_windows: z,
  n_orderable_ops: z,
  is_elementwise: z,
  preserves_length: z,
  is_scalar_like: z,
  is_literal: z,
)r   rl   r,   r,   r-   __repr__  s$   zExprMetadata.__repr__c                 C  s   | j  o| j S ri   )r   rn   rl   r,   r,   r-   is_filtration#  ro   zExprMetadata.is_filtrationc              
   C  s2   | j r	d}t|t| jtj| j| jdddddS )N4Can't apply aggregations to scalar-like expressions.FTr   )rn   r   r   rb   rg   rk   r   r   rm   rA   r,   r,   r-   with_aggregation'     zExprMetadata.with_aggregationc              
   C  s6   | j r	d}t|t| jtj| j| jd dddddS )Nr   r:   FTr   )rn   r   r   rb   rg   rq   r   r   r   r,   r,   r-   with_orderable_aggregation6  s   z'ExprMetadata.with_orderable_aggregationc              
   C  s(   t | jtj| j| j| j| j| j| j	dS )Nr   )
r   rb   rg   ru   r   r   r   rt   rn   rs   rl   r,   r,   r-   with_elementwise_opF  s   z ExprMetadata.with_elementwise_opc              
   C  s4   | j r	d}t|t| jtj| j| j| jddddS )Nz;Can't apply window (e.g. `rank`) to scalar-like expression.Fr   )	rn   r   r   rb   rg   r   r   r   r   r   r,   r,   r-   with_windowR  s   zExprMetadata.with_windowc              
   C  s8   | j r	d}t|t| jtj| j| jd | jddddS )NzNCan't apply orderable window (e.g. `diff`, `shift`) to scalar-like expression.r:   Fr   )	rn   r   r   rb   rg   rp   r   r   r   r   r,   r,   r-   with_orderable_windowd  s   z"ExprMetadata.with_orderable_windowc              
   C  sx   | j r	d}t|| js| jrd}t|| j}|s&| jtjur&d}t|| jjr.|d8 }t	| j
tjd|dddddS )NCannot nest `over` statements.oCannot use `over` on expressions which are elementwise
(e.g. `abs`) or which change length (e.g. `drop_nulls`).aP  Cannot use `order_by` in `over` on expression which isn't orderable.
If your expression is orderable, then make sure that `over(order_by=...)`
comes immediately after the order-dependent expression.

Hint: instead of
  - `(nw.col('price').diff() + 1).over(order_by='date')`
write:
  + `nw.col('price').diff().over(order_by='date') + 1`
r:   TFr   )r   r   rt   r   r   r   rg   r   rr   r   rb   r   )rm   rA   r   r,   r,   r-   with_ordered_overt  s0   	zExprMetadata.with_ordered_overc              
   C  sH   | j r	d}t|| js| jrd}t|t| jtjd| jdddddS )Nr   r   TFr   )	r   r   rt   r   r   rb   rg   r   r   r   r,   r,   r-   with_partitioned_over  s"   z"ExprMetadata.with_partitioned_overc              
   C  s2   | j r	d}t|t| jtj| j| jdddddS )NECan't apply filtration (e.g. `drop_nulls`) to scalar-like expression.Fr   )rn   r   r   rb   rg   r   r   r   r   r,   r,   r-   with_filtration  r   zExprMetadata.with_filtrationc              
   C  s6   | j r	d}t|t| jtj| j| jd dddddS )Nr   r:   Fr   )rn   r   r   rb   rg   r   r   r   r   r,   r,   r-   with_orderable_filtration  s   z&ExprMetadata.with_orderable_filtrationc                   C  s   t tjtjddddS )NFT)rt   r   rn   )r   r   r   rg   rk   r,   r,   r,   r-   aggregation  s   zExprMetadata.aggregationc                   C  s   t tjtjdddddS )NFT)rt   r   rs   rn   )r   r   r   rg   rj   r,   r,   r,   r-   literal  s   zExprMetadata.literalc                   C     t tjtjS ri   )r   r   r   rg   ru   r,   r,   r,   r-   selector_single     zExprMetadata.selector_singlec                   C  r   ri   )r   r   r   rg   ru   r,   r,   r,   r-   selector_multi_named  r   z!ExprMetadata.selector_multi_namedc                   C  r   ri   )r   r   r   rg   ru   r,   r,   r,   r-   selector_multi_unnamed  r   z#ExprMetadata.selector_multi_unnamedlhsr   rhsr!   c                C  s   t ||ddddS )NTFrL   allow_multi_outputto_single_outputcombine_metadata)rw   r   r   r,   r,   r-   from_binary_op  s   
zExprMetadata.from_binary_opr3   c                 G  s   t |ddddS )NFTr   r   )rw   r3   r,   r,   r-   from_horizontal_op  s   zExprMetadata.from_horizontal_opN)rb   r   r   rg   r   rM   r   r   r   rM   rt   rM   rn   rM   rs   rM   r&   r   )r   r   r   r   r&   r   )r&   rQ   r{   )r&   r   )r   r   r   r!   r&   r   )r3   r!   r&   r   )r|   r}   r~   r   	__slots__r   r   r   r   r   r   r   r   r   r   r   r   r   r   staticmethodr   r   r   r   r   r   r   r   r,   r,   r,   r-   r      sJ    








#

	
r   r   IntoExpr | object | Noner   r   c              
   G  s2  d}t j}d}d}d}d}	d}
d}t|D ]b\}}t|tr | r$t|r+d}d}
d}qt|rw|j}|j	 rR|j}|dkrF|sFd}t
||sR|dkrP||@ n|}||jO }||j7 }||jO }|	|jM }	|
|jM }
||jM }|t|j7 }q|dkrd}t||r|rd}t|t|tj||||	|
|dS )	a  Combine metadata from `args`.

    Arguments:
        args: Arguments, maybe expressions, literals, or Series.
        str_as_lit: Whether to interpret strings as literals or as column names.
        allow_multi_output: Whether to allow multi-output inputs.
        to_single_output: Whether the result is always single-output, regardless
            of the inputs (e.g. `nw.sum_horizontal`).
    r   FTz]Multi-output expressions (e.g. nw.col('a', 'b'), nw.all()) are not supported in this context.r:   zXLength-changing expressions can only be used in isolation, or followed by an aggregationzVCannot combine length-changing expressions with length-preserving ones or aggregationsr   )r   r   	enumerater+   rQ   r2   r.   ra   rb   r   r   r   r   r   rt   rn   rs   r   r   r   r   rg   r   )rL   r   r   r   n_filtrationsresult_expansion_kindresult_has_windowsresult_n_orderable_opsresult_preserves_lengthresult_is_elementwiseresult_is_scalar_likeresult_is_literaliargmetadatarb   rA   r,   r,   r-   r     sd   






r   r!   function_namerQ   r   c                   s:   ddl m  t fdd|D sd|  d}t|d S )Nr   r   c                 3  s.    | ]}t |r|jjpt|t fV  qd S ri   )r.   ra   r   r+   rQ   r\   r]   r   r,   r-   	<genexpr>Q  s
    
z4check_expressions_preserve_length.<locals>.<genexpr>zBExpressions which aggregate or change length cannot be passed to 'z'.)r1   r    allr   )r   r   rA   r,   r   r-   !check_expressions_preserve_lengthK  s   r   kwargsc                  O  s    t | | }tdd |D S )Nc                 s  s     | ]}t |o|jjV  qd S ri   )r.   ra   rn   r   r,   r,   r-   r   ^  s    z,all_exprs_are_scalar_like.<locals>.<genexpr>)r   valuesr   )r   r   r3   r,   r,   r-   all_exprs_are_scalar_likeY  s   r   n_ary_functionCallable[..., CompliantExprAny]
comparandsc                  sX   fdd|D }fdd|D }t dd |D   fddt||D }|| S )Nc                 3  s    | ]
}t  |d V  qdS )rL   N)rX   r\   	comparand)rI   rL   r,   r-   r   g  s
    
z(apply_n_ary_operation.<locals>.<genexpr>c                   s   g | ]	}t j| d qS )r   )rg   rz   r   r   r,   r-   r`   k  s    z)apply_n_ary_operation.<locals>.<listcomp>c                 s  s    | ]}|j  V  qd S ri   r   )r\   kindr,   r,   r-   r   p  s    c                 3  s6    | ]\}} rt |rt|r||n|V  qd S ri   )r   rn   	broadcast)r\   compliant_exprr   )r   r,   r-   r   q  s    
)anyr   )rI   r   rL   r   compliant_exprskindsr,   )r   rI   rL   r-   apply_n_ary_operationa  s   

r   N)r%   r   r&   r'   )r%   r   r&   r0   )r3   r4   r&   r5   )r3   r   r&   rC   )rI   r   rJ   rK   rL   rM   r&   rN   )rY   r   r6   r   rZ   r7   r&   r[   )r%   rg   r&   r   )
r   r   rL   rM   r   rM   r   rM   r&   r   )r   r!   r   rQ   r&   r   )r   r!   r   r!   r&   rM   )
rI   r   r   r   r   rK   rL   rM   r&   r   )=
__future__r   enumr   r   	itertoolsr   typingr   r   r   r	   r
   r   narwhals._utilsr   r   narwhals.dependenciesr   r   narwhals.exceptionsr   r   collections.abcr   typing_extensionsr   r   narwhals._compliantr   r   narwhals._compliant.typingr   r   r   r   r   r   r*   r   r1   r    narwhals.typingr!   r"   r#   r$   r.   r2   rB   rH   rX   rf   rg   rn   r   r   r   r   r   r   r,   r,   r,   r-   <module>   s@     






F  
#
N
