-
Notifications
You must be signed in to change notification settings - Fork 128
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Adds tests and code for new `open_file`, `read_sequences`, and `write_sequences` functions loosely based on a proposed API [1]. These functions transparently handle compressed inputs and outputs using the xopen library. The `open_file` function is a context manager that lightly wraps the `xopen` function and also supports either path strings or existing IO buffers. Both the read and write functions use this context manager to open files. This manager enables the common use case of writing to the same handle many times inside a for loop, by replacing the standard `open` call with `open_file`. Doing so, we maintain a Pythonic interface that also supports compressed file formats and path-or-buffer inputs. This context manager also enables input and output of any other file type in compressed formats (e.g., metadata, sequence indices, etc.). Note that the `read_sequences` and `write_sequences` functions do not infer the format of sequence files (e.g., FASTA, GenBank, etc.). Inferring file formats requires peeking at the first record in each given input, but peeking is not supported by piped inputs that we want to support (e.g., piped gzip inputs from xopen). There are also no internal use cases for Augur to read multiple sequences of different formats, so I can't currently justify the complexity required to support type inference. Instead, I opted for the same approach used by BioPython where the calling code must know the type of input file being passed. This isn't an unreasonable expectation for Augur's internal code. I also considered inferring file type by filename extensions like xopen infers compression modes. Filename extensions are less standardized across bioinformatics than we would like for this type of inference to work robustly. Tests ignore BioPython and pycov warnings to minimize warning fatigue for issues we cannot address during test-driven development. [1] #645
- Loading branch information
Showing
4 changed files
with
292 additions
and
1 deletion.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,106 @@ | ||
#!/usr/bin/env python3 | ||
"""Interfaces for reading and writing data also known as input/output (I/O) | ||
""" | ||
import Bio.SeqIO | ||
import Bio.SeqRecord | ||
from contextlib import contextmanager | ||
from pathlib import Path | ||
from xopen import xopen | ||
|
||
|
||
@contextmanager | ||
def open_file(path_or_buffer, mode="r", **kwargs): | ||
"""Opens a given file path and returns the handle. | ||
Transparently handles compressed inputs and outputs. | ||
Parameters | ||
---------- | ||
path_or_buffer : str or Path-like or IO buffer | ||
Name of the file to open or an existing IO buffer | ||
mode : str | ||
Mode to open file (read or write) | ||
Returns | ||
------- | ||
IO | ||
File handle object | ||
""" | ||
try: | ||
with xopen(path_or_buffer, mode, **kwargs) as handle: | ||
yield handle | ||
except TypeError: | ||
yield path_or_buffer | ||
|
||
|
||
def read_sequences(*paths, format="fasta"): | ||
"""Read sequences from one or more paths. | ||
Automatically infer compression mode (e.g., gzip, etc.) and return a stream | ||
of sequence records in the requested format (e.g., "fasta", "genbank", etc.). | ||
Parameters | ||
---------- | ||
paths : list of str or Path-like objects | ||
One or more paths to sequence files of any type supported by BioPython. | ||
format : str | ||
Format of input sequences matching any of those supported by BioPython | ||
(e.g., "fasta", "genbank", etc.). | ||
Yields | ||
------ | ||
Bio.SeqRecord.SeqRecord | ||
Sequence record from the given path(s). | ||
""" | ||
for path in paths: | ||
# Open the given path as a handle, inferring the file's compression. | ||
# This way we can pass a handle to BioPython's SeqIO interface | ||
# regardless of the compression mode. | ||
with open_file(path) as handle: | ||
sequences = Bio.SeqIO.parse(handle, format) | ||
|
||
for sequence in sequences: | ||
yield sequence | ||
|
||
|
||
def write_sequences(sequences, path_or_buffer, format="fasta"): | ||
"""Write sequences to a given path in the given format. | ||
Automatically infer compression mode (e.g., gzip, etc.) based on the path's | ||
filename extension. | ||
Parameters | ||
---------- | ||
sequences : iterable of Bio.SeqRecord.SeqRecord objects | ||
A list-like collection of sequences to write | ||
path_or_buffer : str or Path-like object or IO buffer | ||
A path to a file to write the given sequences in the given format. | ||
format : str | ||
Format of input sequences matching any of those supported by BioPython | ||
(e.g., "fasta", "genbank", etc.) | ||
Returns | ||
------- | ||
int : | ||
Number of sequences written out to the given path. | ||
""" | ||
with open_file(path_or_buffer, "wt") as handle: | ||
# Bio.SeqIO supports writing to the same handle multiple times for specific | ||
# file formats. For the formats we use, this function call should work for | ||
# both a newly opened file handle or one that is provided by the caller. | ||
# For more details see: | ||
# https://github.com/biopython/biopython/blob/25f5152f4aeefe184a323db25694fbfe0593f0e2/Bio/SeqIO/__init__.py#L233-L251 | ||
sequences_written = Bio.SeqIO.write( | ||
sequences, | ||
handle, | ||
format | ||
) | ||
|
||
return sequences_written |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,179 @@ | ||
#!/usr/bin/env python3 | ||
from Bio import SeqIO | ||
from Bio.Seq import Seq | ||
from Bio.SeqRecord import SeqRecord | ||
import bz2 | ||
import gzip | ||
import lzma | ||
from pathlib import Path | ||
import pytest | ||
import random | ||
import sys | ||
|
||
from augur.io import open_file, read_sequences, write_sequences | ||
|
||
|
||
def random_seq(k): | ||
"""Generate a single random sequence of nucleotides of length k. | ||
""" | ||
return "".join(random.choices(("A","T","G","C"), k=k)) | ||
|
||
def generate_sequences(n, k=10): | ||
"""Generate n random sequences of length k. | ||
""" | ||
return ( | ||
SeqRecord(Seq(random_seq(k)), id=f"SEQ_{i}") | ||
for i in range(1, n + 1) | ||
) | ||
|
||
@pytest.fixture | ||
def sequences(): | ||
return list(generate_sequences(3)) | ||
|
||
@pytest.fixture | ||
def sequences_generator(): | ||
return generate_sequences(3) | ||
|
||
@pytest.fixture | ||
def fasta_filename(tmpdir, sequences): | ||
filename = str(tmpdir / "sequences.fasta") | ||
SeqIO.write(sequences, filename, "fasta") | ||
return filename | ||
|
||
@pytest.fixture | ||
def additional_fasta_filename(tmpdir, sequences): | ||
filename = str(tmpdir / "additional_sequences.fasta") | ||
SeqIO.write(sequences, filename, "fasta") | ||
return filename | ||
|
||
@pytest.fixture | ||
def gzip_fasta_filename(tmpdir, sequences): | ||
filename = str(tmpdir / "sequences.fasta.gz") | ||
|
||
with gzip.open(filename, "wt") as oh: | ||
SeqIO.write(sequences, oh, "fasta") | ||
|
||
return filename | ||
|
||
@pytest.fixture | ||
def bzip2_fasta_filename(tmpdir, sequences): | ||
filename = str(tmpdir / "sequences.fasta.bz2") | ||
|
||
with bz2.open(filename, "wt") as oh: | ||
SeqIO.write(sequences, oh, "fasta") | ||
|
||
return filename | ||
|
||
@pytest.fixture | ||
def lzma_fasta_filename(tmpdir, sequences): | ||
filename = str(tmpdir / "sequences.fasta.xz") | ||
|
||
with lzma.open(filename, "wt") as oh: | ||
SeqIO.write(sequences, oh, "fasta") | ||
|
||
return filename | ||
|
||
@pytest.fixture | ||
def genbank_reference(): | ||
return "tests/builds/zika/config/zika_outgroup.gb" | ||
|
||
|
||
class TestReadSequences: | ||
def test_read_sequences_from_single_file(self, fasta_filename): | ||
sequences = read_sequences(fasta_filename, format="fasta") | ||
assert len(list(sequences)) == 3 | ||
|
||
def test_read_sequences_from_multiple_files(self, fasta_filename, additional_fasta_filename): | ||
sequences = read_sequences(fasta_filename, additional_fasta_filename, format="fasta") | ||
assert len(list(sequences)) == 6 | ||
|
||
def test_read_sequences_from_multiple_files_or_buffers(self, fasta_filename, additional_fasta_filename): | ||
with open(fasta_filename) as fasta_handle: | ||
sequences = read_sequences(fasta_handle, additional_fasta_filename, format="fasta") | ||
assert len(list(sequences)) == 6 | ||
|
||
def test_read_single_fasta_record(self, fasta_filename): | ||
record = next(read_sequences(fasta_filename, format="fasta")) | ||
assert record.id == "SEQ_1" | ||
|
||
def test_read_single_genbank_record(self, genbank_reference): | ||
reference = next(read_sequences(genbank_reference, format="genbank")) | ||
assert reference.id == "KX369547.1" | ||
|
||
def test_read_single_genbank_record_from_a_path(self, genbank_reference): | ||
reference = next(read_sequences(Path(genbank_reference), format="genbank")) | ||
assert reference.id == "KX369547.1" | ||
|
||
def test_read_sequences_from_single_gzip_file(self, gzip_fasta_filename): | ||
sequences = read_sequences(gzip_fasta_filename, format="fasta") | ||
assert len(list(sequences)) == 3 | ||
|
||
def test_read_sequences_from_single_lzma_file(self, lzma_fasta_filename): | ||
sequences = read_sequences(lzma_fasta_filename, format="fasta") | ||
assert len(list(sequences)) == 3 | ||
|
||
def test_read_sequences_from_single_bzip2_file(self, bzip2_fasta_filename): | ||
sequences = read_sequences(bzip2_fasta_filename, format="fasta") | ||
assert len(list(sequences)) == 3 | ||
|
||
def test_read_sequences_from_multiple_files_with_different_compression(self, fasta_filename, gzip_fasta_filename, lzma_fasta_filename): | ||
sequences = read_sequences(fasta_filename, gzip_fasta_filename, lzma_fasta_filename, format="fasta") | ||
assert len(list(sequences)) == 9 | ||
|
||
|
||
class TestWriteSequences: | ||
def test_write_sequences(self, tmpdir, sequences): | ||
output_filename = Path(tmpdir) / Path("new_sequences.fasta") | ||
sequences_written = write_sequences(sequences, output_filename, "fasta") | ||
assert sequences_written == len(sequences) | ||
|
||
def test_write_genbank_sequence(self, tmpdir, genbank_reference): | ||
output_filename = Path(tmpdir) / Path("new_sequences.fasta") | ||
|
||
reference = SeqIO.read(genbank_reference, "genbank") | ||
sequences_written = write_sequences([reference], output_filename, "genbank") | ||
assert sequences_written == 1 | ||
|
||
def test_write_sequences_from_generator(self, tmpdir, sequences_generator): | ||
output_filename = Path(tmpdir) / Path("new_sequences.fasta") | ||
sequences_written = write_sequences(sequences_generator, output_filename, "fasta") | ||
assert sequences_written == 3 | ||
|
||
def test_write_single_set_of_sequences_to_gzip_file(self, tmpdir, sequences): | ||
output_filename = Path(tmpdir) / Path("new_sequences.fasta.gz") | ||
sequences_written = write_sequences(sequences, output_filename, "fasta") | ||
assert sequences_written == len(sequences) | ||
|
||
with gzip.open(output_filename, "rt") as handle: | ||
assert sequences_written == len([line for line in handle if line.startswith(">")]) | ||
|
||
def test_write_single_set_of_sequences_to_bzip2_file(self, tmpdir, sequences): | ||
output_filename = Path(tmpdir) / Path("new_sequences.fasta.bz2") | ||
sequences_written = write_sequences(sequences, output_filename, "fasta") | ||
assert sequences_written == len(sequences) | ||
|
||
with bz2.open(output_filename, "rt") as handle: | ||
assert sequences_written == len([line for line in handle if line.startswith(">")]) | ||
|
||
def test_write_single_set_of_sequences_to_lzma_file(self, tmpdir, sequences): | ||
output_filename = Path(tmpdir) / Path("new_sequences.fasta.xz") | ||
sequences_written = write_sequences(sequences, output_filename, "fasta") | ||
assert sequences_written == len(sequences) | ||
|
||
with lzma.open(output_filename, "rt") as handle: | ||
assert sequences_written == len([line for line in handle if line.startswith(">")]) | ||
|
||
def test_write_sequences_by_external_handle(self, tmpdir, sequences): | ||
output_filename = Path(tmpdir) / Path("new_sequences.fasta") | ||
|
||
with open_file(output_filename, "w") as handle: | ||
total_sequences_written = 0 | ||
for sequence in sequences: | ||
sequences_written = write_sequences( | ||
sequence, | ||
handle | ||
) | ||
total_sequences_written += sequences_written | ||
|
||
with open(output_filename, "r") as handle: | ||
assert total_sequences_written == len([line for line in handle if line.startswith(">")]) |