
    bMhM              	          d dl Z d dlZd dlZd dlmZmZ d dlmc m	Z
 d dlZd dlmZmZmZmZmZ d dlmZ d dlmZ  ej        ddgd           d	             Z ej        d ej        d
 e
j        d
                    gd           d             Zd Z G d d          Z G d d          Z G d d          Z e
j        d
           G d d                      Z  e
j        d
           G d de                       Z! G d de           Z" G d de!          Z# G d d          Z$ G d d           Z% G d! d"          Z&dS )#    N)NumExprClobberingErrorUndefinedVariableError)	DataFrameIndex
MultiIndexSeries
date_range)NUMEXPR_INSTALLEDpythonpandasc                     | S N xs    b/var/www/html/test/jupyter/venv/lib/python3.11/site-packages/pandas/tests/frame/test_query_eval.py<lambda>r      s    1     )paramsidsc                     | j         S r   paramrequests    r   parserr      s
    =r   numexpr)marksc                     | S r   r   r   s    r   r   r      s    ! r   c                     | j         S r   r   r   s    r   enginer!      s    
 =r   c                 D    | dk    rt          j        d|             d S d S )Nr   zcannot evaluate with parser=)pytestskip)r   s    r   skip_if_no_pandas_parserr%   %   s4    ;6;;<<<<< r   c                       e Zd Zej        d             Zej        d             Zej        d             Zd Zd Z	d Z
d ZdS )	
TestCompatc                 (    t          dg di          S )NA         r   selfs    r   dfzTestCompat.df+   s    #yyy)***r   c                 $    ||j         dk             S )Nr   r)   r0   r1   s     r   	expected1zTestCompat.expected1/   s    "$(|r   c                     |j         dz   S Nr+   r3   r4   s     r   	expected2zTestCompat.expected23   s    taxr   c                     |                     d          }t          j        ||           |                    d          }t          j        ||d           d S )NA>0A+1Fcheck_namesquerytmassert_frame_equalevalassert_series_equalr0   r1   r5   r8   results        r   test_query_defaultzTestCompat.test_query_default7   sS     %
fi000
vyeDDDDDDr   c                     |                     dd           }t          j        ||           |                    dd           }t          j        ||d           d S )Nr:   r!   r;   Fr<   r>   rD   s        r   test_query_NonezTestCompat.test_query_None?   s]    %--
fi000t,,
vyeDDDDDDr   c                     |                     dd          }t          j        ||           |                    dd          }t          j        ||d           d S )Nr:   r   rH   r;   Fr<   r>   rD   s        r   test_query_pythonzTestCompat.test_query_pythonE   s]    %11
fi000x00
vyeDDDDDDr   c                    t           r\|                    dd          }t          j        ||           |                    dd          }t          j        ||d           d S d}t          j        t          |          5  |                    dd           d d d            n# 1 swxY w Y   t          j        t          |          5  |                    dd           d d d            d S # 1 swxY w Y   d S )	Nr:   r   rH   r;   Fr<   z~'numexpr' is not installed or an unsupported version. Cannot use engine='numexpr' for query/eval if 'numexpr' is not installedmatch)	r
   r?   r@   rA   rB   rC   r#   raisesImportError)r0   r1   r5   r8   rE   msgs         r   test_query_numexprzTestCompat.test_query_numexprK   s    	1XXeIX66F!&)444WWU9W55F"69%HHHHHH! 
 {#666 2 2y1112 2 2 2 2 2 2 2 2 2 2 2 2 2 2{#666 1 1i0001 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1s$   B%%B),B)C00C47C4N)__name__
__module____qualname__r#   fixturer1   r5   r8   rF   rI   rK   rR   r   r   r   r'   r'   *   s        ^+ + ^+ ^  ^ ^  ^E E EE E EE E E1 1 1 1 1r   r'   c                       e Zd Zej                            dddg          ej                            dg d          d                         Zd Zd Zd	 Z	d
 Z
d Zd Zd Zd ZdS )TestDataFrameEvaln   i  zop_str,op,rop))+__add____radd__)-__sub____rsub__)*__mul____rmul__)/__truediv____rtruediv__c                 t   t          dt          |          t          d                    }d|j        d<   |                                }t          t          j        |j        |                              |d          t          d                    }t          d| d	          }t          d
| d	          }	t          j        |	|           |dv r0 t          ||          |          }	t          j        |	|           d S |dv r0 t          ||          |          }	t          j        |	|           d S d S )Nr+   abcdindexcolumnsr,   r   rk   zbase z dfzm )r[   ra   )r^   rd   )r   rangelistilocmeannptilevaluesreshaperB   r@   rA   getattr)
r0   op_stroproprY   r1   mbaseexpectedrE   s
             r   test_opszTestDataFrameEval.test_ops_   sF    qa$v,,???
GGIIGAHa  ((B//f
 
 
 ++++,, &6&&&''
fh/// $WR__Q''F!&(33333 :%WR%%a((F!&(33333 r   c                 n   t          dt          j                            d                              d          i          }t          j        |j        dd<   dt          j        |j        dd                   z
  }dt          j        |          z
  j        dd         }t          j	        ||           d S )Nr)   r,   ia  r      r+      )
r   rr   randomdefault_rngstandard_normalnanrp   isnanr@   rA   )r0   r1   r|   rE   s       r   test_dataframe_sub_numexpr_pathz1TestDataFrameEval.test_dataframe_sub_numexpr_path   s     RY22155EEeLLMNNv!rx"...bhrll"(2.
fh/////r   c                 Z   t          g dg dd          }d}t          j        t          |          5  |                    d            d d d            n# 1 swxY w Y   t          j        t          |          5  |                    d           d d d            d S # 1 swxY w Y   d S )Nr*   )abr   r)   Bz%expr must be a string to be evaluatedrM   c                     | j         dk    S )Nr   )r   r   s    r   r   z6TestDataFrameEval.test_query_non_str.<locals>.<lambda>   s    qscz r   o   r   r#   rO   
ValueErrorr?   r0   r1   rQ   s      r   test_query_non_strz$TestDataFrameEval.test_query_non_str   s.   YYY___==>>5]:S111 	+ 	+HH))***	+ 	+ 	+ 	+ 	+ 	+ 	+ 	+ 	+ 	+ 	+ 	+ 	+ 	+ 	+ ]:S111 	 	HHSMMM	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	s#   AAA=B  B$'B$c                     t          dg di          }d}t          j        t          |          5  |                    d           d d d            d S # 1 swxY w Y   d S )Nr)   r*   zexpr cannot be an empty stringrM    r   r   s      r   test_query_empty_stringz)TestDataFrameEval.test_query_empty_string   s    YYY'((.]:S111 	 	HHRLLL	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	s   AAAc                 h   t          t          j                            d                              d          t          d                    }ddi}ddi}|                    d||g	          |d         |d         z   k    sJ t          j        d||g	          |d         |d         z   k    sJ d S )
Nr,   
   r,   abrm   r   r+   r   a + b	resolvers)r   rr   r   r   r   ro   rB   pd)r0   r1   dict1dict2s       r   test_eval_resolvers_as_listz-TestDataFrameEval.test_eval_resolvers_as_list   s    I!!!$$44W==tDzz
 
 
 aawww5%.w99U3Z%PS*=TTTTTww5%.999U3Z%PS*=TTTTTTTr   c                 8   t          t          j                            d                              d          t          d                    }ddi}|                    d|g          }|d         |d	         |d         z  z   }t          j        ||           d S )
Nr,   r   r   rm   cz	a + b * cr   r   r   )	r   rr   r   r   r   ro   rB   r@   rC   )r0   r1   r   rE   r|   s        r   test_eval_resolvers_combinedz.TestDataFrameEval.test_eval_resolvers_combined   s    I!!!$$44W==tDzz
 
 
 a 88c7RWuSz11
