DM-22073: Add support for writing matplotlib figures

.gitignore

@@ -15,3 +15,4 @@ bin/*
 .coverage
 .pytest_cache/
 .mypy_cache/
+.vscode/

config/datastores/formatters.yaml

@@ -34,3 +34,4 @@ Packages: lsst.daf.butler.formatters.pickleFormatter.PickleFormatter
 PropertyList: lsst.daf.butler.formatters.pickleFormatter.PickleFormatter
 PropertySet: lsst.daf.butler.formatters.pickleFormatter.PickleFormatter
 NumpyArray: lsst.daf.butler.formatters.pickleFormatter.PickleFormatter
+Plot: lsst.daf.butler.formatters.matplotlibFormatter.MatplotlibFormatter

config/storageClasses.yaml

@@ -134,3 +134,5 @@ storageClasses:
     pytype: lsst.ip.isr.StrayLightData
   Thumbnail:
     pytype: numpy.ndarray
+  Plot:
+    pytype: matplotlib.figure.Figure

python/lsst/daf/butler/formatters/matplotlibFormatter.py

@@ -0,0 +1,43 @@
+# This file is part of daf_butler.
+#
+# Developed for the LSST Data Management System.
+# This product includes software developed by the LSST Project
+# (http://www.lsst.org).
+# See the COPYRIGHT file at the top-level directory of this distribution
+# for details of code ownership.
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program.  If not, see <http://www.gnu.org/licenses/>.
+
+"""Support for writing matplotlib figures."""
+
+__all__ = ("MatplotlibFormatter",)
+
+from lsst.daf.butler.formatters.fileFormatter import FileFormatter
+
+
+class MatplotlibFormatter(FileFormatter):
+    """Interface for writing matplotlib figures.
+    """
+
+    extension = ".png"
+    """Matplotlib figures are always written in PNG format."""
+
+    def _readFile(self, path, pytype=None):
+        # docstring inherited from FileFormatter._readFile
+        raise NotImplementedError(
+            f"matplotlib figures cannot be read by the butler; path is {path}")
+
+    def _writeFile(self, inMemoryDataset):
+        # docstring inherited from FileFormatter._writeFile
+        inMemoryDataset.savefig(self.fileDescriptor.location.path)

tests/test_matplotlibFormatter.py

@@ -0,0 +1,85 @@
+# This file is part of daf_butler.
+#
+# Developed for the LSST Data Management System.
+# This product includes software developed by the LSST Project
+# (http://www.lsst.org).
+# See the COPYRIGHT file at the top-level directory of this distribution
+# for details of code ownership.
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program.  If not, see <http://www.gnu.org/licenses/>.
+
+"""Tests for MatplotlibFormatter.
+"""
+
+import unittest
+import tempfile
+import os
+import shutil
+
+try:
+    import matplotlib
+    matplotlib.use("Agg")  # noqa:E402
+    from matplotlib import pyplot
+except ImportError:
+    pyplot = None
+
+import numpy as np
+from lsst.daf.butler import Butler, DatasetType
+import filecmp
+import urllib
+
+TESTDIR = os.path.abspath(os.path.dirname(__file__))
+
+
+@unittest.skipIf(pyplot is None,
+                 "skipping test because matplotlib import failed")
+class MatplotlibFormatterTestCase(unittest.TestCase):
+    """Test for MatplotlibFormatter.
+    """
+
+    def setUp(self):
+        self.root = tempfile.mkdtemp(dir=TESTDIR)
+        Butler.makeRepo(self.root)
+
+    def tearDown(self):
+        if os.path.exists(self.root):
+            shutil.rmtree(self.root, ignore_errors=True)
+
+    def testMatplotlibFormatter(self):
+        butler = Butler(self.root, run="testrun")
+        datasetType = DatasetType("test_plot", [], "Plot",
+                                  universe=butler.registry.dimensions)
+        butler.registry.registerDatasetType(datasetType)
+        pyplot.imshow(np.random.randn(3, 4))
+        ref = butler.put(pyplot.gcf(), datasetType)
+        parsed = urllib.parse.urlparse(butler.getUri(ref))
+        with tempfile.NamedTemporaryFile(suffix=".png") as file:
+            pyplot.gcf().savefig(file.name)
+            self.assertTrue(
+                filecmp.cmp(
+                    parsed.path,
+                    file.name,
+                    shallow=True
+                )
+            )
+        self.assertTrue(butler.datasetExists(ref))
+        with self.assertRaises(ValueError):
+            butler.get(ref)
+        butler.remove(ref)
+        with self.assertRaises(LookupError):
+            butler.datasetExists(ref)
+
+
+if __name__ == "__main__":
+    unittest.main()