From 541ab16a255a5287c331843d8180ed6b9ef10e00 Mon Sep 17 00:00:00 2001 From: "Edward S. Rice" Date: Mon, 16 Sep 2024 16:19:02 +0200 Subject: [PATCH] Allow Reader/Writer to be context managers (#311) * Allow Reader/Writer to be context managers 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. * Remove enter/exit from Writer not necessary because they will be inherited --- cyvcf2/cyvcf2.pyx | 6 ++++++ cyvcf2/tests/test_reader.py | 6 ++++++ 2 files changed, 12 insertions(+) diff --git a/cyvcf2/cyvcf2.pyx b/cyvcf2/cyvcf2.pyx index a729997..32e6f83 100644 --- a/cyvcf2/cyvcf2.pyx +++ b/cyvcf2/cyvcf2.pyx @@ -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: diff --git a/cyvcf2/tests/test_reader.py b/cyvcf2/tests/test_reader.py index 84d2249..43d954d 100644 --- a/cyvcf2/tests/test_reader.py +++ b/cyvcf2/tests/test_reader.py @@ -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)