<s> """<STR_LIT>""" <EOL> from __future__ import absolute_import <EOL> import re <EOL> from numexpr . necompiler import typecode_to_kind <EOL> from numexpr . necompiler import expressionToAST , typeCompileAst <EOL> from numexpr . necompiler import stringToExpression , NumExpr , getExprNames <EOL> from numexpr . expressions import ExpressionNode <EOL> from . utilsextension import get_nested_field <EOL> from . utils import lazyattr <EOL> from six . moves import zip <EOL> _no_matching_opcode = re . compile ( r"<STR_LIT>" ) <EOL> def _unsupported_operation_error ( exception ) : <EOL> """<STR_LIT>""" <EOL> message = exception . args [ <NUM_LIT:0> ] <EOL> op , types = _no_matching_opcode . search ( message ) . groups ( ) <EOL> newmessage = "<STR_LIT>" % op <EOL> newmessage += '<STR_LIT:U+002CU+0020>' . join ( [ typecode_to_kind [ t ] for t in types [ <NUM_LIT:1> : ] ] ) <EOL> return exception . __class__ ( newmessage ) <EOL> def _check_indexable_cmp ( getidxcmp ) : <EOL> """<STR_LIT>""" <EOL> def newfunc ( exprnode , indexedcols ) : <EOL> result = getidxcmp ( exprnode , indexedcols ) <EOL> if result [ <NUM_LIT:0> ] is not None : <EOL> try : <EOL> typeCompileAst ( expressionToAST ( exprnode ) ) <EOL> except NotImplementedError as nie : <EOL> raise _unsupported_operation_error ( nie ) <EOL> return result <EOL> newfunc . __name__ = getidxcmp . __name__ <EOL> newfunc . __doc__ = getidxcmp . __doc__ <EOL> return newfunc <EOL> @ _check_indexable_cmp <EOL> def _get_indexable_cmp ( exprnode , indexedcols ) : <EOL> """<STR_LIT>""" <EOL> not_indexable = ( None , None , None ) <EOL> turncmp = { '<STR_LIT>' : '<STR_LIT>' , <EOL> '<STR_LIT>' : '<STR_LIT>' , <EOL> '<STR_LIT>' : '<STR_LIT>' , <EOL> '<STR_LIT>' : '<STR_LIT>' , <EOL> '<STR_LIT>' : '<STR_LIT>' , } <EOL> def get_cmp ( var , const , op ) : <EOL> var_value , const_value = var . value , const . value <EOL> if ( var . astType == '<STR_LIT>' and var_value in indexedcols <EOL> and const . astType in [ '<STR_LIT>' , '<STR_LIT>' ] ) : <EOL> if const . astType == '<STR_LIT>' : <EOL> const_value = ( const_value , ) <EOL> return ( var_value , op , const_value ) <EOL> return None <EOL> def is_indexed_boolean ( node ) : <EOL> return ( node . astType == '<STR_LIT>' <EOL> and node . astKind == '<STR_LIT:bool>' <EOL> and node . value in indexedcols ) <EOL> if is_indexed_boolean ( exprnode ) : <EOL> return ( exprnode . value , '<STR_LIT>' , True ) <EOL> if exprnode . astType == '<STR_LIT>' and exprnode . value == '<STR_LIT>' : <EOL> child = exprnode . children [ <NUM_LIT:0> ] <EOL> if is_indexed_boolean ( child ) : <EOL> return ( child . value , '<STR_LIT>' , False ) <EOL> if child . astKind == "<STR_LIT:bool>" : <EOL> return ( child , '<STR_LIT>' , None ) <EOL> if exprnode . astType != '<STR_LIT>' : <EOL> return not_indexable <EOL> cmpop = exprnode . value <EOL> if cmpop not in turncmp : <EOL> return not_indexable <EOL> left , right = exprnode . children <EOL> cmp_ = get_cmp ( left , right , cmpop ) <EOL> if cmp_ : <EOL> return cmp_ <EOL> cmp_ = get_cmp ( right , left , turncmp [ cmpop ] ) <EOL> if cmp_ : <EOL> return cmp_ <EOL> return not_indexable <EOL> def _equiv_expr_node ( x , y ) : <EOL> """<STR_LIT>""" <EOL> if not isinstance ( x , ExpressionNode ) and not isinstance ( y , ExpressionNode ) : <EOL> return x == y <EOL> elif ( type ( x ) is not type ( y ) or not isinstance ( x , ExpressionNode ) <EOL> or not isinstance ( y , ExpressionNode ) <EOL> or x . value != y . value or x . astKind != y . astKind <EOL> or len ( x . children ) != len ( y . children ) ) : <EOL> return False <EOL> for xchild , ychild in zip ( x . children , y . children ) : <EOL> if not _equiv_expr_node ( xchild , ychild ) : <EOL> return False <EOL> return True <EOL> def _get_idx_expr_recurse ( exprnode , indexedcols , idxexprs , strexpr ) : <EOL> """<STR_LIT>""" <EOL> not_indexable = ( [ ] , [ '<STR_LIT>' ] ) <EOL> op_conv = { <EOL> '<STR_LIT>' : '<STR_LIT:&>' , <EOL> '<STR_LIT>' : '<STR_LIT:|>' , <EOL> '<STR_LIT>' : '<STR_LIT>' , <EOL> } <EOL> negcmp = { <EOL> '<STR_LIT>' : '<STR_LIT>' , <EOL> '<STR_LIT>' : '<STR_LIT>' , <EOL> '<STR_LIT>' : '<STR_LIT>' , <EOL> '<STR_LIT>' : '<STR_LIT>' , <EOL> } <EOL> def fix_invert ( idxcmp , exprnode , indexedcols ) : <EOL> invert = False <EOL> while idxcmp [ <NUM_LIT:1> ] == "<STR_LIT>" : <EOL> invert ^= True <EOL> exprnode = idxcmp [ <NUM_LIT:0> ] <EOL> idxcmp = _get_indexable_cmp ( exprnode , indexedcols ) <EOL> return idxcmp , exprnode , invert <EOL> idxcmp = _get_indexable_cmp ( exprnode , indexedcols ) <EOL> idxcmp , exprnode , invert = fix_invert ( idxcmp , exprnode , indexedcols ) <EOL> if idxcmp [ <NUM_LIT:0> ] : <EOL> if invert : <EOL> var , op , value = idxcmp <EOL> if op == '<STR_LIT>' and value in [ True , False ] : <EOL> value ^= True <EOL> else : <EOL> op = negcmp [ op ] <EOL> expr = ( var , ( op , ) , ( value , ) ) <EOL> invert = False <EOL> else : <EOL> expr = ( idxcmp [ <NUM_LIT:0> ] , ( idxcmp [ <NUM_LIT:1> ] , ) , ( idxcmp [ <NUM_LIT:2> ] , ) ) <EOL> return [ expr ] <EOL> if invert : <EOL> return not_indexable <EOL> if exprnode . astType != '<STR_LIT>' or exprnode . value not in [ '<STR_LIT>' , '<STR_LIT>' ] : <EOL> return not_indexable <EOL> left , right = exprnode . children <EOL> lcolvar , lop , llim = _get_indexable_cmp ( left , indexedcols ) <EOL> rcolvar , rop , rlim = _get_indexable_cmp ( right , indexedcols ) <EOL> op = exprnode . value <EOL> if ( lcolvar is not None and rcolvar is not None <EOL> and _equiv_expr_node ( lcolvar , rcolvar ) and op == '<STR_LIT>' ) : <EOL> if lop in [ '<STR_LIT>' , '<STR_LIT>' ] and rop in [ '<STR_LIT>' , '<STR_LIT>' ] : <EOL> expr = ( lcolvar , ( lop , rop ) , ( llim , rlim ) ) <EOL> return [ expr ] <EOL> if lop in [ '<STR_LIT>' , '<STR_LIT>' ] and rop in [ '<STR_LIT>' , '<STR_LIT>' ] : <EOL> expr = ( rcolvar , ( rop , lop ) , ( rlim , llim ) ) <EOL> return [ expr ] <EOL> lexpr = _get_idx_expr_recurse ( left , indexedcols , idxexprs , strexpr ) <EOL> rexpr = _get_idx_expr_recurse ( right , indexedcols , idxexprs , strexpr ) <EOL> def add_expr ( expr , idxexprs , strexpr ) : <EOL> """<STR_LIT>""" <EOL> if isinstance ( expr , list ) : <EOL> idxexprs . append ( expr [ <NUM_LIT:0> ] ) <EOL> lenexprs = len ( idxexprs ) <EOL> if lenexprs == <NUM_LIT:1> : <EOL> strexpr [ : ] = [ "<STR_LIT>" ] <EOL> else : <EOL> strexpr [ : ] = [ <EOL> "<STR_LIT>" % ( strexpr [ <NUM_LIT:0> ] , op_conv [ op ] , lenexprs - <NUM_LIT:1> ) ] <EOL> if lexpr != not_indexable and ( op == "<STR_LIT>" or rexpr != not_indexable ) : <EOL> add_expr ( lexpr , idxexprs , strexpr ) <EOL> if rexpr != not_indexable : <EOL> add_expr ( rexpr , idxexprs , strexpr ) <EOL> return ( idxexprs , strexpr ) <EOL> if rexpr != not_indexable and op == "<STR_LIT>" : <EOL> add_expr ( rexpr , idxexprs , strexpr ) <EOL> return ( idxexprs , strexpr ) <EOL> return not_indexable <EOL> def _get_idx_expr ( expr , indexedcols ) : <EOL> """<STR_LIT>""" <EOL> return _get_idx_expr_recurse ( expr , indexedcols , [ ] , [ '<STR_LIT>' ] ) <EOL> class CompiledCondition ( object ) : <EOL> """<STR_LIT>""" <EOL> @ lazyattr <EOL> def index_variables ( self ) : <EOL> """<STR_LIT>""" <EOL> idxexprs = self . index_expressions <EOL> idxvars = [ ] <EOL> for expr in idxexprs : <EOL> idxvar = expr [ <NUM_LIT:0> ] <EOL> if idxvar not in idxvars : <EOL> idxvars . append ( idxvar ) <EOL> return frozenset ( idxvars ) <EOL> def __init__ ( self , func , params , idxexprs , strexpr , ** kwargs ) : <EOL> self . function = func <EOL> """<STR_LIT>""" <EOL> self . parameters = params <EOL> """<STR_LIT>""" <EOL> self . index_expressions = idxexprs <EOL> """<STR_LIT>""" <EOL> self . string_expression = strexpr <EOL> """<STR_LIT>""" <EOL> self . kwargs = kwargs <EOL> """<STR_LIT>""" <EOL> def __repr__ ( self ) : <EOL> return ( "<STR_LIT>" <EOL> % ( self . index_expressions , self . string_expression , <EOL> self . index_variables ) ) <EOL> def with_replaced_vars ( self , condvars ) : <EOL> """<STR_LIT>""" <EOL> exprs = self . index_expressions <EOL> exprs2 = [ ] <EOL> for expr in exprs : <EOL> idxlims = expr [ <NUM_LIT:2> ] <EOL> limit_values = [ ] <EOL> for idxlim in idxlims : <EOL> if isinstance ( idxlim , tuple ) : <EOL> idxlim = condvars [ idxlim [ <NUM_LIT:0> ] ] <EOL> idxlim = idxlim . tolist ( ) <EOL> limit_values . append ( idxlim ) <EOL> var , ops , _ = expr <EOL> exprs2 . append ( ( var , ops , tuple ( limit_values ) ) ) <EOL> newcc = CompiledCondition ( <EOL> self . function , self . parameters , exprs2 , self . string_expression , <EOL> ** self . kwargs ) <EOL> return newcc <EOL> def _get_variable_names ( expression ) : <EOL> """<STR_LIT>""" <EOL> names = [ ] <EOL> stack = [ expression ] <EOL> while stack : <EOL> node = stack . pop ( ) <EOL> if node . astType == '<STR_LIT>' : <EOL> names . append ( node . value ) <EOL> elif hasattr ( node , '<STR_LIT>' ) : <EOL> stack . extend ( node . children ) <EOL> return list ( set ( names ) ) <EOL> def compile_condition ( condition , typemap , indexedcols ) : <EOL> """<STR_LIT>""" <EOL> expr = stringToExpression ( condition , typemap , { } ) <EOL> if expr . astKind != '<STR_LIT:bool>' : <EOL> raise TypeError ( "<STR_LIT>" <EOL> % condition ) <EOL> idxexprs = _get_idx_expr ( expr , indexedcols ) <EOL> if isinstance ( idxexprs , list ) : <EOL> strexpr = [ '<STR_LIT>' ] <EOL> else : <EOL> idxexprs , strexpr = idxexprs <EOL> strexpr = strexpr [ <NUM_LIT:0> ] <EOL> varnames = _get_variable_names ( expr ) <EOL> signature = [ ( var , typemap [ var ] ) for var in varnames ] <EOL> try : <EOL> func = NumExpr ( expr , signature ) <EOL> except NotImplementedError as nie : <EOL> raise _unsupported_operation_error ( nie ) <EOL> _ , ex_uses_vml = getExprNames ( condition , { } ) <EOL> kwargs = { '<STR_LIT>' : ex_uses_vml } <EOL> params = varnames <EOL> return CompiledCondition ( func , params , idxexprs , strexpr , ** kwargs ) <EOL> def call_on_recarr ( func , params , recarr , param2arg = None , ** kwargs ) : <EOL> """<STR_LIT>""" <EOL> args = [ ] <EOL> for param in params : <EOL> if param2arg : <EOL> arg = param2arg ( param ) <EOL> else : <EOL> arg = param <EOL> if hasattr ( arg , '<STR_LIT>' ) : <EOL> arg = get_nested_field ( recarr , arg . pathname ) <EOL> args . append ( arg ) <EOL> return func ( * args , ** kwargs ) </s>
