Add method _filter_clusters

sanger-pathogens · Apr 28, 2016 · 3fe8933 · 3fe8933
1 parent b4ef535
commit 3fe8933
Show file tree

Hide file tree

Showing 2 changed files with 64 additions and 0 deletions.
diff --git a/ariba/summary.py b/ariba/summary.py
@@ -114,6 +114,36 @@ def _gather_output_rows(self):
         return rows
 
 
+    @classmethod
+    def _filter_clusters(cls, rows):
+        '''Removes any cluster where every sample has assembled == "no"'''
+        found_a_yes = set()
+        first_filename = True
+        all_clusters = set()
+
+        for filename in rows:
+            for cluster in rows[filename]:
+                if first_filename:
+                    all_clusters.add(cluster)
+
+                if cluster in found_a_yes:
+                    continue
+
+                if rows[filename][cluster]['assembled'] == 'yes':
+                    found_a_yes.add(cluster)
+                    break
+
+            first_filename = False
+
+        to_delete = all_clusters.difference(found_a_yes)
+
+        for filename in rows:
+            for cluster in to_delete:
+                del rows[filename][cluster]
+
+        return rows
+
+
     @classmethod
     def _write_csv(cls, filenames, rows, outfile, phandango=False):
         lines = []

diff --git a/ariba/tests/summary_test.py b/ariba/tests/summary_test.py
@@ -120,6 +120,40 @@ def test_gather_output_rows(self):
         self.assertEqual(expected, got)
 
 
+    def test_filter_clusters(self):
+        '''Test _filter_clusters'''
+        rows = {
+            'file1': {
+                'cluster1': {'assembled': 'yes'},
+                'cluster2': {'assembled': 'yes'},
+                'cluster3': {'assembled': 'no'},
+                'cluster4': {'assembled': 'no'},
+            },
+            'file2': {
+                'cluster1': {'assembled': 'yes'},
+                'cluster2': {'assembled': 'no'},
+                'cluster3': {'assembled': 'yes'},
+                'cluster4': {'assembled': 'no'},
+            }
+        }
+
+        expected = {
+            'file1': {
+                'cluster1': {'assembled': 'yes'},
+                'cluster2': {'assembled': 'yes'},
+                'cluster3': {'assembled': 'no'},
+            },
+            'file2': {
+                'cluster1': {'assembled': 'yes'},
+                'cluster2': {'assembled': 'no'},
+                'cluster3': {'assembled': 'yes'},
+            }
+        }
+
+        got = summary.Summary._filter_clusters(rows)
+        self.assertEqual(expected, got)
+
+
     def test_write_csv(self):
         '''Test _write_csv'''
         tmp_out = 'tmp.out.tsv'