vx00000r   c                     t          dddgi          }|                    d          }t          ddgddgd          }t          j        ||           d S )Na1YNzc = ((a1 == 'Y') & True)TF)r   r   )r   rB   r@   rA   )r0   r1   resr|   s       r   test_eval_object_dtype_binopz.TestDataFrameEval.test_eval_object_dtype_binop   s`    sCj)**gg011S#JdE]CCDD
c8,,,,,r   c                    |dk    r5t           j                            d          }|                    |           t	          t          j        g d          t          j        g d          d          }|                    d||          }t          t          j        g d	                    }t          j
        ||           d S )
Nr   z/numexpr does not support extension array dtypes)reasonr*   rZ   r      r   r   za / br!   r   )g      ?g?      ?)r#   markxfailapplymarkerr   r   arrayrB   r   r@   rC   )r0   r!   r   r   r   r1   rE   r|   s           r   test_extension_array_evalz+TestDataFrameEval.test_extension_array_eval   s    Y;$$H %  D %%%RXiii00rx			7J7JKKLL??"(#5#5#56677
vx00000r   c                     t          dgdgd          }|                    d||          }t          dg          }t          j        ||           d S )Ny      ?       @y      ?      ?r   za/br   y      ?      ?)r   rB   r   r@   rC   )r0   r!   r   r1   rE   r|   s         r   test_complex_evalz#TestDataFrameEval.test_complex_eval   sZ    fXVH5566vf==:,''
vx00000r   N)rS   rT   rU   r#   r   parametrizer}   r   r   r   r   r   r   r   r   r   r   r   rX   rX   ]   s        [S1d),,[	
 	
 	
 4 4  -,480 0 0	 	 	  U U U1 1 1- - -
1 
1 
11 1 1 1 1r   rX   c                   &    e Zd Zd Zd Zd Zd ZdS ) TestDataFrameQueryWithMultiIndexc                 \   t          |           t          j                            d                              ddgd          }t          j                            d                              ddgd          }t          j        ||gdd	g
          }t          t          j                            d                              d          |          }t          |j
                            d          j        |d          }|                    d||          }|                    d||          }	||dk             }
t          j        ||
           t          j        |	|
           |                    d||          }|                    d||          }	||dk             }
t          j        ||
           t          j        |	|
           |                    d||          }|                    d||          }	||                    dg                   }
t          j        ||
           t          j        |	|
           |                    d||          }|                    d||          }	||                    dg                    }
t          j        ||
           t          j        |	|
           |                    d||          }|                    d||          }	||                    dg                   }
t          j        ||
           t          j        |	|
           |                    d||          }|                    d||          }	||                    dg                    }
t          j        ||
           t          j        |	|
           d S )Nr,   redgreenr   sizeeggshamcolorfoodnamesr   rj   rj   namezcolor == "red"r   r!   z"red" == colorzcolor != "red"z"red" != colorzcolor == ["red"]z["red"] == colorzcolor != ["red"]z["red"] != colorz["red"] in colorz"red" in colorz["red"] not in colorz"red" not in colorr%   rr   r   r   choicer   from_arraysr   r   r   rj   get_level_valuesrt   r?   r@   rA   isinr0   r   r!   r   r   rj   r1   indres1res2exps              r    test_query_with_named_multiindexzATestDataFrameQueryWithMultiIndex.test_query_with_named_multiindex   sz    (((I!!!$$++UG,<2+FFI!!!$$++VUO"+EE&1vgv5FGGGry,,Q//??HHPUVVVH%%g..5U
 
 

 xx(xGGxx(xGG
dC(((
dC((( xx(xGGxx(xGG
dC(((
dC((( xx*6&xIIxx*6&xII5'""#
dC(((
dC(((xx*6&xIIxx*6&xII#((E7###$
dC(((
dC((( xx*6&xIIxx(xGG5'""#
dC(((
dC(((xx.vfxMMxx,VFxKK#((E7###$
dC(((
dC(((((r   c                    t          |           t          j                            d                              ddgd          }t          j                            d                              ddgd          }t          j        ||g          }t          t          j                            d                              d          |	          }t          |j
                            d
          j        |	          }|                    d||          }|                    d||          }	||dk             }
t          j        ||
           t          j        |	|
           |                    d||          }|                    d||          }	||dk             }
t          j        ||
           t          j        |	|
           |                    d||          }|                    d||          }	||                    dg                   }
t          j        ||
           t          j        |	|
           |                    d||          }|                    d||          }	||                    dg                    }
t          j        ||
           t          j        |	|
           |                    d||          }|                    d||          }	||                    dg                   }
t          j        ||
           t          j        |	|
           |                    d||          }|                    d||          }	||                    dg                    }
t          j        ||
           t          j        |	|
           t          |j
                            d          j        |	          }|                    d||          }|                    d||          }	||dk             }
t          j        ||
           t          j        |	|
           |                    d||          }|                    d||          }	||dk             }
t          j        ||
           t          j        |	|
           |                    d||          }|                    d||          }	||                    dg                   }
t          j        ||
           t          j        |	|
           |                    d||          }|                    d ||          }	||                    dg                    }
t          j        ||
           t          j        |	|
           |                    d!||          }|                    d"||          }	||                    dg                   }
t          j        ||
           t          j        |	|
           |                    d#||          }|                    d$||          }	||                    dg                    }
t          j        ||
           t          j        |	|
           d S )%Nr,   r   r   r   r   r   r   r   r   r   ilevel_0 == "red"r   z"red" == ilevel_0ilevel_0 != "red"z"red" != ilevel_0zilevel_0 == ["red"]z["red"] == ilevel_0zilevel_0 != ["red"]z["red"] != ilevel_0z["red"] in ilevel_0z"red" in ilevel_0z["red"] not in ilevel_0z"red" not in ilevel_0r+   zilevel_1 == "eggs"z"eggs" == ilevel_1zilevel_1 != "eggs"z"eggs" != ilevel_1zilevel_1 == ["eggs"]z["eggs"] == ilevel_1zilevel_1 != ["eggs"]z["eggs"] != ilevel_1z["eggs"] in ilevel_1z"eggs" in ilevel_1z["eggs"] not in ilevel_1z"eggs" not in ilevel_1r   r   s              r   "test_query_with_unnamed_multiindexzCTestDataFrameQueryWithMultiIndex.test_query_with_unnamed_multiindex  s    (((I!!!$$++UG,<2+FFI!!!$$++VUO"+EE&1v..ry,,Q//??HHPUVVVRX..q118FFFxx+F6xJJxx+F6xJJ
dC(((
dC((( xx+F6xJJxx+F6xJJ
dC(((
dC((( xx-fVxLLxx-fVxLL5'""#
dC(((
dC(((xx-fVxLLxx-fVxLL#((E7###$
dC(((
dC((( xx-fVxLLxx+F6xJJ5'""#
dC(((
dC(((xx1&xPPxx/vxNN#((E7###$
dC(((
dC((( RX..q118FFFxx,VFxKKxx,VFxKK
dC(((
dC((( xx,VFxKKxx,VFxKK
dC(((
dC((( xx.vfxMMxx.vfxMM6(##$
dC(((
dC(((xx.vfxMMxx.vfxMM#((F8$$$%
dC(((
dC((( xx.vfxMMxx,VFxKK6(##$
dC(((
dC(((xx26&xQQxx0xOO#((F8$$$%
dC(((
dC(((((r   c                    t          |           t          j                            d                              ddgd          }t          j        d          }t          j        ||g          }d dg|_        t          t          j                            d          
                    d          |          }|                    d	||
          }t          |j                            d          j        |d          }||dk             }	t!          j        ||	           |                    d||
          }t          |j                            d          j        |d          }||dk             }	t!          j        ||	           |                    d||
          }t          |j                            d          j        |          }||dk             }	t!          j        ||	           |                    d||
          }t          |j                            d          j        |          }||dk             }	t!          j        ||	           d S )Nr,   r   r   r   r   ratingr   r   zrating == 1r   r   r+   zrating != 1r   r   r   )r%   rr   r   r   r   aranger   r   r   r   r   r?   r   rj   r   rt   r@   rA   )
r0   r   r!   r   r   rj   r1   r   r   r   s
             r   *test_query_with_partially_named_multiindexzKTestDataFrameQueryWithMultiIndex.test_query_with_partially_named_multiindex^  s+    (((I!!!$$++UG,<2+FFIbMM&1v..X&ry,,Q//??HHPUVVVhh}VFhCCH%%h//6e(
 
 
 l
c3'''hh}VFhCCH%%h//6e(
 
 
 l
c3'''hh*6&hIIRX..q118FFF
c3'''hh*6&hIIRX..q118FFF
c3'''''r   c           	      t   t          t          j        d          t          j        d t          d          D             ddg                    }|                                }d }|j                                        }|j	        | ||j	        d           ||j	        d          |d	}|
                                D ]w\  }}t          |t                    r|                    ||                   sJ 8t          |t                    rt          j        |||                    it#          d
          d S )Nr   r-   c                 ,    g | ]}t          d           S )r   )rn   ).0_s     r   
<listcomp>z^TestDataFrameQueryWithMultiIndex.test_query_multiindex_get_index_resolvers.<locals>.<listcomp>  s    ---qr---r   r,   spamr   r   r   c                 f    |                      |          }|                                }| |_        |S r   )r   	to_seriesrj   )milevellevel_valuesss       r   r   z]TestDataFrameQueryWithMultiIndex.test_query_multiindex_get_index_resolvers.<locals>.to_series  s2    ..u55L&&((AAGHr   )rj   rk   r   r   clevel_0z object must be a Series or Index)r   rr   onesr   r   rn   _get_index_resolversrk   r   rj   items
isinstancer   is_r   r@   rC   AssertionError)r0   r1   r   r   
col_seriesr|   kvs           r   )test_query_multiindex_get_index_resolverszJTestDataFrameQueryWithMultiIndex.test_query_multiindex_get_index_resolvers}  sa   GG(--E!HH---ff5E  
 
 
 ++--		 	 	 Z))++
X!Ibh//Ibh//"
 
 OO%% 	I 	IDAq!U## IuuXa[))))))Av&& I&q(1+6666$%GHHH	I 	Ir   N)rS   rT   rU   r   r   r   r   r   r   r   r   r      sZ        0) 0) 0)dV) V) V)p( ( (>I I I I Ir   r   c                   F   e Zd Zej        d             Zej        d             Zd Zd Zd Z	d Z
d Zd Zd	 Zd
 Zd Zd Zd Zd Zd Zd Zd Zd Zd Zd Zd Zd Zd Zd Zej                            ddej         gdej!        gg          d             Z"d Z#d Z$dS ) TestDataFrameQueryNumExprPandasc                     dS Nr   r   r/   s    r   r!   z&TestDataFrameQueryNumExprPandas.engine      yr   c                     dS Nr   r   r/   s    r   r   z&TestDataFrameQueryNumExprPandas.parser      xr   c                    t          |           t          t          j                            d                              d                    }t          dd          |d<   t          dd          |d<   t          d	d          |d
<   |                    d||          }||j        dk     d|j	        k     z           }t          j        ||           d S )Nr,   r   r-   1/1/2012r   periodsdates11/1/2013dates21/1/2014dates3z"@df.dates1 < 20130101 < @df.dates3r   20130101)r%   r   rr   r   r   r   r	   r?   r  r  r@   rA   r0   r!   r   r1   r   expecs         r   %test_date_query_with_attribute_accesszETestDataFrameQueryNumExprPandas.test_date_query_with_attribute_access  s     (((ry,,Q//??GGHH!*a8888!*a8888!*a8888hh0  
 
 BI
*zBI/EFG
c5)))))r   c                    t          t          j                            d                              d                    }t          dd          |d<   t          dd          |d<   t          d	d          |d
<   |                    d||          }||j        dk     d|j        k     z           }t          j
        ||           d S )Nr,   r   r   r   r   r  r  r  r  r  dates1 < 20130101 < dates3r   r  r   rr   r   r   r   r	   r?   r  r  r@   rA   r  s         r   #test_date_query_no_attribute_accesszCTestDataFrameQueryNumExprPandas.test_date_query_no_attribute_access  s    ry,,Q//??GGHH!*a8888!*a8888!*a8888hh3F6hRRBI
*zBI/EFG
c5)))))r   c                    d}t          t          j                            d                              |df                    }t          d|          |d<   t          d|          |d<   t          d	|          |d
<   t          j        |j        t          j                            d                              |          dk    df<   t          j        |j        t          j                            d                              |          dk    d
f<   |	                    d||          }||j
        dk     d|j        k     z           }t          j        ||           d S )Nr   r,   r-   r   r   r  r  r  r  r  r   r  r   r  r   rr   r   r   r   r	   r   NaTlocr?   r  r  r@   rA   r0   r!   r   rY   r1   r   r  s          r   test_date_query_with_NaTz8TestDataFrameQueryNumExprPandas.test_date_query_with_NaT  s2   ry,,Q//??AGGHH!*a8888!*a8888!*a8888EGVry$$Q''..q11C7ABEGVry$$Q''..q11C7ABhh3F6hRRBI
*zBI/EFG
c5)))))r   c                    d}t          t          j                            d                              |df                    }t          d|          |d<   t          d|          |d<   |                    dd	d	
          }|J |                    d||          }||j        dk     d|j	        k     z           }t          j        ||           d S )Nr   r,   r-   r   r   r  r  r  Tinplacedropindex < 20130101 < dates3r   r  r   rr   r   r   r   r	   	set_indexr?   rj   r  r@   rA   r0   r!   r   rY   r1   return_valuer   r  s           r   test_date_index_queryz5TestDataFrameQueryNumExprPandas.test_date_index_query  s    ry,,Q//??AGGHH!*a8888!*a8888||Hd|FF###hh26&hQQBHz)j29.DEF
c5)))))r   c                    d}t          t          j                            d                              |df                                        dt          i          }t          d|          |d<   t          d|          |d	<   t          j	        |j
        d
<   |                    ddd          }|J |                    d||          }||j        dk     d|j        k     z           }t          j        ||           d S )Nr   r,   r-   r   r   r   r  r  r  r   r   Tr  r  r   r  r   rr   r   r   r   astypeobjectr	   r   r  rp   r  r?   rj   r  r@   rA   r  s           r   test_date_index_query_with_NaTz>TestDataFrameQueryNumExprPandas.test_date_index_query_with_NaT  s    ry,,Q//??AGGHHOOK
 
 "*a8888!*a8888||Hd|FF###hh26&hQQBHz)j29.DEF
c5)))))r   c                     d}i }t          d|          |d<   t          d|          |d<   t          |          }t          j        |j        t
          j                            d                              |          dk    df<   |                    dd	d	
          }|J |	                    d||          }||j
                                        dk     d|j        k     z           }t          j        ||           d S )Nr   r   r   r  r  r  r,   r   Tr  r  r   r  )r	   r   r   r  r  rr   r   r   r  r?   rj   r   r  r@   rA   )	r0   r!   r   rY   dr1   r  r   r  s	            r   )test_date_index_query_with_NaT_duplicateszITestDataFrameQueryNumExprPandas.test_date_index_query_with_NaT_duplicates  s     Q777( Q777(q\\EGVry$$Q''..q11C7AB||Hd|FF###hh3F6hRRBH&&((:5*ry:PQR
c5)))))r   c                    d}t          t          d|          t          j        |          d          }|                    d||          }t          |          dk    sJ |                    d||          }t          j        ||           d	}d
D ]Q}t          j	        t          |          5  |                    d| d||           d d d            n# 1 swxY w Y   Rd S )Nr   r   r   )datesnondatezdates == nondater   r   zdates != nondatez=Invalid comparison between dtype=datetime64\[ns\] and ndarray)<><=>=rM   zdates z nondate)r   r	   rr   r   r?   lenr@   rA   r#   rO   	TypeError)r0   r!   r   rY   r1   rE   rQ   rx   s           r   test_date_query_with_non_datez=TestDataFrameQueryNumExprPandas.test_date_query_with_non_date  sT    Q777BIaLLQQ
 
 ,VFKK6{{a,VFKK
fb)))N( 	N 	NBy444 N N."...vfMMMN N N N N N N N N N N N N N N	N 	Ns   1CC	!C	c                    t          t          d          t          dd          t          dd          d          }d}t          j        t          |          5  |                    d	||
           d d d            d S # 1 swxY w Y   d S )Nr   r-      rZ      )ir[   rzinvalid syntaxrM   zi - +r   )r   rn   r#   rO   SyntaxErrorr?   )r0   r!   r   r1   rQ   s        r   test_query_syntax_errorz7TestDataFrameQueryNumExprPandas.test_query_syntax_error   s    U2YYU1b\\aMMNN];c222 	< 	<HHWVFH;;;	< 	< 	< 	< 	< 	< 	< 	< 	< 	< 	< 	< 	< 	< 	< 	< 	< 	<s   B  BBc                    t          |           t          t          j                            d                              d          t          d                    }d\  }}|                    d||          }||j        |j	        k             }t          j        ||           |                    d||          }|||j	        k             }t          j        ||           t          j        t          d	
          5  |                    d||           d d d            n# 1 swxY w Y   t          j        t          d
          5  |                    d||           d d d            d S # 1 swxY w Y   d S )Nr,   )   r,   r   rm   )r+   r,   za > br   z@a > b!local variable 'c' is not definedrM   z@a > b > @czname 'c' is not definedz
@a > b > c)r%   r   rr   r   r   r   ro   r?   r   r   r@   rA   r#   rO   r   )r0   r!   r   r1   r   r   r   r|   s           r   test_query_scopez0TestDataFrameQueryNumExprPandas.test_query_scope  s    (((I!!!$$44W==tDzz
 
 
 1hhwvfh==bdRTk?
c8,,,hhxvh>>a"$h<
c8,,, ]"*M
 
 
 	B 	B HH]6&HAAA	B 	B 	B 	B 	B 	B 	B 	B 	B 	B 	B 	B 	B 	B 	B ]19RSSS 	A 	AHH\&H@@@	A 	A 	A 	A 	A 	A 	A 	A 	A 	A 	A 	A 	A 	A 	A 	A 	A 	As$   ?D$$D(+D(
E00E47E4c                 J   dx}}t          t          j                            d                              ||df          t          d                    }t          j        t          d          5  |	                    d	||
           d d d            d S # 1 swxY w Y   d S )Nr   r,   r-   r   abcrm   zname 'sin' is not definedrM   sin > 5r   )
r   rr   r   r   integersro   r#   rO   r   r?   )r0   r!   r   rY   rz   r1   s         r   test_query_doesnt_pickup_localz>TestDataFrameQueryNumExprPandas.test_query_doesnt_pickup_local   s    
AI!!!$$--aq!f-==tE{{
 
 

 ]19TUUU 	> 	>HHYvfH===	> 	> 	> 	> 	> 	> 	> 	> 	> 	> 	> 	> 	> 	> 	> 	> 	> 	>s   2BBBc                 f   dx}}t          t          j                            d                              ||df          t          d                    }d|j        _        d}t          j	        t          |	          5  |                    d
||           d d d            d S # 1 swxY w Y   d S )Nr   r,   r-   r   r=  rm   sinzVariables in expression.+rM   r>  r   )r   rr   r   r   r?  ro   rj   r   r#   rO   r   r?   )r0   r!   r   rY   rz   r1   rQ   s          r   test_query_builtinz2TestDataFrameQueryNumExprPandas.test_query_builtin*  s    
AI!!!$$--aq!f-==tE{{
 
 
 )]1=== 	> 	>HHYvfH===	> 	> 	> 	> 	> 	> 	> 	> 	> 	> 	> 	> 	> 	> 	> 	> 	> 	>s    B&&B*-B*c                    t          t          j                            d                              d          g d          }t          j        |                    d||          ||j        |j	        k                         t          j        |                    d||          ||j        |j	        z   |j	        |j
        z  k                        d S )Nr,   r   r   r   r   rm   a < br   za + b > b * c)r   rr   r   r   r   r@   rA   r?   r   r   r   r0   r!   r   r1   s       r   
test_queryz*TestDataFrameQueryNumExprPandas.test_query5  s    I!!!$$44W==
 
 
 	HHWVFH;;Rrt_	
 	
 	
 	HH_VFHCCrtbd{RTBD[()	
 	
 	
 	
 	
r   c                    t          t          j                            d                              dd          t          t          d          d          g d          }|                    d	||
          }||j        dk     |j	        |j
        k     z           }t          j        ||           |                    d||
          }||j        |j
        k              }t          j        ||           d S )Nr,   r   r   r   blobr   rE  ri   z(blob < 5) & (a < b)r   r   zblob < b)r   rr   r   r   r?  r   rn   r?   rj   r   r   r@   rA   r  s         r   test_query_index_with_namez:TestDataFrameQueryNumExprPandas.test_query_index_with_nameB  s    I!!!$$--bw-??b		///#OO
 
 

 hh-fVhLLBHqLRTBD[12
c5)))hhz&h@@28bd?#
c5)))))r   c                    t          t          j                            d                              dd          t          d          g d          }|                    d||          }||j        |j        k              }t          j
        ||           |                    d	||          }||j        d
k              }t          j
        ||           d S )Nr,   r   r   r   rE  ri   z	index < br   z	index < 5r   )r   rr   r   r   r?  rn   r?   rj   r   r@   rA   r  s         r   test_query_index_without_namez=TestDataFrameQueryNumExprPandas.test_query_index_without_nameQ  s    I!!!$$--bw-??))#OO
 
 
 hh{6&hAA28bd?#
c5))) hh{6&hAA28a< 
c5)))))r   c                     t          |           t          t          j                            d                              d                    }t          t          j                            d                              d                    }||dk    |dk    z           }|                    d||          }t          j        ||           t          j
        d||          }t          j        ||           t          j
        d||          }||dk    |dk    z  ||dk             dk    z           }t          j        ||           t          j
        d||          }|                    d	||          }t          j        ||           d S )
Nr,   r   r   z(@df > 0) & (@df2 > 0)r   zdf[df > 0 and df2 > 0]z)df[df > 0 and df2 > 0 and df[df > 0] > 0]zdf[(df>0) & (df2>0)](@df>0) & (@df2>0))r%   r   rr   r   r   r   r?   r@   rA   r   rB   )r0   r!   r   r1   df2r|   rE   s          r   test_nested_scopez1TestDataFrameQueryNumExprPandas.test_nested_scopeb  st    (((ry,,Q//??GGHH	--a00@@HHIIrAv#'*+26&QQ
fh///1&PPP
fh///7v
 
 
 rAv#'*baj1n=>
fh////vNNN8808OO
fh/////r   c                    t          t          j                            d                              d                    }t          j        t          d          5  |                    d||           d d d            d S # 1 swxY w Y   d S )Nr,   r   name 'df' is not definedrM   zdf > 0r   )	r   rr   r   r   r   r#   rO   r   r?   rG  s       r   *test_nested_raises_on_local_self_referencezJTestDataFrameQueryNumExprPandas.test_nested_raises_on_local_self_referencey  s    ry,,Q//??GGHH ]19STTT 	= 	=HHXfVH<<<	= 	= 	= 	= 	= 	= 	= 	= 	= 	= 	= 	= 	= 	= 	= 	= 	= 	=s   BBBc                    t          |           t          t          j                            d                              d          t          d                    }d}||j        |k              }|                    d||          }t          j
        ||           ||j        |j        k              }|                    d||          }t          j
        ||           d S )	Nr,   )d   r   
abcdefghijrm   r+   za < @br   rF  )r%   r   rr   r   r   r   ro   r   r?   r@   rA   r   )r0   r!   r   r1   r   expectrE   s          r   test_local_syntaxz1TestDataFrameQueryNumExprPandas.test_local_syntax  s     (((I!!!$$44Y??&&
 
 
 BD1H(6&AA
ff---BD24K'&@@
ff-----r   c                    t          |           t          d          }t          t          j                            d                              dt          |          f          |          }|                    d||          }|j	        |j
        k     |j
        |j        k     z  |j
                            |j	                   z  |j                            |j
                   z  }||         }t          j        ||           d S )Nr=  r,   rW  rm   z!a < b < c and a not in b not in cr   )r%   ro   r   rr   r   r   r   r.  r?   r   r   r   r   r@   rA   )r0   r!   r   colsr1   r   r   r  s           r   test_chained_cmp_and_inz7TestDataFrameQueryNumExprPandas.test_chained_cmp_and_in  s     (((E{{I!!!$$44c3t995EFFPT
 
 
 hh/v  
 
 tbd{rtbd{+rtyy.>>"$))BD//AQQ3
c5)))))r   c                 $   t          |           t          t          j                            d                              dd          d          }t          t          j                            d                              dd          d          }t          ||d	          }|j        |j        d
z
  	                    |                   }|
                    d||          }t          j        ||           t          t          j                            d                              dd          d          }|j        |d
z
  	                    |                   }|
                    d||          }t          j        ||           d S )Nr,   r-      r   r   rK  r   r   r   r+   z
b - 1 in ar   z@b - 1 in a)r%   r   rr   r   r   r?  r   r  r   r   r?   r@   rA   )r0   r!   r   r   r   r1   r|   rE   s           r   test_local_variable_with_inz;TestDataFrameQueryNumExprPandas.test_local_variable_with_in  sV    (((29((++44QR4@@sKKK29((++44Rb4AALLLQQ''((624!8//!,,-,vfEE
h///29((++44Rb4AALLL61q5,,q//*-vFF
h/////r   c                     t          |           d}t          dg di          }|                    d||          }||j        dk             }t	          j        ||           d S )Nr+   r   )r   r   r   r   @crb  z	a == "@c"r   rb  r%   r   r?   r   r@   rA   )r0   r!   r   r   r1   rE   r|   s          r   test_at_inside_stringz5TestDataFrameQueryNumExprPandas.test_at_inside_string  sn     (((===>??+fVDDbddl#
fh/////r   c                 t   | j         | j        }}t          |           t          t          j                            d                              d          t          d                    }t          j	        t          d          5  |                    d||           d d d            d S # 1 swxY w Y   d S )	Nr,   r   r   rm   r:  rM   za == @cr   )r!   r   r%   r   rr   r   r   ro   r#   rO   r   r?   rG  s       r   test_query_undefined_localz:TestDataFrameQueryNumExprPandas.test_query_undefined_local  s    dk (((ry,,Q//66w??dTTT]"*M
 
 
 	> 	> HHYvfH===	> 	> 	> 	> 	> 	> 	> 	> 	> 	> 	> 	> 	> 	> 	> 	> 	> 	>s   B--B14B1c                 @   d}t           j        ddd         }t          |t           j                            d                              |j                  d          }d|j        _        |	                    d||          }||d         d	k             }t          j        ||           t          |t           j                            d                              |j                  d          }|	                    d
||          }|j        |j        |j        d	k                      }t          j        ||           t          |t           j                            d                              |j                  d          }d|j        _        |	                    d||          }||j        d	k             }t          j        ||           |	                    d||          }|j        |j        |j        d	k                      }t          j        ||           d S )Nr+   r9  e   r,   )rj   r   rj   z	index > 5r   r   zilevel_0 > 5r   r   za > 5)rr   r_r   r   r   r   r   rj   r   r?   r@   rA   r  r   )r0   r!   r   rY   r   r1   rE   r|   s           r   :test_index_resolvers_come_after_columns_with_the_same_namezZTestDataFrameQueryNumExprPandas.test_index_resolvers_come_after_columns_with_the_same_name  s    E"S)bi33A66FFqvNNOO
 
  +fVDDbkAo&
fh///bi33A66FFqvNNOO
 
 .GG6"(28a<01
fh///QRY%:%:1%=%=%M%Maf%U%UVVWW'&@@bdQh<
fh///+fVDD6"(28a<01
fh/////r   zop, f==!=c                    d}t          t          j                            d                              |          t          j                            d                              |          d          }t          j        |j        d d ddf<   d| d}| ||j        t          j                           }|                    |||          }	t          j	        |	|           d S )Nr   r,   r   r   a z infr   )
r   rr   r   r   infr  r   r?   r@   rA   )
r0   rx   fr!   r   rY   r1   qr|   rE   s
             r   test_infz(TestDataFrameQueryNumExprPandas.test_inf  s    Y**1--44Q77Y**1--44Q77 
 
 sssAvMMMaabfoo&!F6::
fh/////r   c                    |}t          ddd|d          }t          |          }t          |          }|                    d          }t          j        ||           t          |          }|j                            t                    |_        |                                                    d          }t          j        ||           d S )Nz
2019-01-011dr   time)startfreqr   tzr   r   z"2018-01-03 00:00:00+00" < time)	r	   r   r?   r@   rA   rk   r!  r"  reset_index)r0   tz_aware_fixturerx  df_indexr|   r1   rE   s          r   test_check_tz_aware_index_queryz?TestDataFrameQueryNumExprPandas.test_check_tz_aware_index_query  s    T2"6
 
 
 8,,,X&&&;<<
fh///X&&#+226::!!''(IJJ
fh/////r   c           	         d}t          dt          j                            d                              |          z  t          j                            d                              |          d          }||d                             d          dk             }|                    d||          }t          j        ||           t          t          j        t          j                            d                              |          d	k     t          j	        t          j                            d          
                    |                    t          j                            d          
                    |          d          }||d                                                  }|                    d
||          }t          j        ||           d S )Nr   r,   r   r   intr   za.astype('int') == 0r   r   za.notnull())r   rr   r   r   r!  r?   r@   rA   wherer   r   notnull)r0   r!   r   rY   r1   r|   rE   s          r   test_method_calls_in_queryz:TestDataFrameQueryNumExprPandas.test_method_calls_in_query  s   ..q1188;;;Y**1--44Q77 
 
 bgnnU++q010OO
fh///XI))!,,33A66<FI))!,,<<Q?? 
 Y**1--==a@@ 	
 	
 bgoo''(-vFF
fh/////r   N)%rS   rT   rU   r#   rV   r!   r   r	  r  r  r  r#  r&  r0  r7  r;  r@  rC  rH  rL  rN  rR  rU  rZ  r]  r`  rd  rf  rj  r   r   operatoreqnerr  r|  r  r   r   r   r   r     s       ^  ^ ^  ^
* 
* 
** * *
* 
* 
*	* 	* 	** * ** * *N N N"< < <A A A4> > >	> 	> 	>
 
 
* * ** * *"0 0 0.= = =. . . * * *0 0 00 0 0> > >0 0 0> [Whk':T8;<O&PQQ0 0 RQ00 0 0 0 0 0 0 0r   r   c                   x    e Zd Zej        d             Zej        d             Zd Zd Zd Z	d Z
d Zd Zd	 Zd
S )TestDataFrameQueryNumExprPythonc                     dS r   r   r/   s    r   r!   z&TestDataFrameQueryNumExprPython.engine  r   r   c                     dS Nr   r   r/   s    r   r   z&TestDataFrameQueryNumExprPython.parser  r   r   c                    t          t          j                            d                              d                    }t          dd          |d<   t          dd          |d<   t          d	d          |d
<   |                    d||          }||j        dk     d|j        k     z           }t          j
        ||           d S )Nr,   r   r   r   r   r  r  r  r  r  )(dates1 < 20130101) & (20130101 < dates3)r   r  r  r  s         r   r  zCTestDataFrameQueryNumExprPython.test_date_query_no_attribute_access"  s    ry,,Q//??GGHH!*a8888!*a8888!*a8888hh7v  
 
 BI
