Merge pull request #3698 from neutrinoceros/backport_3687

Backport #3410 + #3687 to yt-4.0.x (image filename validation fixes)

yt/visualization/_commons.py

@@ -1,13 +1,14 @@
 import os
-from typing import Optional
+import sys
+import warnings
+from typing import Optional, Type
 
 import matplotlib
 from packaging.version import Version
 
-from yt.utilities.logger import ytLogger as mylog
-
 from ._mpl_imports import (
     FigureCanvasAgg,
+    FigureCanvasBase,
     FigureCanvasPdf,
     FigureCanvasPS,
     FigureCanvasSVG,
@@ -23,13 +24,33 @@
 if MPL_VERSION >= Version("3.4"):
     DEFAULT_FONT_PROPERTIES["math_fontfamily"] = "cm"
 
+SUPPORTED_FORMATS = frozenset(FigureCanvasBase.get_supported_filetypes().keys())
+SUPPORTED_CANVAS_CLASSES = frozenset(
+    (FigureCanvasAgg, FigureCanvasPdf, FigureCanvasPS, FigureCanvasSVG)
+)
+
 
-AGG_FORMATS = [".png", ".jpg", ".jpeg", ".raw", ".rgba", ".tif", ".tiff"]
-SUPPORTED_FORMATS = AGG_FORMATS + [".eps", ".ps", ".pdf", ".svg"]
+def get_canvas_class(suffix: str) -> Type[FigureCanvasBase]:
+    s = normalize_extension_string(suffix)
+    if s not in SUPPORTED_FORMATS:
+        raise ValueError(f"Unsupported file format '{suffix}'.")
+    for cls in SUPPORTED_CANVAS_CLASSES:
+        if s in cls.get_supported_filetypes():
+            return cls
+    raise RuntimeError(
+        "Something went terribly wrong. "
+        f"File extension '{suffix}' is supposed to be supported "
+        "but no compatible backend was found."
+    )
 
 
 def normalize_extension_string(s: str) -> str:
-    return f".{s.lstrip('.')}"
+    if sys.version_info < (3, 9):
+        if s.startswith("."):
+            return s[1:]
+        return s
+    else:
+        return s.removeprefix(".")
 
 
 def validate_image_name(filename, suffix: Optional[str] = None) -> str:
@@ -39,28 +60,27 @@ def validate_image_name(filename, suffix: Optional[str] = None) -> str:
     Otherwise, suffix is appended to the filename, replacing any existing extension.
     """
     name, psuffix = os.path.splitext(filename)
+    psuffix = normalize_extension_string(psuffix)
+
+    if suffix is not None:
+        suffix = normalize_extension_string(suffix)
+
     if psuffix in SUPPORTED_FORMATS:
-        if suffix is not None:
-            suffix = normalize_extension_string(suffix)
         if suffix in SUPPORTED_FORMATS and suffix != psuffix:
-            mylog.warning(
-                "Received two valid image formats '%s' (from `filename`) "
-                "and '%s' (from `suffix`). The former is ignored.",
-                psuffix,
-                suffix,
+            warnings.warn(
+                f"Received two valid image formats {psuffix!r} (from filename) "
+                f"and {suffix!r} (from suffix). The former is ignored."
             )
-            return f"{name}{suffix}"
+            return f"{name}.{suffix}"
         return str(filename)
 
     if suffix is None:
-        suffix = ".png"
-
-    suffix = normalize_extension_string(suffix)
+        suffix = "png"
 
     if suffix not in SUPPORTED_FORMATS:
-        raise ValueError("Unsupported file format '{suffix}'.")
+        raise ValueError(f"Unsupported file format {suffix!r}")
 
-    return f"{filename}{suffix}"
+    return f"{filename}.{suffix}"
 
 
 def get_canvas(figure, filename):
@@ -72,19 +92,4 @@ def get_canvas(figure, filename):
             f"Can not determine canvas class from filename '{filename}' "
             f"without an extension."
         )
-
-    if suffix in AGG_FORMATS:
-        return FigureCanvasAgg(figure)
-
-    if suffix == ".pdf":
-        return FigureCanvasPdf(figure)
-
-    if suffix == ".svg":
-        return FigureCanvasSVG(figure)
-
-    if suffix in (".eps", ".ps"):
-        return FigureCanvasPS(figure)
-
-    raise ValueError(
-        f"No matching canvas for filename '{filename}' with extension '{suffix}'."
-    )
+    return get_canvas_class(suffix)(figure)

yt/visualization/_mpl_imports.py

@@ -1,4 +1,5 @@
 import matplotlib
+from matplotlib.backend_bases import FigureCanvasBase
 from matplotlib.backends.backend_agg import FigureCanvasAgg
 from matplotlib.backends.backend_pdf import FigureCanvasPdf
 from matplotlib.backends.backend_ps import FigureCanvasPS

yt/visualization/plot_container.py

@@ -6,7 +6,7 @@
 import warnings
 from collections import defaultdict
 from functools import wraps
-from typing import Any, Dict, Optional
+from typing import Any, Dict, List, Optional, Tuple, Union
 
 import numpy as np
 from matplotlib.cm import get_cmap
@@ -559,23 +559,23 @@ def set_figure_size(self, size):
     @validate_plot
     def save(
         self,
-        name: Optional[str] = None,
-        suffix: str = ".png",
+        name: Optional[Union[str, List[str], Tuple[str, ...]]] = None,
+        suffix: Optional[str] = None,
         mpl_kwargs: Optional[Dict[str, Any]] = None,
     ):
         """saves the plot to disk.
 
         Parameters
         ----------
-        name : string or tuple
+        name : string or tuple, optional
            The base of the filename. If name is a directory or if name is not
            set, the filename of the dataset is used. For a tuple, the
            resulting path will be given by joining the elements of the
            tuple
-        suffix : string
+        suffix : string, optional
            Specify the image type by its suffix. If not specified, the output
-           type will be inferred from the filename. Defaults to PNG.
-        mpl_kwargs : dict
+           type will be inferred from the filename. Defaults to '.png'.
+        mpl_kwargs : dict, optional
            A dict of keyword arguments to be passed to matplotlib.
 
         >>> slc.save(mpl_kwargs={"bbox_inches": "tight"})

yt/visualization/profile_plotter.py

@@ -3,6 +3,7 @@
 import os
 from collections import OrderedDict
 from functools import wraps
+from typing import Any, Dict, Optional
 
 import matplotlib
 import numpy as np
@@ -270,45 +271,53 @@ def __init__(
         ProfilePlot._initialize_instance(self, profiles, label, plot_spec, y_log)
 
     @validate_plot
-    def save(self, name=None, suffix=".png", mpl_kwargs=None):
+    def save(
+        self,
+        name: Optional[str] = None,
+        suffix: Optional[str] = None,
+        mpl_kwargs: Optional[Dict[str, Any]] = None,
+    ):
         r"""
         Saves a 1d profile plot.
 
         Parameters
         ----------
-        name : str
+        name : str, optional
             The output file keyword.
-        suffix : string
+        suffix : string, optional
             Specify the image type by its suffix. If not specified, the output
-            type will be inferred from the filename. Defaults to PNG.
-        mpl_kwargs : dict
+            type will be inferred from the filename. Defaults to '.png'.
+        mpl_kwargs : dict, optional
             A dict of keyword arguments to be passed to matplotlib.
         """
         if not self._plot_valid:
             self._setup_plots()
-        unique = set(self.plots.values())
-        if len(unique) < len(self.plots):
+
+        # Mypy is hardly convinced that we have a `plots` and a `profile` attr
+        # at this stage, so we're lasily going to deactivate it locally
+        unique = set(self.plots.values())  # type: ignore
+        if len(unique) < len(self.plots):  # type: ignore
             iters = zip(range(len(unique)), sorted(unique))
         else:
-            iters = self.plots.items()
+            iters = self.plots.items()  # type: ignore
 
         if name is None:
-            if len(self.profiles) == 1:
-                name = str(self.profiles[0].ds)
+            if len(self.profiles) == 1:  # type: ignore
+                name = str(self.profiles[0].ds)  # type: ignore
             else:
                 name = "Multi-data"
 
         name = validate_image_name(name, suffix)
         prefix, suffix = os.path.splitext(name)
 
-        xfn = self.profiles[0].x_field
+        xfn = self.profiles[0].x_field  # type: ignore
         if isinstance(xfn, tuple):
             xfn = xfn[1]
 
         names = []
         for uid, plot in iters:
-            if isinstance(uid, tuple):
-                uid = uid[1]
+            if isinstance(uid, tuple):  # type: ignore
+                uid = uid[1]  # type: ignore
             uid_name = f"{prefix}_1d-Profile_{xfn}_{uid}{suffix}"
             names.append(uid_name)
             mylog.info("Saving %s", uid_name)
@@ -1282,18 +1291,20 @@ def annotate_text(self, xpos=0.0, ypos=0.0, text=None, **text_kwargs):
         return self
 
     @validate_plot
-    def save(self, name=None, suffix=".png", mpl_kwargs=None):
+    def save(
+        self, name: Optional[str] = None, suffix: Optional[str] = None, mpl_kwargs=None
+    ):
         r"""
         Saves a 2d profile plot.
 
         Parameters
         ----------
-        name : str
+        name : str, optional
             The output file keyword.
-        suffix : string
+        suffix : string, optional
            Specify the image type by its suffix. If not specified, the output
-           type will be inferred from the filename. Defaults to PNG.
-        mpl_kwargs : dict
+           type will be inferred from the filename. Defaults to '.png'.
+        mpl_kwargs : dict, optional
            A dict of keyword arguments to be passed to matplotlib.
 
         >>> plot.save(mpl_kwargs={"bbox_inches": "tight"})

yt/visualization/tests/test_commons.py

@@ -21,18 +21,27 @@ def test_default(name, expected):
     "name, suffix, expected",
     [
         ("noext", ".png", "noext.png"),
+        ("noext", None, "noext.png"),
         ("nothing.png", ".png", "nothing.png"),
+        ("nothing.png", None, "nothing.png"),
         ("nothing.png", ".pdf", "nothing.pdf"),
         ("nothing.pdf", ".pdf", "nothing.pdf"),
+        ("nothing.pdf", None, "nothing.pdf"),
         ("nothing.pdf", ".png", "nothing.png"),
         ("version.1.2.3", ".png", "version.1.2.3.png"),
+        ("version.1.2.3", None, "version.1.2.3.png"),
         ("version.1.2.3", ".pdf", "version.1.2.3.pdf"),
     ],
 )
+@pytest.mark.filterwarnings(
+    r"ignore:Received two valid image formats '\w+' \(from filename\) "
+    r"and '\w+' \(from suffix\). The former is ignored.:UserWarning"
+)
 def test_custom_valid_ext(name, suffix, expected):
-    alt_suffix = suffix.replace(".", "")
     result1 = validate_image_name(name, suffix=suffix)
-    result2 = validate_image_name(name, suffix=alt_suffix)
-
     assert result1 == expected
-    assert result2 == expected
+
+    if suffix is not None:
+        alt_suffix = suffix.replace(".", "")
+        result2 = validate_image_name(name, suffix=alt_suffix)
+        assert result2 == expected

yt/visualization/tests/test_save.py

@@ -1,19 +1,30 @@
 import os
+import re
+
+import pytest
 
 import yt
 
+# avoid testing every supported format as some backends may be buggy on testing platforms
+FORMATS_TO_TEST = [".eps", ".pdf", ".svg"]
+
 
-def test_save_to_path(tmp_path):
+@pytest.fixture(scope="session")
+def simple_sliceplot():
     ds = yt.testing.fake_amr_ds()
     p = yt.SlicePlot(ds, "z", "Density")
+    yield p
+
+
+def test_save_to_path(simple_sliceplot, tmp_path):
+    p = simple_sliceplot
     p.save(f"{tmp_path}/")
     assert len(list((tmp_path).glob("*.png"))) == 1
 
 
-def test_save_to_missing_path(tmp_path):
+def test_save_to_missing_path(simple_sliceplot, tmp_path):
     # the missing layer should be created
-    ds = yt.testing.fake_amr_ds()
-    p = yt.SlicePlot(ds, "z", "Density")
+    p = simple_sliceplot
 
     # using forward slashes should work even on windows !
     save_path = os.path.join(tmp_path / "out") + "/"
@@ -22,13 +33,65 @@ def test_save_to_missing_path(tmp_path):
     assert len(list((tmp_path / "out").glob("*.png"))) == 1
 
 
-def test_save_to_missing_path_with_file_prefix(tmp_path):
+def test_save_to_missing_path_with_file_prefix(simple_sliceplot, tmp_path):
     # see issue
     # https://github.com/yt-project/yt/issues/3210
-    ds = yt.testing.fake_amr_ds()
-    p = yt.SlicePlot(ds, "z", "Density")
+    p = simple_sliceplot
     p.save(tmp_path.joinpath("out", "saymyname"))
     assert (tmp_path / "out").exists()
     output_files = list((tmp_path / "out").glob("*.png"))
     assert len(output_files) == 1
     assert output_files[0].stem.startswith("saymyname")  # you're goddamn right
+
+
+@pytest.mark.parametrize("ext", FORMATS_TO_TEST)
+def test_suffix_from_filename(ext, simple_sliceplot, tmp_path):
+    p = simple_sliceplot
+
+    target = (tmp_path / "myfile").with_suffix(ext)
+    # this shouldn't raise a warning, see issue
+    # https://github.com/yt-project/yt/issues/3667
+    p.save(target)
+    assert target.is_file()
+
+
+@pytest.mark.parametrize("ext", FORMATS_TO_TEST)
+def test_suffix_clashing(ext, simple_sliceplot, tmp_path):
+    if ext == ".png":
+        pytest.skip()
+
+    p = simple_sliceplot
+
+    target = (tmp_path / "myfile").with_suffix(ext)
+    expected_warning = re.compile(
+        r"Received two valid image formats '%s' \(from filename\) "
+        r"and 'png' \(from suffix\)\. The former is ignored\." % ext[1:]
+    )
+
+    with pytest.warns(UserWarning, match=expected_warning):
+        p.save(target, suffix="png")
+    output_files = list(tmp_path.glob("*.png"))
+    assert len(output_files) == 1
+    assert output_files[0].stem.startswith("myfile")
+    assert not list((tmp_path / "out").glob(f"*.{ext}"))
+
+
+def test_invalid_format_from_filename(simple_sliceplot, tmp_path):
+    p = simple_sliceplot
+    target = (tmp_path / "myfile").with_suffix(".nope")
+
+    p.save(target)
+    output_files = list(tmp_path.glob("*"))
+    assert len(output_files) == 1
+    # the output filename may contain a generated part
+    # it's not exactly clear if it's desirable or intented in this case
+    # so we just check conditions that should hold in any case
+    assert output_files[0].name.startswith("myfile.nope")
+    assert output_files[0].name.endswith(".png")
+
+
+def test_invalid_format_from_suffix(simple_sliceplot, tmp_path):
+    p = simple_sliceplot
+    target = tmp_path / "myfile"
+    with pytest.raises(ValueError, match=r"Unsupported file format 'nope'"):
+        p.save(target, suffix="nope")