Skip to content

Commit

Permalink
Add initial I/O interface and tests
Browse files Browse the repository at this point in the history
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
huddlej committed Mar 10, 2021
1 parent 8df4b4d commit 8a20b4f
Show file tree
Hide file tree
Showing 4 changed files with 292 additions and 1 deletion.
106 changes: 106 additions & 0 deletions augur/io.py
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
5 changes: 5 additions & 0 deletions pytest.python3.ini
Original file line number Diff line number Diff line change
@@ -1,4 +1,9 @@
[pytest]
# ignore biopython's deprecation warnings about alphabet that are outside of our control
filterwarnings =
ignore::PendingDeprecationWarning:Bio.Alphabet
ignore:Coverage disabled

addopts =
# do not capture any output---necessary for interactive breakpoints
-s
Expand Down
3 changes: 2 additions & 1 deletion setup.py
Original file line number Diff line number Diff line change
Expand Up @@ -52,7 +52,8 @@
"jsonschema >=3.0.0, ==3.*",
"packaging >=19.2",
"pandas >=1.0.0, ==1.*",
"phylo-treetime ==0.8.*"
"phylo-treetime ==0.8.*",
"xopen >=1.0.1, ==1.*"
],
extras_require = {
'full': [
Expand Down
179 changes: 179 additions & 0 deletions tests/test_io.py
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(">")])

0 comments on commit 8a20b4f

Please sign in to comment.