*zBI/EFG
c5)))))r   c                    d}t          t          j                            d                              |df                    }t          d|          |d<   t          d|          |d<   t          d	|          |d
<   t          j        |j        t          j                            d                              |          dk    df<   t          j        |j        t          j                            d                              |          dk    d
f<   |	                    d||          }||j
        dk     d|j        k     z           }t          j        ||           d S )Nr   r,   r-   r   r   r  r  r  r  r  r   r  r   r  r  r  s          r   r  z8TestDataFrameQueryNumExprPython.test_date_query_with_NaT-  s;   ry,,Q//??AGGHH!*a8888!*a8888!*a8888EGVry$$Q''..q11C7ABEGVry$$Q''..q11C7ABhh7v  
 
 BI
*zBI/EFG
c5)))))r   c                    d}t          t          j                            d                              |df                    }t          d|          |d<   t          d|          |d<   |                    dd	d	
          }|J |                    d||          }||j        dk     d|j	        k     z           }t          j        ||           d S )Nr   r,   r-   r   r   r  r  r  Tr  ((index < 20130101) & (20130101 < dates3)r   r  r  r  s           r   r  z5TestDataFrameQueryNumExprPython.test_date_index_query;  s    ry,,Q//??AGGHH!*a8888!*a8888||Hd|FF###hh6vf  
 
 BHz)j29.DEF
