Adding to ascii file (#7)

Author: josenavas

qiita_files/demux.py

@@ -461,6 +461,43 @@ def to_per_sample_ascii(demux, samples=None):
         yield samp, to_ascii(demux, samples=[samp])
 
 
+def to_ascii_file(demux_fp, output_fp, samples=None, out_format='fastq'):
+    """Writes the sequences on FASTQ or FASTA format
+
+    Parameters
+    ----------
+    demux_fp : str
+        The demux file path
+    output_fp : str
+        The output file path
+    samples : list of str, optional
+        Samples to pull out. If None, then all samples will be examined.
+        Defaults to None.
+    out_format: {'fastq', 'fasta'}, optional
+        The format in which the output file should be written. Default: FASTQ
+
+    Raises
+    ------
+    ValueError
+        If `out_format` is not 'fastq' or 'fasta'
+    """
+    if out_format == 'fastq':
+        formatter = format_fastq_record
+    elif out_format == 'fasta':
+        formatter = format_fasta_record
+    else:
+        raise ValueError("'out_format' should be either 'fastq' or 'fasta', "
+                         "found: %s" % out_format)
+
+    with open_file(demux_fp, 'r') as demux:
+        if samples is None:
+            samples = list(demux.keys())
+        samples = [s.encode() for s in samples]
+        with open(output_fp, 'wb') as out:
+            for rec in _to_ascii(demux, samples, formatter):
+                out.write(rec)
+
+
 def _to_file(demux_fp, sample, fp, formatter):
     with open_file(demux_fp, 'r') as demux:
         with open(fp, 'wb') as out:

qiita_files/tests/test_demux.py

@@ -22,7 +22,7 @@
                                _per_sample_lengths, _summarize_lengths,
                                _set_attr_stats, _construct_datasets, to_hdf5,
                                to_ascii, stat, to_per_sample_ascii,
-                               to_per_sample_files)
+                               to_per_sample_files, to_ascii_file)
 
 
 class BufferTests(TestCase):
@@ -328,6 +328,60 @@ def test_to_per_sample_ascii(self):
         obs = [(s[0], list(s[1])) for s in to_per_sample_ascii(self.hdf5_file)]
         self.assertEqual(obs, exp)
 
+    def test_to_ascii_file(self):
+        with tempfile.NamedTemporaryFile('r+', suffix='.fq',
+                                         delete=False) as f:
+            f.write(fqdata_variable_length)
+
+        self.to_remove.append(f.name)
+
+        with tempfile.NamedTemporaryFile('r+', suffix='.demux',
+                                         delete=False) as demux_f:
+            pass
+
+        self.to_remove.append(demux_f.name)
+
+        with h5py.File(demux_f.name, 'r+') as demux:
+            to_hdf5(f.name, demux)
+
+        with tempfile.NamedTemporaryFile('r+', suffix='.fq',
+                                         delete=False) as obs_fq:
+            pass
+        self.to_remove.append(obs_fq.name)
+
+        to_ascii_file(demux_f.name, obs_fq.name)
+        with open(obs_fq.name, 'rb') as obs_f:
+            obs = obs_f.read()
+        exp = (b'@a_0 orig_bc=abc new_bc=abc bc_diffs=0\nxyz\n+\nABC\n'
+               b'@b_0 orig_bc=abw new_bc=wbc bc_diffs=4\nqwe\n+\nDFG\n'
+               b'@b_1 orig_bc=abw new_bc=wbc bc_diffs=4\nqwexx\n+\nDEF#G\n')
+        self.assertEqual(obs, exp)
+
+        with tempfile.NamedTemporaryFile('r+', suffix='.fa',
+                                         delete=False) as obs_fa:
+            pass
+        self.to_remove.append(obs_fa.name)
+
+        to_ascii_file(demux_f.name, obs_fa.name, out_format='fasta')
+        with open(obs_fa.name, 'rb') as obs_f:
+            obs = obs_f.read()
+        exp = (b'>a_0 orig_bc=abc new_bc=abc bc_diffs=0\nxyz\n'
+               b'>b_0 orig_bc=abw new_bc=wbc bc_diffs=4\nqwe\n'
+               b'>b_1 orig_bc=abw new_bc=wbc bc_diffs=4\nqwexx\n')
+        self.assertEqual(obs, exp)
+
+        with tempfile.NamedTemporaryFile('r+', suffix='.fq',
+                                         delete=False) as obs_fq:
+            pass
+        self.to_remove.append(obs_fq.name)
+
+        to_ascii_file(demux_f.name, obs_fq.name, samples=['b'])
+        with open(obs_fq.name, 'rb') as obs_f:
+            obs = obs_f.read()
+        exp = (b'@b_0 orig_bc=abw new_bc=wbc bc_diffs=4\nqwe\n+\nDFG\n'
+               b'@b_1 orig_bc=abw new_bc=wbc bc_diffs=4\nqwexx\n+\nDEF#G\n')
+        self.assertEqual(obs, exp)
+
     def test_to_files(self):
         # implicitly tested with test_to_per_sample_fasta
         pass