Skip to content

Commit

Permalink
Allow Reader/Writer to be context managers
Browse files Browse the repository at this point in the history
By adding __enter__ and __exit__ methods to the Reader and Writer
classes, they can now be used as context managers. Also added a test to
make sure the Reader context manager works correctly. Did not add a test
to do the same with Writer, because trying to write to a closed Writer
results in a segfault rather than raising an exception that can be
checked by pytest.
  • Loading branch information
esrice committed Sep 13, 2024
1 parent 8624896 commit dad1596
Show file tree
Hide file tree
Showing 2 changed files with 18 additions and 0 deletions.
12 changes: 12 additions & 0 deletions cyvcf2/cyvcf2.pyx
Original file line number Diff line number Diff line change
Expand Up @@ -293,6 +293,12 @@ cdef class VCF(HTSFile):
if threads is not None:
self.set_threads(threads)

def __enter__(self):
return self

def __exit__(self, type, value, tb):
self.close()

def set_threads(self, int n):
v = hts_set_threads(self.hts, n)
if v < 0:
Expand Down Expand Up @@ -2414,6 +2420,12 @@ cdef class Writer(VCF):
bcf_hdr_sync(self.hdr)
self.header_written = False

def __enter__(self):
return self

def __exit__(self, type, value, tb):
self.close()

@staticmethod
def _infer_file_mode(fname, mode=None):
if mode is not None:
Expand Down
6 changes: 6 additions & 0 deletions cyvcf2/tests/test_reader.py
Original file line number Diff line number Diff line change
Expand Up @@ -1410,3 +1410,9 @@ def test_num_records_no_index(path):
vcf = VCF(os.path.join(HERE, path))
with pytest.raises(ValueError, match="must be indexed"):
vcf.num_records

def test_reader_context_manager():
with VCF(VCF_PATH) as vcf:
pass
with pytest.raises(Exception, match="attempt to iterate over closed"):
next(vcf)

0 comments on commit dad1596

Please sign in to comment.