c5)))))r   c                    d}t          t          j                            d                              |df                                        dt          i          }t          d|          |d<   t          d|          |d	<   t          j	        |j
        d
<   |                    ddd          }|J |                    d||          }||j        dk     d|j        k     z           }t          j        ||           d S )Nr   r,   r-   r   r   r   r  r  r  r  Tr  r  r   r  r   r  s           r   r#  z>TestDataFrameQueryNumExprPython.test_date_index_query_with_NaTH  s   ry,,Q//??AGGHHOOK
 
 "*a8888!*a8888||Hd|FF###hh6vf  
 
 BHz)j29.DEF
c5)))))r   c                 D   d}t          t          j                            d                              |df                    }t          d|          |d<   t          d|          |d<   t          j        |j        t          j                            d                              |          d	k    df<   |	                    dd
d
          }|J d}t          j        t          |          5  |                    d||           d d d            d S # 1 swxY w Y   d S )Nr   r,   r-   r   r   r  r  r  r   Tr  "'BoolOp' nodes are not implementedrM   r  r   )r   rr   r   r   r   r	   r   r  r  r  r#   rO   NotImplementedErrorr?   )r0   r!   r   rY   r1   r  rQ   s          r   r&  zITestDataFrameQueryNumExprPython.test_date_index_query_with_NaT_duplicatesY  sg   ry,,Q//??AGGHH!*a8888!*a8888EGVry$$Q''..q11C7AB||Hd|FF###3].c::: 	P 	PHH0HOOO	P 	P 	P 	P 	P 	P 	P 	P 	P 	P 	P 	P 	P 	P 	P 	P 	P 	Ps   /DDDc                    d}t          j        d||          }|dk    sJ t          t          j                            d                              d                    }t          t          j                            d                              d                    }d}t          j        t          |          5  |
                    d||           d d d            n# 1 swxY w Y   t          j        t          d	          5  |
                    d
||           d d d            n# 1 swxY w Y   ||dk    |dk    z           }t          j        d||          }t          j        ||           ||dk    |dk    z  ||dk             dk    z           }t          j        d||          }t          j        ||           d S )Nr+   zx + 1r   r,   r   z5The '@' prefix is only supported by the pandas parserrM   rP  rT  z(df>0) & (df2>0)r   zdf[(df > 0) & (df2 > 0)]z+df[(df > 0) & (df2 > 0) & (df[df > 0] > 0)])r   rB   r   rr   r   r   r   r#   rO   r6  r?   r   r@   rA   )	r0   r!   r   r   rE   r1   rQ  rQ   r|   s	            r   rR  z1TestDataFrameQueryNumExprPython.test_nested_scopee  sX   ???{{{{ry,,Q//??GGHH	--a00@@HHII G];c222 	I 	IHH)&HHHH	I 	I 	I 	I 	I 	I 	I 	I 	I 	I 	I 	I 	I 	I 	I ]19STTT 	G 	GHH'vHFFF	G 	G 	G 	G 	G 	G 	G 	G 	G 	G 	G 	G 	G 	G 	G rAv#'*+3F6RRR
h///rAv#'*baj1n=>9&QW
 
 
 	h/////s$   =C""C&)C&D--D14D1c                 \   t          g dg dd          }d}t          j        t          |          5  |                    d           d d d            n# 1 swxY w Y   d}t          j        t          |          5  |                    d           d d d            d S # 1 swxY w Y   d S )	Nr*   r   )minmaxzMVariables in expression \"\(min\) == \(1\)\" overlap with builtins: \('min'\)rM   zmin == 1zMVariables in expression \"\(max\) == \(1\)\" overlap with builtins: \('max'\)zmax == 1)r   r#   rO   r   r?   )r0   r1   regex_to_matchs      r   +test_query_numexpr_with_min_and_max_columnszKTestDataFrameQueryNumExprPython.test_query_numexpr_with_min_and_max_columns  s@   yyy;;<<0 	 ]1HHH 	! 	!HHZ   	! 	! 	! 	! 	! 	! 	! 	! 	! 	! 	! 	! 	! 	! 	!0 	 ]1HHH 	! 	!HHZ   	! 	! 	! 	! 	! 	! 	! 	! 	! 	! 	! 	! 	! 	! 	! 	! 	! 	!s#   AAA>B!!B%(B%N)rS   rT   rU   r#   rV   r!   r   r  r  r  r#  r&  rR  r  r   r   r   r  r    s        ^  ^ ^  ^	* 	* 	** * ** * ** * *"
P 
P 
P0 0 06! ! ! ! !r   r  c                   T    e Zd Zej        d             Zej        d             Zd ZdS )TestDataFrameQueryPythonPandasc                     dS r  r   r/   s    r   r!   z%TestDataFrameQueryPythonPandas.engine  r   r   c                     dS r   r   r/   s    r   r   z%TestDataFrameQueryPythonPandas.parser  r   r   c                 F   dx}}t          t          j                            d                              ||df          t          d                    }d|j        _        ||j        dk             }|                    d	||
          }t          j
        ||           d S Nr   r,   r-   r   r=  rm   rB  r   r>  r   r   rr   r   r   r?  ro   rj   r   r?   r@   rA   r0   r!   r   rY   rz   r1   r|   rE   s           r   rC  z1TestDataFrameQueryPythonPandas.test_query_builtin      
AI!!!$$--aq!f-==tE{{
 
 
 bhl#)F6BB
h/////r   NrS   rT   rU   r#   rV   r!   r   rC  r   r   r   r  r    [        ^  ^ ^  ^	0 	0 	0 	0 	0r   r  c                   T    e Zd Zej        d             Zej        d             Zd ZdS )TestDataFrameQueryPythonPythonc                     dS r  r   r/   s    r   r!   z%TestDataFrameQueryPythonPython.engine  r   r   c                     dS r  r   r/   s    r   r   z%TestDataFrameQueryPythonPython.parser  r   r   c                 F   dx}}t          t          j                            d                              ||df          t          d                    }d|j        _        ||j        dk             }|                    d	||
          }t          j
        ||           d S r  r  r  s           r   rC  z1TestDataFrameQueryPythonPython.test_query_builtin  r  r   Nr  r   r   r   r  r    r  r   r  c            
       2   e Zd Zd Zd Zd Zd Zd Zd Ze	j
                            ddej        gd	ej        gd
ej        gdej        gg          d             Zd Zd Ze	j
                            dg dg dg dg dg dg dg          d             ZdS )TestDataFrameQueryStringsc           	         t          t          j                            d                              d          dg          }t          t          d                    |d<   ||j        dk             }|dk    rd}d	}|gdz  |gdz  z   }|d d d
         }d\  }	}
d|	g|
gz   z  }d}t          |||          D ]c\  }}}| d| d| }t          j
        t          |          5  |                    |||d|j        i           d d d            n# 1 swxY w Y   dd S |                    d||          }t          j        ||           |                    d||          }t          j        ||           t          j        |||j                            dg                              ||j        dk             }|                    d||          }t          j        ||           |                    d||          }t          j        ||           t          j        |||j                            dg                               d S )Nr,   r   r+   r   rm   
aabbccddeestringsr   r   z"a"rl   rk  rl  $'(Not)?In' nodes are not implemented rM   )r!   r   
local_dictz"a" == stringsr   zstrings == "a"zstrings != "a"z"a" != strings)r   rr   r   r   r   r   ro   r  zipr#   rO   r  r?   r@   rA   r   r0   r   r!   r1   rY  collstlhsrhsr  r  opsrQ   rx   exr   s                   r   test_str_query_methodz/TestDataFrameQueryStrings.test_str_query_method  s   ry,,Q//??HHSVRWXXXtL11229BJ#%&XCC%!)seai'Cddd)CFBtrd{#C9C #Cc 2 2  R((b((3((]#6cBBB  HH%%$-rz#:	                    ((+F6(JJC!#v...((+F6(JJC!#v...!#r"*//3%*@*@'ABBB
c)*F((+F6(JJC!#v...((+F6(JJC!#v...!#r2:??C5+A+A*A'BCCCCCs   "!DD	D	c                    t          t          j                            d                              d          dg          }t          t          d                    |d<   ||j                            ddg                   }|dk    rd}d	}|gdz  |gdz  z   }|d d d
         }d\  }	}
d|	g|
gz   z  }d}t          |||          D ][\  }}}| d| d| }t          j        t          |          5  |                    |||           d d d            n# 1 swxY w Y   \d S |                    d||          }t          j        ||           |                    d||          }t          j        ||           ||j                            ddg                    }|                    d||          }t          j        ||           |                    d||          }t          j        ||           d S )Nr,   r  r   rm   r  r  r   r   z
["a", "b"]rl   r  r  r  rM   r   zstrings == ["a", "b"]z["a", "b"] == stringszstrings != ["a", "b"]z["a", "b"] != strings)r   rr   r   r   r   r   ro   r  r   r  r#   rO   r  r?   r@   rA   r  s                   r   test_str_list_query_methodz4TestDataFrameQueryStrings.test_str_list_query_method  sb   ry,,Q//??HHSVRWXXXtL11229BJOOS#J//0XCC%!)seai'Cddd)CFBtrd{#C9C #Cc 2 2 ? ?R((b((3((]#6cBBB ? ?HHRvH>>>? ? ? ? ? ? ? ? ? ? ? ? ? ? ?? ?
 ((26&(QQC!#v...((26&(QQC!#v...#s4445F((26&(QQC!#v...((26&(QQC!#v.....s   3DD	D	c           	         t          t          d          t          d          t          j                            d                              dd          t          j                            d                              dd          d          }|d	k    r|                    d
||          }||j                            |j	                           }t          j        ||           |                    d||          }||j                            |j	                  |j        |j        k     z           }t          j        ||           d S d}t          j        t           |          5  |                    d
||           d d d            n# 1 swxY w Y   d}t          j        t           |          5  |                    d||           d d d            d S # 1 swxY w Y   d S )Naaaabbbbccccaabbccddeeffr,   r      r   	   r   r   r   r%  r   za in br   za in b and c < dr  rM   r  )r   ro   rr   r   r   r?  r?   r   r   r   r@   rA   r   r%  r#   rO   r  )r0   r   r!   r1   r   r  rQ   s          r   test_query_with_string_columnsz8TestDataFrameQueryStrings.test_query_with_string_columns  si   .)).))Y**1--66qr6BBY**1--66qr6BB	 
 
 X((8F6(BBCrtyy'E!#u---((-fV(LLCrtyy"$+67E!#u-----9C2#>>> A A&@@@A A A A A A A A A A A A A A A 8C2#>>> K K+F6JJJK K K K K K K K K K K K K K K K K Ks$   -FFF:G  G$'G$c           	      6   t          t          d          t          d          t          j                            d                              dd          t          j                            d                              dd          d          }|                    d	||
          }||j        |j        k             }t          j
        ||           |                    d||
          }||j        |j        k             }t          j
        ||           d S )Nr  r  r,   r   r  r   r  r  za == br   za != b)r   ro   rr   r   r   r?  r?   r   r   r@   rA   )r0   r   r!   r1   r   r   s         r   test_object_array_eq_nez1TestDataFrameQueryStrings.test_object_array_eq_ne"  s    .)).))Y**1--66qr6BBY**1--66qr6BB	 
 
 hhxvh>>
c3'''hhxvh>>
c3'''''r   c                    t          |           d t          dd          D             dz  }t          ddd          }t          d	dd          }t          t	          j        dd
                              d          ||                    |          d          }||j        dk             }|	                    d||          }t          j        ||           d S )Nc                 (    g | ]}d D ]
}d| d| S ))loadexitzpage r  r   )r   rY   acts      r   r   zLTestDataFrameQueryStrings.test_query_with_nested_strings.<locals>.<listcomp>5  sK     
 
 
"#>N
 
7:A
 
 
 
r   r+   rZ   r,   z2014-01-01 0:00:0130sr   )rw  r   z2014-02-01 1:00:01   )idevent	timestampz"page 1 load"z'"page 1 load"' in eventr   )r%   rn   r	   r   rr   r   repeatappendr  r?   r@   rA   )	r0   r   r!   eventsstamps1stamps2r1   r|   r   s	            r   test_query_with_nested_stringsz8TestDataFrameQueryStrings.test_query_with_nested_strings3  s     (((
 
',Q{{
 
 
 1qIII1qIIIi1oo,,Q//$^^G44 
 
 bh/12hh5fVhTT
h,,,,,r   c                     t          |           t          g dg dd          }|                    d||          }||j        dk             }t	          j        ||           d S )N)r   r   test & testr*   r   za == "test & test"r   r  rc  )r0   r   r!   r1   r   r  s         r   (test_query_with_nested_special_characterzBTestDataFrameQueryStrings.test_query_with_nested_special_characterF  sr     (((666YYYGGHHhh+F6hJJ24=()
c5)))))r   zop, funcr*  r+  r,  r-  c                    t          t          j                            d                              t          d          d                    }t          t          j        |j                            }t          ||d          }|	                    d| d||          }| ||j
        d                   }	t          j        ||	           d S )	Nr,   abcder9  )Xr   zX z "d"r   r%  )r   rr   r   r   r   ro   r   r   r   r?   r  r@   rA   )
r0   r   r!   rx   funcr   r   r1   r   r|   s
             r   test_query_lex_compare_stringsz8TestDataFrameQueryStrings.test_query_lex_compare_stringsM  s     29((++224=="EEFF29QV$$%%QQ''((hh}B}}}VFhCCdd24oo&
c8,,,,,r   c                 f   d}t           j                            d                              ddt	          |          f                              t                    }t          ||          }|                    d||          }||j	        |j
        z           }t          j        ||           d S )N)bidbidsizeaskasksizer,   r+   r   rm   z	bid & askr   )rr   r   r   r?  r.  r!  boolr   r?   r  r  r@   rA   )r0   r   r!   rk   datar1   r   r|   s           r   "test_query_single_element_booleansz<TestDataFrameQueryStrings.test_query_single_element_booleans_  s    4y$$Q''00!S\\9J0KKRRSWXXtW---hh{6&hAAbfrvo&
c8,,,,,r   c                     t          |           t          g dg dd          }||j        dk             }d}|                    d||          }t	          j        ||           d S )N)BUD USr  IBM USr  )gl[@gGzn[@gf@g33333f@)SymbolPricer  zSymbol == @symbr   )r%   r   r  r?   r@   rA   )r0   r   r!   r1   esymbr5  s          r   !test_query_string_scalar_variablez;TestDataFrameQueryStrings.test_query_string_scalar_variableg  s     (((BBB999 
 
 ryH$%HH&vfHEE
a#####r   in_list)Nasdfghjk)r  Nr  )r  r  N)NNr  )r  NN)NNNc                    d}d}d t          |          D             }t          d|id          }|j                            d          |_        t          d|id          }|                    d||	          }||d         d
k             }|                    d||	          }	t          j        ||           t          j        ||           t          j        ||	           t          j        ||	           d S )Nr   r   c                 &    i | ]\  }}|d k    ||S )r  r   )r   r4  values      r   
<dictcomp>zMTestDataFrameQueryStrings.test_query_string_null_elements.<locals>.<dictcomp>  s#    SSSE5F??Au???r   r   stringdtypeint64za == 'asdf'r   r  za <= 'asdf')	enumerater   rj   r!  r?   r@   rA   )
r0   r  r   r!   r|   df_expectedr1   r   r   res3s
             r   test_query_string_null_elementsz9TestDataFrameQueryStrings.test_query_string_null_elementst  s    SSYw-?-?SSShx@@@'-44W==W~X666xxfVxDD"S'V#$xxfVxDD
