Close mzML output handle explicitly

Author: joewandy

tests/test_mzml_writer.py

@@ -0,0 +1,69 @@
+import vimms.MzmlWriter as mzml_writer_module
+from vimms.MzmlWriter import MzmlWriter
+
+
+class _ContextManager:
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        return False
+
+
+class _FakePsimsMzMLWriter:
+    opened_handles = []
+
+    def __init__(self, out_handle):
+        self.opened_handles.append(out_handle)
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        return False
+
+    def controlled_vocabularies(self):
+        pass
+
+    def file_description(self, *args, **kwargs):
+        pass
+
+    def sample_list(self, *args, **kwargs):
+        pass
+
+    def software_list(self, *args, **kwargs):
+        pass
+
+    def scan_settings_list(self, *args, **kwargs):
+        pass
+
+    def instrument_configuration_list(self, *args, **kwargs):
+        pass
+
+    def data_processing_list(self, *args, **kwargs):
+        pass
+
+    def run(self, *args, **kwargs):
+        return _ContextManager()
+
+    def spectrum_list(self, *args, **kwargs):
+        return _ContextManager()
+
+    def chromatogram_list(self, *args, **kwargs):
+        return _ContextManager()
+
+    def write_chromatogram(self, *args, **kwargs):
+        pass
+
+    def close(self):
+        pass
+
+
+def test_mzml_writer_closes_output_handle(monkeypatch, tmp_path):
+    _FakePsimsMzMLWriter.opened_handles = []
+    monkeypatch.setattr(mzml_writer_module, "PsimsMzMLWriter", _FakePsimsMzMLWriter)
+
+    MzmlWriter("test_analysis", {1: []}).write_mzML(tmp_path / "test.mzML")
+
+    assert len(_FakePsimsMzMLWriter.opened_handles) == 1
+    assert _FakePsimsMzMLWriter.opened_handles[0].closed

vimms/MzmlWriter.py

@@ -50,29 +50,28 @@ def write_mzML(self, out_file):
         create_if_not_exist(out_dir)
 
         # start writing mzML here
-        with PsimsMzMLWriter(open(out_file, "wb")) as writer:
-            # add default controlled vocabularies
-            writer.controlled_vocabularies()
-
-            # write other fields like sample list, software list, etc.
-            self._write_info(writer)
-
-            # open the run
-            with writer.run(id=self.analysis_name):
-                self._write_spectra(writer, self.scans)
-
-                # open chromatogram list sections
-                with writer.chromatogram_list(count=1):
-                    tic_rts, tic_intensities = self._get_tic_chromatogram(self.scans)
-                    writer.write_chromatogram(
-                        tic_rts,
-                        tic_intensities,
-                        id="tic",
-                        chromatogram_type="total ion current chromatogram",
-                        time_unit="second",
-                    )
-
-        writer.close()
+        with open(out_file, "wb") as out_handle:
+            with PsimsMzMLWriter(out_handle) as writer:
+                # add default controlled vocabularies
+                writer.controlled_vocabularies()
+
+                # write other fields like sample list, software list, etc.
+                self._write_info(writer)
+
+                # open the run
+                with writer.run(id=self.analysis_name):
+                    self._write_spectra(writer, self.scans)
+
+                    # open chromatogram list sections
+                    with writer.chromatogram_list(count=1):
+                        tic_rts, tic_intensities = self._get_tic_chromatogram(self.scans)
+                        writer.write_chromatogram(
+                            tic_rts,
+                            tic_intensities,
+                            id="tic",
+                            chromatogram_type="total ion current chromatogram",
+                            time_unit="second",
+                        )
 
     def _write_info(self, out):
         """