fix: plot_format leaks its settings when the wrapped block raises

pybloqs/block/image.py

@@ -66,8 +66,10 @@ def plot_format(plot_format: Optional[str] = None, dpi: Optional[int] = None) ->
     """
     old = get_plot_format()
     set_plot_format(plot_format, dpi)
-    yield
-    set_plot_format(*old)
+    try:
+        yield
+    finally:
+        set_plot_format(*old)
 
 
 def get_plot_format() -> Tuple[str, int]:
@@ -92,8 +94,11 @@ def set_plot_format(plot_format: Optional[str] = None, plot_dpi: Optional[int] =
     global _PLOT_MIME_TYPE
     global _PLOT_DPI
     if plot_format is not None:
+        # Look the mime type up before assigning either global, so an unsupported format
+        # leaves the existing settings untouched rather than half-applied.
+        mime_type = _MIME_TYPES[plot_format]
         _PLOT_FORMAT = plot_format
-        _PLOT_MIME_TYPE = _MIME_TYPES[plot_format]
+        _PLOT_MIME_TYPE = mime_type
 
     if plot_dpi is not None:
         _PLOT_DPI = plot_dpi

tests/unit/block/test_image_unit.py

@@ -56,3 +56,22 @@ def test_plot_format_ctx_manager():
         assert i._PLOT_DPI == 1000
     assert old_fmt == i._PLOT_FORMAT
     assert old_dpi == i._PLOT_DPI
+
+
+def test_plot_format_ctx_manager_restores_on_exception():
+    old_fmt, old_dpi, old_mime = i._PLOT_FORMAT, i._PLOT_DPI, i._PLOT_MIME_TYPE
+    with pytest.raises(RuntimeError):
+        with i.plot_format("svg", 1000):
+            raise RuntimeError("failure while rendering")
+    assert i._PLOT_FORMAT == old_fmt
+    assert i._PLOT_DPI == old_dpi
+    assert i._PLOT_MIME_TYPE == old_mime
+
+
+def test_set_plot_format_unsupported_format_leaves_settings_unchanged():
+    old_fmt, old_dpi, old_mime = i._PLOT_FORMAT, i._PLOT_DPI, i._PLOT_MIME_TYPE
+    with pytest.raises(KeyError):
+        i.set_plot_format("jpeg")
+    assert i._PLOT_FORMAT == old_fmt
+    assert i._PLOT_DPI == old_dpi
+    assert i._PLOT_MIME_TYPE == old_mime