Merge pull request #12 from PennChopMicrobiomeProgram/write-totalreads-table

Add ability to write table of total read counts

Author: kylebittinger

dnabclib/main.py

@@ -61,6 +61,9 @@ def main(argv=None):
     p.add_argument(
         "--manifest-file", type=argparse.FileType("w"), help=(
             "Write manifest file for QIIME2"))
+    p.add_argument(
+        "--total-reads-file", type=argparse.FileType("w"), help=(
+            "Write TSV table of total read counts"))
     args = p.parse_args(argv)
 
     samples = list(Sample.load(args.barcode_file))
@@ -78,3 +81,5 @@ def main(argv=None):
 
     if args.manifest_file:
         writer.write_qiime2_manifest(args.manifest_file)
+    if args.total_reads_file:
+        writer.write_read_counts(args.total_reads_file, assigner.read_counts)

dnabclib/writer.py

@@ -30,6 +30,11 @@ def write_qiime2_manifest(self, f):
             fp1 = os.path.abspath(f1.name)
             f.write("{0},{1},forward\n".format(sample.name, fp1))
 
+    def write_read_counts(self, f, read_counts):
+        f.write("SampleID\tNumReads\n")
+        for sample_name, n in read_counts.items():
+            f.write("{0}\t{1}\n".format(sample_name, n))
+
     def _get_output_file(self, sample):
         f = self._open_files.get(sample)
         if f is None:

test/test_main.py

@@ -46,6 +46,7 @@ def setUp(self):
         self.output_dir = os.path.join(self.temp_dir, "output")
         self.summary_fp = os.path.join(self.temp_dir, "summary.json")
         self.manifest_fp = os.path.join(self.temp_dir, "manifest.csv")
+        self.total_reads_fp = os.path.join(self.temp_dir, "read_counts.tsv")
 
     def tearDown(self):
         shutil.rmtree(self.temp_dir)
@@ -58,6 +59,7 @@ def test_regular(self):
             "--barcode-file", self.barcode_fp,
             "--output-dir", self.output_dir,
             "--manifest-file", self.manifest_fp,
+            "--total-reads-file", self.total_reads_fp,
             "--revcomp",
             ])
         self.assertEqual(
@@ -72,6 +74,12 @@ def test_regular(self):
                 self.assertIn(direction, ["forward\n", "reverse\n"])
                 self.assertIn(sample_id, ["SampleA", "SampleB"])
 
+        with open(self.total_reads_fp) as f:
+            self.assertEqual(next(f), "SampleID\tNumReads\n")
+            self.assertEqual(next(f), "SampleA\t1\n")
+            self.assertEqual(next(f), "SampleB\t1\n")
+            self.assertEqual(next(f), "unassigned\t1\n")
+
 
 class SampleNameTests(unittest.TestCase):
     def test_get_sample_names_main(self):

test/test_writer.py

@@ -64,6 +64,12 @@ def test_write(self):
             "h56,{0},forward\n".format(fp),
         ])
 
+        f2 = MockFile()
+        w.write_read_counts(f2, {"s1": 365})
+        self.assertEqual(f2.contents, [
+            "SampleID\tNumReads\n",
+            "s1\t365\n"
+        ])
 
 class MockFile:
     def __init__(self):