From 3f5fad3a27cbc7d0a5be1bfb2a0ca03c18c1279c Mon Sep 17 00:00:00 2001
From: Andrew Murray <radarhere@users.noreply.github.com>
Date: Sat, 19 Feb 2022 14:29:03 +1100
Subject: [PATCH] Fixed BUFR, GRIB and HDF5 stub saving

---
 Tests/test_file_bufrstub.py    | 32 ++++++++++++++++++++++++++++++++
 Tests/test_file_gribstub.py    | 32 ++++++++++++++++++++++++++++++++
 Tests/test_file_hdf5stub.py    | 32 ++++++++++++++++++++++++++++++++
 src/PIL/BufrStubImagePlugin.py |  2 +-
 src/PIL/GribStubImagePlugin.py |  2 +-
 src/PIL/Hdf5StubImagePlugin.py |  2 +-
 6 files changed, 99 insertions(+), 3 deletions(-)

diff --git a/Tests/test_file_bufrstub.py b/Tests/test_file_bufrstub.py
index 11acc1c88..e330404d6 100644
--- a/Tests/test_file_bufrstub.py
+++ b/Tests/test_file_bufrstub.py
@@ -45,3 +45,35 @@ def test_save(tmp_path):
     # Act / Assert: stub cannot save without an implemented handler
     with pytest.raises(OSError):
         im.save(tmpfile)
+
+
+def test_handler(tmp_path):
+    class TestHandler:
+        opened = False
+        loaded = False
+        saved = False
+
+        def open(self, im):
+            self.opened = True
+
+        def load(self, im):
+            self.loaded = True
+            return Image.new("RGB", (1, 1))
+
+        def save(self, im, fp, filename):
+            self.saved = True
+
+    handler = TestHandler()
+    BufrStubImagePlugin.register_handler(handler)
+    with Image.open(TEST_FILE) as im:
+        assert handler.opened
+        assert not handler.loaded
+
+        im.load()
+        assert handler.loaded
+
+        temp_file = str(tmp_path / "temp.bufr")
+        im.save(temp_file)
+        assert handler.saved
+
+    BufrStubImagePlugin._handler = None
diff --git a/Tests/test_file_gribstub.py b/Tests/test_file_gribstub.py
index e4930d8dc..fd427746e 100644
--- a/Tests/test_file_gribstub.py
+++ b/Tests/test_file_gribstub.py
@@ -45,3 +45,35 @@ def test_save(tmp_path):
     # Act / Assert: stub cannot save without an implemented handler
     with pytest.raises(OSError):
         im.save(tmpfile)
+
+
+def test_handler(tmp_path):
+    class TestHandler:
+        opened = False
+        loaded = False
+        saved = False
+
+        def open(self, im):
+            self.opened = True
+
+        def load(self, im):
+            self.loaded = True
+            return Image.new("RGB", (1, 1))
+
+        def save(self, im, fp, filename):
+            self.saved = True
+
+    handler = TestHandler()
+    GribStubImagePlugin.register_handler(handler)
+    with Image.open(TEST_FILE) as im:
+        assert handler.opened
+        assert not handler.loaded
+
+        im.load()
+        assert handler.loaded
+
+        temp_file = str(tmp_path / "temp.grib")
+        im.save(temp_file)
+        assert handler.saved
+
+    GribStubImagePlugin._handler = None
diff --git a/Tests/test_file_hdf5stub.py b/Tests/test_file_hdf5stub.py
index ff3397055..20b4b9619 100644
--- a/Tests/test_file_hdf5stub.py
+++ b/Tests/test_file_hdf5stub.py
@@ -46,3 +46,35 @@ def test_save():
             im.save(dummy_filename)
         with pytest.raises(OSError):
             Hdf5StubImagePlugin._save(im, dummy_fp, dummy_filename)
+
+
+def test_handler(tmp_path):
+    class TestHandler:
+        opened = False
+        loaded = False
+        saved = False
+
+        def open(self, im):
+            self.opened = True
+
+        def load(self, im):
+            self.loaded = True
+            return Image.new("RGB", (1, 1))
+
+        def save(self, im, fp, filename):
+            self.saved = True
+
+    handler = TestHandler()
+    Hdf5StubImagePlugin.register_handler(handler)
+    with Image.open(TEST_FILE) as im:
+        assert handler.opened
+        assert not handler.loaded
+
+        im.load()
+        assert handler.loaded
+
+        temp_file = str(tmp_path / "temp.h5")
+        im.save(temp_file)
+        assert handler.saved
+
+    Hdf5StubImagePlugin._handler = None
diff --git a/src/PIL/BufrStubImagePlugin.py b/src/PIL/BufrStubImagePlugin.py
index 48f21e1b3..9510f733e 100644
--- a/src/PIL/BufrStubImagePlugin.py
+++ b/src/PIL/BufrStubImagePlugin.py
@@ -59,7 +59,7 @@ class BufrStubImageFile(ImageFile.StubImageFile):
 
 
 def _save(im, fp, filename):
-    if _handler is None or not hasattr("_handler", "save"):
+    if _handler is None or not hasattr(_handler, "save"):
         raise OSError("BUFR save handler not installed")
     _handler.save(im, fp, filename)
 
diff --git a/src/PIL/GribStubImagePlugin.py b/src/PIL/GribStubImagePlugin.py
index b9bdd16e3..cc9bc2639 100644
--- a/src/PIL/GribStubImagePlugin.py
+++ b/src/PIL/GribStubImagePlugin.py
@@ -59,7 +59,7 @@ class GribStubImageFile(ImageFile.StubImageFile):
 
 
 def _save(im, fp, filename):
-    if _handler is None or not hasattr("_handler", "save"):
+    if _handler is None or not hasattr(_handler, "save"):
         raise OSError("GRIB save handler not installed")
     _handler.save(im, fp, filename)
 
diff --git a/src/PIL/Hdf5StubImagePlugin.py b/src/PIL/Hdf5StubImagePlugin.py
index 362f2d399..df11cf2a6 100644
--- a/src/PIL/Hdf5StubImagePlugin.py
+++ b/src/PIL/Hdf5StubImagePlugin.py
@@ -59,7 +59,7 @@ class HDF5StubImageFile(ImageFile.StubImageFile):
 
 
 def _save(im, fp, filename):
-    if _handler is None or not hasattr("_handler", "save"):
+    if _handler is None or not hasattr(_handler, "save"):
         raise OSError("HDF5 save handler not installed")
     _handler.save(im, fp, filename)