dK000
dD)))
dD)))
dD)))))r   N)rS   rT   rU   r  r  r  r  r  r  r#   r   r   r  ltgtleger  r  r  r  r   r   r   r  r    sb       'D 'D 'DR!/ !/ !/FK K K4( ( ("- - -&* * * [(+(+8;8;		
 - - -- - -$ $ $ ["""""""""      	

 
* *
 
* * *r   r  c                       e Zd Zej        d             Zd Zd Zej        	                    dg d          d             Z
dS )TestDataFrameEvalWithFramec                     t          t          j                            d                              d          t          d                    S )Nr,   r   r=  rm   )r   rr   r   r   r   ro   r/   s    r   framez TestDataFrameEvalWithFrame.frame  sB    I!!!$$44W==tE{{
 
 
 	
r   c                 ~    |                     d||          }|j        |j        z   }t          j        ||           d S )Nr   r   rB   r   r   r@   rC   r0   r  r   r!   r   rY  s         r   test_simple_exprz+TestDataFrameEvalWithFrame.test_simple_expr  s@    jjj??57"
sF+++++r   c                     |                     d||          }|j        |j        dk              |j        z   }t          j        ||           d S )Nza[a < 1] + br   r+   r  r  s         r   test_bool_arith_exprz/TestDataFrameEvalWithFrame.test_bool_arith_expr  sK    jjvjFF1%/
sF+++++r   rx   )r[   r^   ra   rd   c                     t          ddgddgd          }d}t          j        t          |          5  |                    d| d	||
           d d d            d S # 1 swxY w Y   d S )Nr+   r,   r   r%  r   z:unsupported operand type\(s\) for .+: '.+' and '.+'|CannotrM   rn  z br   )r   r#   rO   r/  rB   )r0   r   r!   rx   r1   rQ   s         r   %test_invalid_type_for_operator_raisesz@TestDataFrameEvalWithFrame.test_invalid_type_for_operator_raises  s    aV3*5566K]9C000 	? 	?GGKKKKvG>>>	? 	? 	? 	? 	? 	? 	? 	? 	? 	? 	? 	? 	? 	? 	? 	? 	? 	?s   AA"%A"N)rS   rT   rU   r#   rV   r  r	  r  r   r   r  r   r   r   r  r    s        ^
 
 ^

, , ,
, , ,
 [T#7#7#788? ? 98? ? ?r   r  c                      e Zd Zej        d             Zd Zd Zd Zd Z	d Z
d Zd Zd	 Zd
 Zd Zd Zd Zd Zd Zd Zd Zd Zd Zd Zd Zd Zd Zd Zd Zd Zej                             dg d          d             Z!ej                             ddd g          ej                             dg d          d!                         Z"d" Z#d#S )$!TestDataFrameQueryBacktickQuotingc           	   #      K   t          i dg ddg ddg ddg dd	g d
dg ddg ddg ddg ddg ddg ddg ddg ddg ddg ddg d d!g d"g d#g d$g d%g d&g d'd(          V  d)S )*z
        Yields a dataframe with strings that may or may not need escaping
        by backticks. The last two columns cannot be escaped by backticks
        and should raise a ValueError.
        r)   r*   B B)r-   r,   r+   C Cr   C  C)r  rZ   r-   C_C)   r  r   D_D D)   r+   rh  E.E)r   r-   r   F-F)r  r+   r   1e1)r,   rZ   r  def)r   r  r,   A (x))rZ   r+   r-   zB(x))r+   r+   r   zB (x))r,   r  rZ        &^ :!€$?(} >    <++*''  )r,   r   r   r   )r   r  r+    A)rZ   r  r    r+   r,   r+   )r   r-   r+   )r  r+   r  )r  r  r   )r,   rZ   r   )r   r  r  )zit'szthat'su   ☺zfoo#barr+   Nr.   r/   s    r   r1   z$TestDataFrameQueryBacktickQuoting.df  ss      YYYyyy yyy 				
 zzz  yyy zzz yyy {{{  			  /			 KKK  iii!" iii#$ "		#)) yy$9999-  
 
 	
 	
 	
 	
 	
r   c                 ~    |                     d          }|d|d         k              }t          j        ||           d S )Nz	1 < `B B`r+   r  r?   r@   rA   r0   r1   r   rY  s       r   #test_single_backtick_variable_queryzETestDataFrameQueryBacktickQuoting.test_single_backtick_variable_query  s>    hh{##A5	M"
c6*****r   c                     |                     d          }|d|d         k     d|d         k     z           }t          j        ||           d S )Nz1 < `B B` and 4 < `C C`r+   r  rZ   r  r"  r#  s       r   !test_two_backtick_variables_queryzCTestDataFrameQueryBacktickQuoting.test_two_backtick_variables_query  sL    hh011QE]q2e9}56
c6*****r   c                 |    |                     d          }|d         |d         z   }t          j        ||           d S )Nz	A + `B B`r)   r  rB   r@   rC   r#  s       r   "test_single_backtick_variable_exprzDTestDataFrameQueryBacktickQuoting.test_single_backtick_variable_expr  >    ggk""C2e9$
sF+++++r   c                 |    |                     d          }|d         |d         z   }t          j        ||           d S )Nz`B B` + `C C`r  r  r(  r#  s       r    test_two_backtick_variables_exprzBTestDataFrameQueryBacktickQuoting.test_two_backtick_variables_expr  s>    ggo&&ERY&
sF+++++r   c                 |    |                     d          }|d         |d         z   }t          j        ||           d S )Nz	`C_C` + Ar  r)   r(  r#  s       r    test_already_underscore_variablezBTestDataFrameQueryBacktickQuoting.test_already_underscore_variable  s>    ggk""ERW$
sF+++++r   c                 |    |                     d          }|d         |d         z   }t          j        ||           d S )NzC_C + `C C`r  r  r(  r#  s       r   test_same_name_but_underscoresz@TestDataFrameQueryBacktickQuoting.test_same_name_but_underscores  s>    ggm$$ERY&
sF+++++r   c                 |    |                     d          }|d         |d         z   }t          j        ||           d S )NzA + `D_D D`r)   r  r(  r#  s       r   !test_mixed_underscores_and_spaceszCTestDataFrameQueryBacktickQuoting.test_mixed_underscores_and_spaces  s>    ggm$$C2g;&
sF+++++r   c                 |    |                     d          }|d         |d         z   }t          j        ||           d S )Nz	A + `C_C`r)   r  r(  r#  s       r   'test_backtick_quote_name_with_no_spaceszITestDataFrameQueryBacktickQuoting.test_backtick_quote_name_with_no_spaces  r*  r   c                     |                     d          }|d         |d         z   |d         z
  }t          j        ||           d S )Nz`E.E` + `F-F` - Ar  r  r)   r(  r#  s       r   test_special_charactersz9TestDataFrameQueryBacktickQuoting.test_special_characters  sH    gg)**ERY&C0
sF+++++r   c                 |    |                     d          }|d         |d         z   }t          j        ||           d S )Nz	A + `1e1`r)   r  r(  r#  s       r   test_start_with_digitz7TestDataFrameQueryBacktickQuoting.test_start_with_digit  r*  r   c                 |    |                     d          }|d         |d         z   }t          j        ||           d S )Nz	A + `def`r)   r  r(  r#  s       r   test_keywordz.TestDataFrameQueryBacktickQuoting.test_keyword   r*  r   c                 ~    |                     d          }||d         dk             }t          j        ||           d S )Nz`A` > 2r)   r,   r"  r#  s       r   test_unneeded_quotingz7TestDataFrameQueryBacktickQuoting.test_unneeded_quoting  s=    hhy!!BsGaK
c6*****r   c                 ~    |                     d          }||d         dk             }t          j        ||           d S )Nz`A (x)` > 2r  r,   r"  r#  s       r   test_parenthesisz2TestDataFrameQueryBacktickQuoting.test_parenthesis
  s>    hh}%%BwK!O$
c6*****r   c                 ~    |                     d          }||d         dk             }t          j        ||           d S )Nz`` > 5r   r   r"  r#  s       r   test_empty_stringz3TestDataFrameQueryBacktickQuoting.test_empty_string  s=    hhx  BrFQJ
c6*****r   c                 ~    |                     d          }||d         dk             }t          j        ||           d S )Nz
`C  C` > 5r  r   r"  r#  s       r   test_multiple_spacesz6TestDataFrameQueryBacktickQuoting.test_multiple_spaces  s>    hh|$$BvJN#
c6*****r   c                 |    |                     d          }|d         |d         z   }t          j        ||           d S )Nz` A` + `  `r  r  r(  r#  s       r   test_start_with_spacesz8TestDataFrameQueryBacktickQuoting.test_start_with_spaces  s>    ggm$$DBtH$
sF+++++r   c                 ~    |                     d          }||d         dk             }t          j        ||           d S )Nu"   `  &^ :!€$?(} >    <++*''  ` > 4r  rZ   r"  r#  s       r   test_lots_of_operators_stringz?TestDataFrameQueryBacktickQuoting.test_lots_of_operators_string  sB    hh;<<B56:;
c6*****r   c                     d}t          j        t          |          5  |                    d           d d d            d S # 1 swxY w Y   d S )Nz(module 'pandas' has no attribute 'thing'rM   z	@pd.thing)r#   rO   AttributeErrorrB   )r0   r1   messages      r   test_missing_attributez8TestDataFrameQueryBacktickQuoting.test_missing_attribute#  s    <]>999 	! 	!GGK   	! 	! 	! 	! 	! 	! 	! 	! 	! 	! 	! 	! 	! 	! 	! 	! 	! 	!   AAAc                     d}t          j        t          |          5  |                    d           d d d            d S # 1 swxY w Y   d S )N6(Could not convert ).*( to a valid Python identifier.)rM   z`it's` > `that's`r#   rO   r6  r?   r   s      r   test_failing_quotez4TestDataFrameQueryBacktickQuoting.test_failing_quote(  s    G];c222 	* 	*HH()))	* 	* 	* 	* 	* 	* 	* 	* 	* 	* 	* 	* 	* 	* 	* 	* 	* 	*rK  c                     d}t          j        t          |          5  |                    d           d d d            d S # 1 swxY w Y   d S )NrM  rM   u	   `☺` > 4rN  r   s      r   $test_failing_character_outside_rangezFTestDataFrameQueryBacktickQuoting.test_failing_character_outside_range-  s    G];c222 	" 	"HH[!!!	" 	" 	" 	" 	" 	" 	" 	" 	" 	" 	" 	" 	" 	" 	" 	" 	" 	"rK  c                     d}t          j        t          |          5  |                    d           d d d            d S # 1 swxY w Y   d S )NzFailed to parse backticksrM   z`foo#bar` > 4rN  r   s      r   test_failing_hashtagz6TestDataFrameQueryBacktickQuoting.test_failing_hashtag2  s    )];c222 	& 	&HH_%%%	& 	& 	& 	& 	& 	& 	& 	& 	& 	& 	& 	& 	& 	& 	& 	& 	& 	&rK  c                 ^   d }|g}|                     d           t          j        t          d          5  |                     d           ddd           n# 1 swxY w Y   t          j        t          d          5  |                     d           ddd           dS # 1 swxY w Y   dS )z
        Only attributes and variables ('named functions') can be called.
        .__call__() is not an allowed attribute because that would allow
        calling anything.
        https://github.com/pandas-dev/pandas/pull/32460
        c                      dS r7   r   )r   s    r   r  zNTestDataFrameQueryBacktickQuoting.test_call_non_named_expression.<locals>.func?  s    1r   z@func()z"Only named functions are supportedrM   z@funcs[0]()Nz@funcs[0].__call__())rB   r#   rO   r/  )r0   r1   r  funcss       r   test_call_non_named_expressionz@TestDataFrameQueryBacktickQuoting.test_call_non_named_expression7  s8   	 	 	 
	]9,PQQQ 	# 	#GGM"""	# 	# 	# 	# 	# 	# 	# 	# 	# 	# 	# 	# 	# 	# 	# ]9,PQQQ 	, 	,GG*+++	, 	, 	, 	, 	, 	, 	, 	, 	, 	, 	, 	, 	, 	, 	, 	, 	, 	,s#   AA A?B""B&)B&c                 :   t          ddgddggddg|          }t          rt          nd }t          j        |          5  |                    d          }d d d            n# 1 swxY w Y   t          g d	g d
gg d|          }t          j        ||           d S )Nr+   r,   r-   rZ   r   r   rk   r  z	c = b - ar   )r-   rZ   r+   rE  r   r
   RuntimeWarningr@   assert_produces_warningrB   rA   )r0   any_numeric_ea_and_arrow_dtyper1   warningrE   r|   s         r   test_ea_dtypesz0TestDataFrameQueryBacktickQuoting.test_ea_dtypesL  s    VaVsCj8V
 
 
 %6?..4'00 	* 	*WW[))F	* 	* 	* 	* 	* 	* 	* 	* 	* 	* 	* 	* 	* 	* 	*YY			"#OO0
 
 

 	fh/////   A!!A%(A%c                 :   t          ddgddggddgd          }t          rt          nd }t          j        |          5  |                    d	          }d d d            n# 1 swxY w Y   t          g d
