Filter column and row according to regex in header

Question

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)

Answer 1

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])

=> ['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).

Answer 2

I see that you are using count while iterating over headers:

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

but enumerate is much nicer for keeping track of count:

cols=[]
print(header[0], end='')
for count,token in enumerate(header):
    if prog.match(token):
        #add the column indices to a list for use later
        cols.append(count)
        print("\t",token, end='')

Then at the bottom when printing the line:

        print(tokens[0], end='')
        for col in cols:
            print("\t",tokens[col], end='')
        print()

you treat tokens[0] separately then the others, now if 0 was in cols then you could instead use .join()

        print('\t'.join(tokens[col] for col in cols))

to use it like that you just have to initialize cols as:

cols = [0] #instead of cols = []

then you can do the same for the headers:

cols=[0]
for count,token in enumerate(header):
    if prog.match(token):
        #add the column indices to a list for use later
        cols.append(count)
print("\t".join(header[col] for col in cols))

The last thing I could see was the use of while line!="": instead of for line in tsv: since files will iterate over lines in python (starting from their current reading position so nothing screws up the headers stuff above):

for line in tsv:
    if line.strip()=="":
        continue
    tokens=line.split()
    if prog.match(tokens[0]):
        print(col_sep.join(tokens[col] for col in cols))

This also means line doesn't have to be defined at the beginning so you could replace:

line=tsv.readline()
header=line.split("\t")

with:

header = tsv.readline().split("\t")

but I can't find ways to reduce the for and if statements without storing entire file in memory, other people could probably be more help on that front, anyway here is the final code from my suggestions:

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
    header=tsv.readline().split("\t")
    cols = [0]
    for count,token in enumerate(header):
        #add the column indices to a list for use later
        if prog.match(token):
            cols.append(count)
    print("\t".join(header[col] for col in cols))
    #in the rest of the line, find the matching rows and 
    #print the columns that match the indices found above
    for line in tsv:
        if line.strip()=="":
            continue
        tokens=line.split()
        if prog.match(tokens[0]):
            print("\t".join(tokens[col] for col in cols))

Answer 3

Just want to add a few things to @DavidMorris great answer.

• Use the csv module to handle character separated values. It will handle the work of producing a list for each row of the input file and can even handle "weird" values (such as values spaning on multiple lines).
• Use a file-object as an iterator. Especially if you actually want to process lines one at a time. This is even more true when using csv as the cleaning of a line is already performed by this module.
• Create a function. It will let you separate the input data from their processing and will help you simplify handling command line arguments (with argparse, getopt or even non-standard modules like docopt). You will also want to put the call to this function under if __name__ == '__main__'.
• Follow PEP8, not that your code is hardly readable, it's more like a few nitpicks: each import statement on its own line and constant spelled in UPPER_SNAKE_CASE.

Putting all that together with @DavidMorris advices lead to:

import re
import sys
import csv
from itertools import compress

def filter_tsv(in_file, out_file, pattern):
    reader = csv.reader(in_file, delimiter='\t')
    writer = csv.writer(out_file, delimiter='\t')
    header = next(reader)
    interest = [pattern.search(column) for column in header]
    # Keep the first column no matter what
    interest[0] = True

    writer.writerow(compress(header, interest))
    for line in reader:
        writer.writerow(compress(line, interest))

if __name__ == '__main__':
    # Put your stuff to handle command-line arguments in here
    with open("input.tsv") as tsv:
        filter_tsv(tsv, sys.stdout, re.compile("[a|c]"))