fix: close CLI capture files after output errors

pslab/cli.py

@@ -21,6 +21,7 @@
 import struct
 import sys
 import time
+from contextlib import suppress
 from itertools import zip_longest
 from typing import List, Tuple
 
@@ -146,24 +147,26 @@ def collect(handler: SerialHandler, args: argparse.Namespace):
 
     output = instrument(handler, args.channels, args.duration)
 
-    if args.file_path is not None:
-        file = open(args.file_path, "w")
-    else:
-        file = sys.stdout
-
-    if not args.json:
-        csv_file = csv.writer(file)
-        csv_file.writerow(output[0])
-        for row in zip_longest(*output[1]):
-            csv_file.writerow(row)
+    file = open(args.file_path, "w") if args.file_path is not None else sys.stdout
+    try:
+        if not args.json:
+            csv_file = csv.writer(file)
+            csv_file.writerow(output[0])
+            for row in zip_longest(*output[1]):
+                csv_file.writerow(row)
+        else:
+            output_dict = dict()
+            for key, val in zip_longest(*output):
+                output_dict[key] = val.tolist()
+            json.dump(output_dict, file)
+    except BaseException:
+        if args.file_path is not None:
+            with suppress(OSError):
+                file.close()
+        raise
     else:
-        output_dict = dict()
-        for key, val in zip_longest(*output):
-            output_dict[key] = val.tolist()
-        json.dump(output_dict, file)
-
-    if args.file_path is not None:
-        file.close()
+        if args.file_path is not None:
+            file.close()
 
 
 def wave(handler: SerialHandler, args: argparse.Namespace):

tests/test_cli_output_cleanup.py

@@ -0,0 +1,134 @@
+"""CLI output resources are released without closing caller-owned streams."""
+
+import argparse
+import builtins
+import io
+import json
+
+import numpy as np
+import pytest
+
+from pslab import cli
+
+
+class FailingOutput:
+    """Delegate file ownership while simulating a failed output write."""
+
+    def __init__(self, file):
+        self.file = file
+        self.error = OSError("output storage unavailable")
+
+    def write(self, text):
+        self.file.write(text[:1])
+        raise self.error
+
+    def close(self):
+        self.file.close()
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *args):
+        self.close()
+
+
+class FailingCloseOutput(FailingOutput):
+    """Expose a close failure after releasing the underlying file."""
+
+    def __init__(self, file):
+        super().__init__(file)
+        self.close_error = OSError("output close failed")
+
+    def close(self):
+        super().close()
+        raise self.close_error
+
+
+@pytest.fixture
+def output_args(monkeypatch):
+    monkeypatch.setitem(
+        cli.INSTRUMENTS,
+        "test-instrument",
+        lambda *args: (["CH1"], [np.array([1.0, 2.0])]),
+    )
+    return argparse.Namespace(
+        instrument="test-instrument", channels=1, duration=1, file_path=None, json=False
+    )
+
+
+@pytest.mark.parametrize("as_json", [False, True])
+def test_failed_file_output_is_closed(tmp_path, monkeypatch, output_args, as_json):
+    """Release the real file even when the write exception retains its frame."""
+    destination = tmp_path / "capture"
+    file = builtins.open(destination, "w")
+    output = FailingOutput(file)
+    monkeypatch.setattr(cli, "open", lambda *args: output, raising=False)
+    output_args.file_path = str(destination)
+    output_args.json = as_json
+    try:
+        with pytest.raises(OSError) as raised:
+            cli.collect(None, output_args)
+        assert raised.value is output.error
+        assert file.closed
+    finally:
+        file.close()
+
+
+@pytest.mark.parametrize("as_json", [False, True])
+def test_failed_stdout_is_not_closed(monkeypatch, output_args, as_json):
+    """Preserve caller ownership after either formatter encounters a write error."""
+    file = io.StringIO()
+    output = FailingOutput(file)
+    monkeypatch.setattr(cli.sys, "stdout", output)
+    output_args.json = as_json
+    with pytest.raises(OSError) as raised:
+        cli.collect(None, output_args)
+    assert raised.value is output.error
+    assert not file.closed
+
+
+@pytest.mark.parametrize("as_json", [False, True])
+def test_successful_file_output_is_closed(tmp_path, monkeypatch, output_args, as_json):
+    """Close successful output while retaining CSV and JSON content."""
+    destination = tmp_path / "capture"
+    file = builtins.open(destination, "w")
+    monkeypatch.setattr(cli, "open", lambda *args: file, raising=False)
+    output_args.file_path = str(destination)
+    output_args.json = as_json
+    cli.collect(None, output_args)
+    assert file.closed
+    if as_json:
+        assert json.loads(destination.read_text()) == {"CH1": [1.0, 2.0]}
+    else:
+        assert destination.read_text().splitlines() == ["CH1", "1.0", "2.0"]
+
+
+@pytest.mark.parametrize("as_json", [False, True])
+def test_close_failure_does_not_replace_write_error(
+    tmp_path, monkeypatch, output_args, as_json
+):
+    file = builtins.open(tmp_path / "capture", "w")
+    output = FailingCloseOutput(file)
+    monkeypatch.setattr(cli, "open", lambda *args: output, raising=False)
+    output_args.file_path = str(tmp_path / "capture")
+    output_args.json = as_json
+    with pytest.raises(OSError) as raised:
+        cli.collect(None, output_args)
+    assert raised.value is output.error
+    assert file.closed
+
+
+@pytest.mark.parametrize("as_json", [False, True])
+def test_close_failure_after_successful_write_is_visible(
+    tmp_path, monkeypatch, output_args, as_json
+):
+    file = builtins.open(tmp_path / "capture", "w")
+    output = FailingCloseOutput(file)
+    monkeypatch.setattr(output, "write", file.write)
+    monkeypatch.setattr(cli, "open", lambda *args: output, raising=False)
+    output_args.file_path = str(tmp_path / "capture")
+    output_args.json = as_json
+    with pytest.raises(OSError) as raised:
+        cli.collect(None, output_args)
+    assert raised.value is output.close_error
+    assert file.closed