#!/usr/bin/env python # This tool takes a tab-delimited text file as input and creates filters on columns based on certain properties. # The tool will skip over invalid lines within the file, informing the user about the number of lines skipped. from __future__ import ( division, print_function, ) import json import re import sys from ast import ( parse, walk, ) AST_NODE_TYPE_WHITELIST = [ "Expr", "Load", "Str", "Num", "BoolOp", "Compare", "And", "Eq", "NotEq", "Or", "GtE", "LtE", "Lt", "Gt", "BinOp", "Add", "Div", "Sub", "Mult", "Mod", "Pow", "LShift", "GShift", "BitAnd", "BitOr", "BitXor", "UnaryOp", "Invert", "Not", "UAdd", "USub", "NotIn", "In", "Is", "IsNot", "List", "Index", "Subscript", "Constant", # Further checks "Name", "Call", "Attribute", ] 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( "|" ) STRING_AND_LIST_METHODS = [name for name in dir("") + dir([]) if not name.startswith("_")] VALID_FUNCTIONS = BUILTIN_AND_MATH_FUNCTIONS + STRING_AND_LIST_METHODS def __check_name(ast_node): name = ast_node.id if re.match(r"^c\d+$", name): return True return name in VALID_FUNCTIONS def __check_attribute(ast_node): attribute_name = ast_node.attr if attribute_name not in STRING_AND_LIST_METHODS: return False return True def __check_call(ast_node): # If we are calling a function or method, it better be a math, # string or list function. ast_func = ast_node.func ast_func_class = ast_func.__class__.__name__ if ast_func_class == "Name": if ast_func.id not in BUILTIN_AND_MATH_FUNCTIONS: return False elif ast_func_class == "Attribute": if not __check_attribute(ast_func): return False else: return False return True def check_expression(text): """ >>> check_expression("c1=='chr1' and c3-c2>=2000 and c6=='+'") True >>> check_expression("eval('1+1')") False >>> check_expression("import sys") False >>> check_expression("[].__str__") False >>> check_expression("__builtins__") False >>> check_expression("'x' in globals") False >>> check_expression("'x' in [1,2,3]") True >>> check_expression("c3=='chr1' and c5>5") True >>> check_expression("c3=='chr1' and c5>-5 and c5<+5") # Unary +/- True >>> check_expression("c3=='chr1' and d5>5") # Invalid d5 reference False >>> check_expression("c3=='chr1' and c5>5 or exec") False >>> check_expression("type(c1) != type(1)") True >>> check_expression("c1.split(',')[1] == '1'") True >>> check_expression("exec 1") False >>> check_expression("str(c2) in [\\\"a\\\",\\\"b\\\"]") True """ try: module = parse(text) except SyntaxError: return False statements = module.body if not len(statements) == 1: return False expression = statements[0] if expression.__class__.__name__ != "Expr": return False for ast_node in walk(expression): ast_node_class = ast_node.__class__.__name__ # Toss out everything that is not a "simple" expression, # imports, error handling, etc... if ast_node_class not in AST_NODE_TYPE_WHITELIST: return False # White-list more potentially dangerous types AST elements. if ast_node_class == "Name": # In order to prevent loading 'exec', 'eval', etc... # put string restriction on names allowed. if not __check_name(ast_node): return False # Check only valid, white-listed functions are called. elif ast_node_class == "Call": if not __check_call(ast_node): return False # Check only valid, white-listed attributes are accessed elif ast_node_class == "Attribute": if not __check_attribute(ast_node): return False return True def get_operands(filter_condition): # Note that the order of all_operators is important items_to_strip = [ "+", "-", "**", "*", "//", "/", "%", "<<", ">>", "&", "|", "^", "~", "<=", "<", ">=", ">", "==", "!=", "<>", " and ", " or ", " not ", " is ", " is not ", " in ", " not in ", ] for item in items_to_strip: if filter_condition.find(item) >= 0: filter_condition = filter_condition.replace(item, " ") operands = set(filter_condition.split(" ")) return operands def stop_err(msg): sys.stderr.write(msg) sys.exit() in_fname = sys.argv[1] out_fname = sys.argv[2] with open(sys.argv[3]) as f: inputs = json.load(f) cond_text = inputs["cond"] try: in_columns = int(sys.argv[4]) assert sys.argv[5] # check to see that the column types variable isn't null in_column_types = sys.argv[5].split(",") except Exception: stop_err("Data does not appear to be tabular. This tool can only be used with tab-delimited data.") num_header_lines = int(sys.argv[6]) # Unescape if input has been escaped mapped_str = { "__lt__": "<", "__le__": "<=", "__eq__": "==", "__ne__": "!=", "__gt__": ">", "__ge__": ">=", "__sq__": "'", "__dq__": '"', "__ob__": "[", "__cb__": "]", } for key, value in mapped_str.items(): cond_text = cond_text.replace(key, value) # Attempt to determine if the condition includes executable stuff and, if so, exit secured = dir() operands = get_operands(cond_text) for operand in operands: try: check = int(operand) except ValueError: if operand in secured: stop_err("Illegal value '%s' in condition '%s'" % (operand, cond_text)) if not check_expression(cond_text): stop_err("Illegal/invalid in condition '%s'" % (cond_text)) # Work out which columns are used in the filter (save using 1 based counting) used_cols = sorted({int(match.group()[1:]) for match in re.finditer(r"c(\d)+", cond_text)}) largest_col_index = max(used_cols) # Prepare the column variable names and wrappers for column data types. Only # cast columns used in the filter. cols, type_casts = [], [] for col in range(1, largest_col_index + 1): col_name = "c%d" % col cols.append(col_name) col_type = in_column_types[col - 1] if col in used_cols: type_cast = "%s(%s)" % (col_type, col_name) else: # If we don't use this column, don't cast it. # Otherwise we get errors on things like optional integer columns. type_cast = col_name type_casts.append(type_cast) col_str = ", ".join(cols) # 'c1, c2, c3, c4' type_cast_str = ", ".join(type_casts) # 'str(c1), int(c2), int(c3), str(c4)' assign = "%s, = line.split( '\\t' )[:%i]" % (col_str, largest_col_index) wrap = "%s = %s" % (col_str, type_cast_str) skipped_lines = 0 invalid_lines = 0 first_invalid_line = 0 invalid_line = None lines_kept = 0 total_lines = 0 out = open(out_fname, "w") # Read and filter input file, skipping invalid lines code = """ for i, line in enumerate( open( in_fname ) ): total_lines += 1 line = line.rstrip( '\\r\\n' ) if i < num_header_lines: lines_kept += 1 print( line, file=out ) continue if not line or line.startswith( '#' ): skipped_lines += 1 continue try: %s %s if %s: lines_kept += 1 print( line, file=out ) except Exception: invalid_lines += 1 if not invalid_line: first_invalid_line = i + 1 invalid_line = line """ % ( assign, wrap, cond_text, ) valid_filter = True try: exec(code) except Exception as e: out.close() if str(e).startswith("invalid syntax"): valid_filter = False stop_err('Filter condition "%s" likely invalid. See tool tips, syntax and examples.' % cond_text) else: stop_err(str(e)) if valid_filter: out.close() valid_lines = total_lines - skipped_lines print("Filtering with %s, " % cond_text) if valid_lines > 0: print( "kept %4.2f%% of %d valid lines (%d total lines)." % (100.0 * lines_kept / valid_lines, valid_lines, total_lines) ) else: print( 'Possible invalid filter condition "%s" or non-existent column referenced. See tool tips, syntax and examples.' % cond_text ) if invalid_lines: print( 'Skipped %d invalid line(s) starting at line #%d: "%s"' % (invalid_lines, first_invalid_line, invalid_line) ) if skipped_lines: print("Skipped %i comment (starting with #) or blank line(s)" % skipped_lines)