Print package name when writer backend is missing (#7980)

monai/data/image_writer.py

@@ -99,24 +99,35 @@ def resolve_writer(ext_name, error_if_not_found=True) -> Sequence:
             As an indexing key it will be converted to a lower case string.
         error_if_not_found: whether to raise an error if no suitable image writer is found.
             if True , raise an ``OptionalImportError``, otherwise return an empty tuple. Default is ``True``.
+            When no writer is available because some backends are not installed, the error message
+            includes the ``pip install`` command of the missing packages.
     """
     if not SUPPORTED_WRITERS:
         init()
     fmt = f"{ext_name}".lower()
     if fmt.startswith("."):
         fmt = fmt[1:]
     avail_writers = []
+    missing_packages: list[str] = []
     default_writers = SUPPORTED_WRITERS.get(EXT_WILDCARD, ())
     for _writer in look_up_option(fmt, SUPPORTED_WRITERS, default=default_writers):
         try:
             _writer()  # this triggers `monai.utils.module.require_pkg` to check the system availability
             avail_writers.append(_writer)
         except OptionalImportError:
+            # keep the package name of the unavailable backend so that it can be reported below
+            _pkg = getattr(_writer, "required_package", None)
+            if _pkg and _pkg not in missing_packages:
+                missing_packages.append(_pkg)
             continue
         except Exception:  # other writer init errors indicating it exists
             avail_writers.append(_writer)
     if not avail_writers and error_if_not_found:
-        raise OptionalImportError(f"No ImageWriter backend found for {fmt}.")
+        err_msg = f"No ImageWriter backend found for {fmt}."
+        if missing_packages:
+            install_hints = ", ".join(f"`pip install {_pkg}`" for _pkg in missing_packages)
+            err_msg += f" Please install one of the following packages: {install_hints}."
+        raise OptionalImportError(err_msg)
     writer_tuple = ensure_tuple(avail_writers)
     SUPPORTED_WRITERS[fmt] = writer_tuple
     return writer_tuple
@@ -175,6 +186,9 @@ class ImageWriter:
     ``resample_if_needed`` method.
     """
 
+    # name of the optional pip package providing this backend, reported when the backend is not installed
+    required_package: str | None = None
+
     def __init__(self, **kwargs):
         """
         The constructor supports adding new instance members.
@@ -371,6 +385,8 @@ class ITKWriter(ImageWriter):
 
     """
 
+    required_package = "itk"
+
     output_dtype: DtypeLike = None
     channel_dim: int | None
 
@@ -549,6 +565,8 @@ class NibabelWriter(ImageWriter):
 
     """
 
+    required_package = "nibabel"
+
     output_dtype: DtypeLike
     affine: Any
 
@@ -691,6 +709,8 @@ class PILWriter(ImageWriter):
         writer.write("test1.png", verbose=True)
     """
 
+    required_package = "pillow"
+
     output_dtype: DtypeLike
     channel_dim: int | None
     scale: int | None

tests/test_save_image.py

@@ -0,0 +1,100 @@
+# Copyright (c) MONAI Consortium
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import shutil
+import sys
+import tempfile
+import unittest
+from unittest.mock import patch
+
+import pytest
+
+from monai.data import image_writer
+from monai.data.image_writer import ImageWriter, resolve_writer
+from monai.utils import OptionalImportError
+
+
+def _reset_writers():
+    """Restore ``SUPPORTED_WRITERS`` to the state created by ``image_writer.init()``."""
+    image_writer.SUPPORTED_WRITERS.clear()
+    image_writer.init()
+
+
+class TestResolveWriterMissingBackend(unittest.TestCase):
+    """The error message should tell which package to install."""
+
+    def setUp(self):
+        _reset_writers()
+
+    def tearDown(self):
+        _reset_writers()
+
+    def test_missing_pillow(self):
+        # ``PILWriter`` is the only backend for png, the pip package name is ``pillow`` rather than ``PIL``
+        with patch.dict(sys.modules, {"PIL": None}):
+            with self.assertRaises(OptionalImportError) as ctx:
+                resolve_writer(".png")
+        msg = str(ctx.exception)
+        self.assertIn("No ImageWriter backend found for png", msg)
+        self.assertIn("`pip install pillow`", msg)
+        self.assertNotIn("`pip install PIL`", msg)
+
+    def test_missing_itk_and_nibabel(self):
+        with patch.dict(sys.modules, {"itk": None, "nibabel": None}):
+            with self.assertRaises(OptionalImportError) as ctx:
+                resolve_writer(".nii.gz")
+        msg = str(ctx.exception)
+        self.assertIn("No ImageWriter backend found for nii.gz", msg)
+        self.assertIn("`pip install itk`", msg)
+        self.assertIn("`pip install nibabel`", msg)
+
+    def test_missing_backend_is_import_error(self):
+        # the exception type is unchanged, only its message is extended
+        with patch.dict(sys.modules, {"itk": None, "nibabel": None}):
+            with self.assertRaises(ImportError):
+                resolve_writer(".nii")
+
+    def test_custom_writer_without_package_hint(self):
+        # a custom backend that provides no package hint keeps the previous message
+        class _CustomWriter(ImageWriter):
+            def __init__(self, *args, **kwargs):
+                raise OptionalImportError("required package `custom-pkg` is not installed.")
+
+        try:
+            image_writer.register_writer("custom", _CustomWriter)
+            with self.assertRaises(OptionalImportError) as ctx:
+                resolve_writer(".custom")
+            self.assertEqual(str(ctx.exception), "No ImageWriter backend found for custom.")
+        finally:
+            image_writer.SUPPORTED_WRITERS.pop("custom", None)
+
+    def test_required_package_of_built_in_writers(self):
+        self.assertEqual(image_writer.ITKWriter.required_package, "itk")
+        self.assertEqual(image_writer.NibabelWriter.required_package, "nibabel")
+        self.assertEqual(image_writer.PILWriter.required_package, "pillow")
+        self.assertIsNone(ImageWriter.required_package)
+
+
+def test_save_image_missing_backend():
+    """``SaveImage`` reports the install command when no writer backend is available."""
+    from monai.transforms import SaveImage
+
+    out_dir = tempfile.mkdtemp()
+    _reset_writers()
+    try:
+        with patch.dict(sys.modules, {"itk": None, "nibabel": None}):
+            with pytest.raises(ImportError, match="pip install itk"):
+                SaveImage(output_dir=out_dir, output_ext=".nii.gz")
+    finally:
+        _reset_writers()
+        shutil.rmtree(out_dir, ignore_errors=True)