Docs/test for undo_monkey_patch

vcf/sample_filter.py

@@ -1,5 +1,3 @@
-#!/usr/bin/env python
-
 # Author: Lenna X. Peterson
 # github.com/lennax
 # arklenna at gmail dot com
@@ -13,6 +11,13 @@
 
 
 class SampleFilter(object):
+    """
+    Modifies the vcf Reader to filter each row by sample as it is parsed.
+    When using the class, be sure to call `undo_monkey_patch()` to restore
+    the original functionality to the Reader.
+
+    """
+
     def __init__(self, infile, outfile=None, filters=None, invert=False):
         # Methods to add to Reader
         def get_filter(self):
@@ -99,5 +104,5 @@ def write(self, outfile=None):
             writer.write_record(row)
 
     def undo_monkey_patch(self):
-        delattr(Reader, 'sample_filter')
         Reader._parse_samples = self._orig_parse_samples
+        delattr(Reader, 'sample_filter')

vcf/test/test_vcf.py

@@ -669,10 +669,10 @@ def testSampleFilterModule(self):
         #print(buf.getvalue())
         # undo monkey patch
         filt.undo_monkey_patch()
+        self.assertTrue('sample_filter' not in dir(vcf.Reader))
         # read output
         reader = vcf.Reader(buf)
         self.assertEqual(reader.samples, ['NA00001'])
-        print(dir(reader))
         rec = reader.next()
         self.assertEqual(len(rec.samples), 1)