Add read/write sequence interface with support for compressed sequences

[huddlej]

Description of proposed changes

This PR adds a new io.py module with three new functions:

• open_file: a context manager that transparently supports reading/writing:
  • compressed files using xopen to infer compression from the filename
  • file paths or file handles
• read_sequences: a function that streams sequence records from one or more input files of the same format (FASTA, GenBank, etc.). Uses open_file to support compressed inputs.
• write_sequences: a function that streams sequence records to a single filename or file handle. Uses open_file to support compressed output.

This PR also refactors the following Augur subcommands to use these new functions:

• parse
• index
• filter
• mask

Python API example

# Read sequences from multiple files with a generator.
sequences = read_sequences(*args.sequences)

# Open a file to write to. If `args.output` ends with
# ".gz", for example, its contents will compressed.
observed_sequence_strains = set()
with open_file(args.output, "wt") as output_handle:
    for sequence in sequences:
        # Track all the strains we've written.
        observed_sequence_strains.add(sequence.id)

        # Write one record at a time to the handle.
        write_sequences(sequence, output_handle)

Command line interface examples

The following command reads in H3N2 HA sequences from a gzip-compressed file and writes out the parsed sequences to an LZMA-compressed file.

augur parse \
  --sequences h3n2_ha.fasta.gz \
  --output-sequences sequences.fasta.xz \
  --output-metadata metadata.tsv \
  --fields strain virus accession date region country division location passage originating_lab submitting_lab age gender

The following command reads LZMA-compressed sequences from the previous command and writes out a gzip-compressed sequence index.

augur index \
    --sequences sequences.fasta.xz \
    --output sequence_index.tsv.gz

Related issue(s)

See the ZenHub Epic for a list of all related issues.
Fixes #644
See #637 for details about how Augur reads and write sequences.
See #645 for the original proposed interface for the new read/write functions.

Testing

This PR adds functional and unit tests for all new and refactored code.

[codecov]

Codecov Report

Merging #652 (bbf963f) into master (8df4b4d) will increase coverage by 1.05%.
The diff coverage is 77.21%.

@@            Coverage Diff             @@
##           master     #652      +/-   ##
==========================================
+ Coverage   30.52%   31.57%   +1.05%     
==========================================
  Files          40       41       +1     
  Lines        5615     5779     +164     
  Branches     1363     1436      +73     
==========================================
+ Hits         1714     1825     +111     
- Misses       3830     3848      +18     
- Partials       71      106      +35     

Impacted Files	Coverage Δ
augur/index.py	80.70% <33.33%> (-1.45%)	⬇️
augur/parse.py	62.50% <47.82%> (+12.50%)	⬆️
augur/filter.py	48.78% <100.00%> (+2.93%)	⬆️
augur/io.py	100.00% <100.00%> (ø)
augur/mask.py	100.00% <100.00%> (ø)
augur/utils.py	37.76% <0.00%> (-0.61%)	⬇️

---

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update 8df4b4d...bbf963f. Read the comment docs.

[kairstenfay]

By request, I took a pass at reviewing your patches here. They're perhaps not the most insightful comments but here's what I have for today :)

[huddlej]

Thank you for the review, @kairstenfay! Do you have any general thoughts on the read/write interfaces? The more I've thought about it, the more I don't like passing the file handle around to write_sequences. These two functions should probably be classes like a SequenceReader and SequenceWriter. This implementation would allow the latter class to track the handle as an attribute. It could also allow the sequence reader to provide fancier functions like streaming out unique sequences, etc.

[kairstenfay]

Thank you for the review, @kairstenfay! Do you have any general thoughts on the read/write interfaces? The more I've thought about it, the more I don't like passing the file handle around to write_sequences. These two functions should probably be classes like a SequenceReader and SequenceWriter. This implementation would allow the latter class to track the handle as an attribute. It could also allow the sequence reader to provide fancier functions like streaming out unique sequences, etc.

I don't have a lot of context or experience using augur, so I hesitate to comment on the interfaces. However, I did think it was a bit odd to pass around the file handler arguments to the write_sequences function. With my superficial understanding, creating SequenceReader & SequenceWriter classes sounds like a good approach to me.

[rneher]

This looks good to me John. I installed this on our cluster. Builds on Monday will use it (and I'll pass in the raw gz files at the top). A bunch of ncov scripts would need adjusting, so I am not changing everything to compressed just yet.

[huddlej]

Thank you for trying this out, @rneher. I'm interested to know what issues you encounter. :)

From the engineering perspective, I'd to make a quick attempt at implementing these two new functions as classes, as discussed with @kairstenfay above. Since my hope is to make this part of the minimal public-facing Python API that we've discussed, I want to make sure this interface is as good as it can be from the start.

The last little lift to finishing this work is to update all augur commands that read or write sequences to use this new interface. I'd prefer not to make this change piece-meal, so users don't get different experiences with different commands.

[rneher]

I merged master into it on scicore and I am running the builds locally w/o problem. But I have only tried compressed files in filter.

[jameshadfield]

Thanks @huddlej -- I spent some time looking through this and like the interface.

We have an open_file function in utils which we no longer use, should we add a deprecation warning to this?

Running nCoV with (xz) compressed metadata worked 💯. This actually surprised me as it means read_metadata handles xz files already!

---

P.S. for future work when we allow augur tree to take compressed input, be aware that VCF inputs, which need a fasta reference, use the TreeTime helper function read_vcf rather than our own helper functions.

[huddlej]

Thank you for looking through this, @jameshadfield!

We have an open_file function in utils which we no longer use, should we add a deprecation warning to this?

The old open_file function was added last March as part of a mask.py refactor and it has only ever been used in that module. We still use this older function in mask.py to open VCF files.

The old function has some functionality that the new function does not including:

1. a check for opening gzipped files in "text" mode. The xopen package implements the "text" mode check already, so we don't need to implement that anymore.
2. an explicit request for UTF-8 encoding. Based on @tsibley's excellent commit message when this UTF-8 encoding was added, we should continue to use this as an explicit default and still allow the calling code to override the default.

As far as how we handle this older function, I'd prefer to replace all internal references to that function with the new function now and then handle deprecation/migration of these types of I/O functions from utils.py to io.py in a later PR.

Running nCoV with (xz) compressed metadata worked 💯. This actually surprised me as it means read_metadata handles xz files already!

Yeah! We get this functionality "for free" by using pandas to read in the metadata. pandas has an interesting I/O library, too, where they effectively implement their own version of the xopen module (I would guess the pandas code predates the xopen module even though xopen was inspired by Heng Li's function of the same name from 2014!).

for future work when we allow augur tree to take compressed input, be aware that VCF inputs, which need a fasta reference, use the TreeTime helper function read_vcf rather than our own helper functions.

That's a good point. I'd vote to implement our own wrapper around any TreeTime functionality for this kind of I/O, so we can more easily replace/update the backend in the future.

[huddlej]

It turns out xopen does not support passing arguments like encoding through to the internal open function call. This may not be an issue for sequence data, since BioPython appears to do its own string decoding, but it could be an issue for other data types. I'm not sure if this is a deal-breaker or not...