Filter column and row according to regex in header I'm writing a script that can grep the rows and columns from a text file that match a regex. I have an implementation but I'm looking to make it more elegant.

input.tsv




        a       b       c       d
a       2       9       6       1
b       6       4       7       0
c       3       7       2       0
d       5       8       2       2



I want to do something like this:


$ grep-row-column.py -E "[a|c]" input.tsv
         a       c
a        2       6
c        3       2



The solution should be able to stream a file (i.e. not read the entire thing into memory).



from __future__ import print_function
import sys,re

filename="input.tsv"
regex="[a|c]"

prog=re.compile(".*"+regex+".*")
with open(filename) as tsv:
    #find the appropriate columns in the header that match the regex
    line=tsv.readline()
    header=line.split("\t")
    count=0
    cols=[]
    print(header[0], end='')
    for token in header:
        if prog.match(token):
            #add the column indices to a list for use later
            cols.append(count)
            print("\t",token, end='')
        count+=1
    print()
    #in the rest of the line, find the matching rows and 
    #print the columns that match the indices found above
    while line!="":
        line=tsv.readline()
        if line.strip()=="":
            continue
        tokens=line.split()
        if prog.match(tokens[0]):
            print(tokens[0], end='')
            for col in cols:
                print("\t",tokens[col], end='')
            print()


Is there a more 'pythonic' way to do this? I couldn't figure out how to avoid all the for loops and if statements.

(Cross-posted from here)
Your approach is basically reasonable, but you aren't taking advantages of Python's more powerful features, or any of the more interesting parts of the standard library.

re.search

A minor point to start with. You're making a regex like so:

prog=re.compile(".*"+regex+".*")


Presumably so that you get an unanchored match (i.e., one which can start at any point in the string).

re has a method re.search which will do unanchored matching by default; you can anchor the match to the start of the string with ^, as you would expect. See search() vs. match() in the documentation.

Manipulating whole data structures

There are a number of specific points below, but the broad thrust of them is that you will have a much nicer time doing large-scale operations on whole data structures, rather than fiddling around with indexes and individual elements. This is true in most languages, but one of the strengths of Python is the rich set of data structures and operations which are available built-in or in the standard library.

itertools.compress

There is a wonderful module in the standard library called itertools. This library is my first thought in Python whenever I want to do something involving sequences of any kind, and it contains a little function called compress which will do a lot of the work for your program.

Basically, you give compress a sequence of data items and a sequence of selectors, and you get back a sequence of those data items for which the selector is true. So in your example, matching your regex against the columns would give you the "selector":

[False, True, False, True, False]


(I'll mention how to generate this simply in a moment). You would then want to modify this to also pass through the row header, so you have the selector:

[True, True, False, True, False]


then, doing itertools.compress for a randomly-selected row:

itertools.compress(['a', 2, 9, 6, 1], [True, True, False, True, False])

=&gt; ['a', 2, 6]


you can then just print that list to print the (filtered) row.

Note that you could also do this filtering yourself, using a list comprehension; however, itertools has the twin advantages of being common (hence readable) and carefully optimized (hence fast).

join

Once you can easily generate lists containing exactly the elements you want to output, you don't have to fiddle around doing print('\t', tokens[col], end='') and similar; you can just do:

print('\t'.join(tokens))


(assuming you have already filtered your tokens using compress or similar).

List comprehensions

List comprehensions are one of Python's most pleasant features. They're far too large a topic for me to sensibly explain in an answer here; that link will give you a good introduction, and there are many other articles and introductions around.

Roughly speaking, list comprehensions (and their cousins, generator expressions) are a nice syntax for creating sequences from sequences, doing "something" (transforming, selecting, etc.) on the way.

For example, you could write a list comprehension like this:

[reg.search(t) for t in tokens]


If reg is [a|c], and tokens is ['', 'a', 'b', 'c', 'd'], this will produce:

[False, True, False, True, False]


which you will recall is exactly what you need to feed to compress to filter to just those columns. (Actually, it will produce a list of Nones and regex match objects, but those evaluate to False/True in a boolean context, so don't worry too much about that).