g dgg dd          }t          j        ||           d S )Nr+   r,   r-   rZ   r   r   Float64rY  z	c = b - 1r   )r-   rZ   r-   rE  rZ  )r0   r1   r^  rE   r|   s        r   test_ea_dtypes_and_scalarz;TestDataFrameQueryBacktickQuoting.test_ea_dtypes_and_scalar[  s    AA'#s9MMM$5?..4'00 	* 	*WW[))F	* 	* 	* 	* 	* 	* 	* 	* 	* 	* 	* 	* 	* 	* 	*YY			"OOO9
 
 
 	fh/////r`  c           	      0   t          ddgddggddg|          }|                    d          }t          t          ddg|	          t          ddg|	          t          ddg|d
         j        	          d          }t	          j        ||           d S )Nr+   r,   r-   rZ   r   r   rY  z	c = 2 - 1r  r   rE  )r   rB   r   r  r@   rA   )r0   r]  r1   rE   r|   s        r   #test_ea_dtypes_and_scalar_operationzETestDataFrameQueryBacktickQuoting.test_ea_dtypes_and_scalar_operationf  s    VaVsCj8V
 
 
 %%QF*HIIIQF*HIIIQF&+*;<<< 
 
 	fh/////r   r  )r  Int64int64[pyarrow]c                    |dk    rt          j        d           t          dt          ddg|          i          }dh}|dk    rt          rt
          nd }t          j        |          5  |                    d          }d d d            n# 1 swxY w Y   t          dt          dg|dg	          i          }t          j	        ||           d S )
Nrg  pyarrowr   r+   r,   r  rf  z	a in @refr  rj   )
r#   importorskipr   r   r
   r[  r@   r\  r?   rA   )r0   r  r1   refr^  rE   r|   s          r   test_query_ea_dtypesz6TestDataFrameQueryBacktickQuoting.test_query_ea_dtypesu  s   $$$	***VQF%8889::c$)W$4$49J$4..PT'00 	+ 	+XXk**F	+ 	+ 	+ 	+ 	+ 	+ 	+ 	+ 	+ 	+ 	+ 	+ 	+ 	+ 	+c61#U1##F#F#FGHH
fh/////s   *BBBr!   r   r   c           	      &   |dk    rt           nd }|dk    rt          st          j        d           |dk    rt          j        d           t          t          g dd          t          g d|          d	          }t          j        |          5  |	                    d
|          }d d d            n# 1 swxY w Y   t          t          ddgdddg          t          ddg|ddg          d	          }t          j
        ||           d S )Nr   znumexpr not installedrg  ri  )r+   r+   r,   rf  r  )r+   r,   r,   r   zA == BrH   r+   r,   r   rj  )r[  r
   r#   r$   rk  r   r   r@   r\  r?   rA   )r0   r  r!   r^  r1   rE   r|   s          r   !test_query_ea_equality_comparisonzCTestDataFrameQueryBacktickQuoting.test_query_ea_equality_comparison  sx    %+i$7$7..TY'8K/000$$$	***'222			QV9W9W9WXX
 
 '00 	7 	7XXhvX66F	7 	7 	7 	7 	7 	7 	7 	7 	7 	7 	7 	7 	7 	7 	7QF'!Q@@@QF%1v>>> 
 
 	fh/////s   B77B;>B;c                 $   t           j                            d          }t          dt          j        d d gd          it
                    }|                    d          }t          dg it
                    }t          j        ||           d S )NUTCr   T)utcr  za > @now)	r   	Timestampnowr   to_datetimer"  r?   r@   rA   )r0   rt  r1   rE   r|   s        r   test_all_nat_in_objectz8TestDataFrameQueryBacktickQuoting.test_all_nat_in_object  s    lu%%R^T4LdCCCDFSSS*%%c2Yf555
fh/////r   N)$rS   rT   rU   r#   rV   r1   r$  r&  r)  r,  r.  r0  r2  r4  r6  r8  r:  r<  r>  r@  rB  rD  rF  rJ  rO  rQ  rS  rW  r_  rc  re  r   r   rm  ro  rv  r   r   r   r  r    sD       ^
 
 ^
B+ + +
+ + +
, , ,
, , ,
, , ,
, , ,
, , ,
, , ,
, , ,
, , ,
, , ,
+ + +
+ + +
+ + +
+ + +
, , ,
+ + +
! ! !
* * *
" " "
& & &
, , ,*0 0 0	0 	0 	00 0 0 [W&J&J&JKK
0 
0 LK
0 [X)'<==[W&J&J&JKK0 0 LK >=0(0 0 0 0 0r   r  )'r  numpyrr   r#   pandas.errorsr   r   pandas.util._test_decoratorsutil_test_decoratorstdr   r   r   r   r   r   r	   pandas._testing_testingr@   pandas.core.computation.checkr
   rV   r   r   
skip_if_nor!   r%   r'   rX   r   r   r  r  r  r  r  r  r   r   r   <module>r     s                * ) ) ) ) ) ) ) )                        ; ; ; ; ; ; (+===  >= lfl9MBM)4L4LMMMN   	 = = =
01 01 01 01 01 01 01 01fs1 s1 s1 s1 s1 s1 s1 s1lGI GI GI GI GI GI GI GIT yw0 w0 w0 w0 w0 w0 w0 w0t yu! u! u! u! u!&E u! u! u!p0 0 0 0 0%D 0 0 0*0 0 0 0 0%D 0 0 0*S* S* S* S* S* S* S* S*l? ? ? ? ? ? ? ?4s0 s0 s0 s0 s0 s0 s0 s0 s0 s0r   