/tools/stats/filtering.py

https://bitbucket.org/cistrome/cistrome-harvard/ · Python · 276 lines · 220 code · 31 blank · 25 comment · 51 complexity · 5fc3570c46b488c81789837c3f542c5f MD5 · raw file

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