From dad159671c7edd2e3394da257f735533168de131 Mon Sep 17 00:00:00 2001 From: "Edward S. Rice" Date: Fri, 13 Sep 2024 12:26:40 +0200 Subject: [PATCH 1/2] 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. --- cyvcf2/cyvcf2.pyx | 12 ++++++++++++ cyvcf2/tests/test_reader.py | 6 ++++++ 2 files changed, 18 insertions(+) diff --git a/cyvcf2/cyvcf2.pyx b/cyvcf2/cyvcf2.pyx index a729997..6229d87 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: @@ -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: 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) From 87026b93cf1ee8c965c8ef0ce8c6970b3f782f6b Mon Sep 17 00:00:00 2001 From: "Edward S. Rice" Date: Fri, 13 Sep 2024 18:28:18 +0200 Subject: [PATCH 2/2] Remove enter/exit from Writer not necessary because they will be inherited --- cyvcf2/cyvcf2.pyx | 6 ------ 1 file changed, 6 deletions(-) diff --git a/cyvcf2/cyvcf2.pyx b/cyvcf2/cyvcf2.pyx index 6229d87..32e6f83 100644 --- a/cyvcf2/cyvcf2.pyx +++ b/cyvcf2/cyvcf2.pyx @@ -2420,12 +2420,6 @@ 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: