/tools/stats/filtering.py

  1#!/usr/bin/env python
  2# This tool takes a tab-delimited text file as input and creates filters on columns based on certain properties.
  3# The tool will skip over invalid lines within the file, informing the user about the number of lines skipped.
  4
  5from __future__ import division
  6import sys, re, os.path
  7from galaxy import eggs
  8
  9from ast import parse, Module, walk
 10
 11# Older py compatibility
 12try:
 13    set()
 14except:
 15    from sets import Set as set
 16
 17AST_NODE_TYPE_WHITELIST = [
 18    'Expr',    'Load',
 19    'Str',     'Num',    'BoolOp',
 20    'Compare', 'And',    'Eq',
 21    'NotEq',   'Or',     'GtE',
 22    'LtE',     'Lt',     'Gt',
 23    'BinOp',   'Add',    'Div',
 24    'Sub',     'Mult',   'Mod',
 25    'Pow',     'LShift', 'GShift',
 26    'BitAnd',  'BitOr',  'BitXor',
 27    'UnaryOp', 'Invert', 'Not',
 28    'NotIn',   'In',     'Is',
 29    'IsNot',   'List',
 30    'Index',   'Subscript',
 31    # Further checks
 32    'Name',    'Call',    'Attribute',
 33]
 34
 35
 36BUILTIN_AND_MATH_FUNCTIONS = 'abs|all|any|bin|chr|cmp|complex|divmod|float|hex|int|len|long|max|min|oct|ord|pow|range|reversed|round|sorted|str|sum|type|unichr|unicode|log|exp|sqrt|ceil|floor'.split('|')
 37STRING_AND_LIST_METHODS = [ name for name in dir('') + dir([]) if not name.startswith('_') ]
 38VALID_FUNCTIONS = BUILTIN_AND_MATH_FUNCTIONS + STRING_AND_LIST_METHODS
 39
 40
 41def __check_name( ast_node ):
 42    name = ast_node.id
 43    if re.match(r'^c\d+$', name):
 44        return True
 45    return name in VALID_FUNCTIONS
 46
 47
 48def __check_attribute( ast_node ):
 49    attribute_name = ast_node.attr
 50    if attribute_name not in STRING_AND_LIST_METHODS:
 51        return False
 52    return True
 53
 54
 55def __check_call( ast_node ):
 56    # If we are calling a function or method, it better be a math,
 57    # string or list function.
 58    ast_func = ast_node.func
 59    ast_func_class = ast_func.__class__.__name__
 60    if ast_func_class == 'Name':
 61        if ast_func.id not in BUILTIN_AND_MATH_FUNCTIONS:
 62            return False
 63    elif ast_func_class == 'Attribute':
 64        if not __check_attribute( ast_func ):
 65            return False
 66    else:
 67        return False
 68
 69    return True
 70
 71
 72def check_expression( text ):
 73    """
 74
 75    >>> check_expression("c1=='chr1' and c3-c2>=2000 and c6=='+'")
 76    True
 77    >>> check_expression("eval('1+1')")
 78    False
 79    >>> check_expression("import sys")
 80    False
 81    >>> check_expression("[].__str__")
 82    False
 83    >>> check_expression("__builtins__")
 84    False
 85    >>> check_expression("'x' in globals")
 86    False
 87    >>> check_expression("'x' in [1,2,3]")
 88    True
 89    >>> check_expression("c3=='chr1' and c5>5")
 90    True
 91    >>> check_expression("c3=='chr1' and d5>5")  # Invalid d5 reference
 92    False
 93    >>> check_expression("c3=='chr1' and c5>5 or exec")
 94    False
 95    >>> check_expression("type(c1) != type(1)")
 96    True
 97    >>> check_expression("c1.split(',')[1] == '1'")
 98    True
 99    >>> check_expression("exec 1")
100    False
101    >>> check_expression("str(c2) in [\\\"a\\\",\\\"b\\\"]")
102    True
103    """
104    try:
105        module = parse( text )
106    except SyntaxError:
107        return False
108
109    if not isinstance(module, Module):
110        return False
111    statements = module.body
112    if not len( statements ) == 1:
113        return False
114    expression = statements[0]
115    if expression.__class__.__name__ != 'Expr':
116        return False
117
118    for ast_node in walk( expression ):
119        ast_node_class = ast_node.__class__.__name__
120
121        # Toss out everything that is not a "simple" expression,
122        # imports, error handling, etc...
123        if ast_node_class not in AST_NODE_TYPE_WHITELIST:
124            return False
125
126        # White-list more potentially dangerous types AST elements.
127        if ast_node_class == 'Name':
128            # In order to prevent loading 'exec', 'eval', etc...
129            # put string restriction on names allowed.
130            if not __check_name( ast_node ):
131                return False
132        # Check only valid, white-listed functions are called.
133        elif ast_node_class == 'Call':
134            if not __check_call( ast_node ):
135                return False
136        # Check only valid, white-listed attributes are accessed
137        elif ast_node_class == 'Attribute':
138            if not __check_attribute( ast_node ):
139                return False
140
141    return True
142
143def get_operands( filter_condition ):
144    # Note that the order of all_operators is important
145    items_to_strip = ['+', '-', '**', '*', '//', '/', '%', '<<', '>>', '&', '|', '^', '~', '<=', '<', '>=', '>', '==', '!=', '<>', ' and ', ' or ', ' not ', ' is ', ' is not ', ' in ', ' not in ']
146    for item in items_to_strip:
147        if filter_condition.find( item ) >= 0:
148            filter_condition = filter_condition.replace( item, ' ' )
149    operands = set( filter_condition.split( ' ' ) )
150    return operands
151
152def stop_err( msg ):
153    sys.stderr.write( msg )
154    sys.exit()
155
156in_fname = sys.argv[1]
157out_fname = sys.argv[2]
158cond_text = sys.argv[3]
159try:
160    in_columns = int( sys.argv[4] )
161    assert sys.argv[5]  #check to see that the column types variable isn't null
162    in_column_types = sys.argv[5].split( ',' )
163except:
164    stop_err( "Data does not appear to be tabular.  This tool can only be used with tab-delimited data." )
165num_header_lines = int( sys.argv[6] )
166
167# Unescape if input has been escaped
168mapped_str = {
169    '__lt__': '<',
170    '__le__': '<=',
171    '__eq__': '==',
172    '__ne__': '!=',
173    '__gt__': '>',
174    '__ge__': '>=',
175    '__sq__': '\'',
176    '__dq__': '"',
177    '__ob__': '[',
178    '__cb__': ']',
179}
180for key, value in mapped_str.items():
181    cond_text = cond_text.replace( key, value )
182    
183# Attempt to determine if the condition includes executable stuff and, if so, exit
184secured = dir()
185operands = get_operands(cond_text)
186for operand in operands:
187    try:
188        check = int( operand )
189    except:
190        if operand in secured:
191            stop_err( "Illegal value '%s' in condition '%s'" % ( operand, cond_text ) )
192
193if not check_expression(cond_text):
194    stop_err( "Illegal/invalid in condition '%s'" % ( cond_text ) )
195
196# Work out which columns are used in the filter (save using 1 based counting)
197used_cols = sorted(set(int(match.group()[1:]) \
198                   for match in re.finditer('c(\d)+', cond_text))) 
199largest_col_index = max(used_cols)
200
201# Prepare the column variable names and wrappers for column data types. Only 
202# cast columns used in the filter.
203cols, type_casts = [], []
204for col in range( 1, largest_col_index + 1 ):
205    col_name = "c%d" % col
206    cols.append( col_name )
207    col_type = in_column_types[ col - 1 ]
208    if col in used_cols:
209        type_cast = "%s(%s)" % ( col_type, col_name )
210    else:
211        #If we don't use this column, don't cast it.
212        #Otherwise we get errors on things like optional integer columns.
213        type_cast = col_name
214    type_casts.append( type_cast )
215 
216col_str = ', '.join( cols )    # 'c1, c2, c3, c4'
217type_cast_str = ', '.join( type_casts )  # 'str(c1), int(c2), int(c3), str(c4)'
218assign = "%s, = line.split( '\\t' )[:%i]" % ( col_str, largest_col_index )
219wrap = "%s = %s" % ( col_str, type_cast_str )
220skipped_lines = 0
221invalid_lines = 0
222first_invalid_line = 0
223invalid_line = None
224lines_kept = 0
225total_lines = 0
226out = open( out_fname, 'wt' )
227    
228# Read and filter input file, skipping invalid lines
229code = '''
230for i, line in enumerate( file( in_fname ) ):
231    total_lines += 1
232    line = line.rstrip( '\\r\\n' )
233
234    if i < num_header_lines:
235        lines_kept += 1
236        print >> out, line
237        continue
238
239    if not line or line.startswith( '#' ):
240        skipped_lines += 1
241        continue
242    try:
243        %s
244        %s
245        if %s:
246            lines_kept += 1
247            print >> out, line
248    except:
249        invalid_lines += 1
250        if not invalid_line:
251            first_invalid_line = i + 1
252            invalid_line = line
253''' % ( assign, wrap, cond_text )
254valid_filter = True
255try:
256    exec code
257except Exception, e:
258    out.close()
259    if str( e ).startswith( 'invalid syntax' ):
260        valid_filter = False
261        stop_err( 'Filter condition "%s" likely invalid. See tool tips, syntax and examples.' % cond_text )
262    else:
263        stop_err( str( e ) )
264
265if valid_filter:
266    out.close()
267    valid_lines = total_lines - skipped_lines
268    print 'Filtering with %s, ' % cond_text
269    if valid_lines > 0:
270        print 'kept %4.2f%% of %d valid lines (%d total lines).' % ( 100.0*lines_kept/valid_lines, valid_lines, total_lines )
271    else:
272        print 'Possible invalid filter condition "%s" or non-existent column referenced. See tool tips, syntax and examples.' % cond_text
273    if invalid_lines:
274        print 'Skipped %d invalid line(s) starting at line #%d: "%s"' % ( invalid_lines, first_invalid_line, invalid_line )
275    if skipped_lines:
276        print 'Skipped %i comment (starting with #) or blank line(s)' % skipped_lines