Merge pull request #16709 from QuLogic/fix-png-write-io

timhoffm · web-flow · commit caa677e2ea0a · 2020-03-08T14:23:15.000+01:00
Fix saving PNGs to file objects in some places
diff --git a/lib/matplotlib/image.py b/lib/matplotlib/image.py
@@ -666,7 +666,7 @@ def write_png(self, fname):
         from matplotlib import _png
         im = self.to_rgba(self._A[::-1] if self.origin == 'lower' else self._A,
                           bytes=True, norm=True)
-        with open(fname, "wb") as file:
+        with cbook.open_file_cm(fname, "wb") as file:
             _png.write_png(im, file)
 
     def set_data(self, A):
diff --git a/lib/matplotlib/mathtext.py b/lib/matplotlib/mathtext.py
@@ -3433,7 +3433,7 @@ def to_png(self, filename, texstr, color='black', dpi=120, fontsize=14):
         from matplotlib import _png
         rgba, depth = self.to_rgba(
             texstr, color=color, dpi=dpi, fontsize=fontsize)
-        with open(filename, "wb") as file:
+        with cbook.open_file_cm(filename, "wb") as file:
             _png.write_png(rgba, file)
         return depth
 
diff --git a/lib/matplotlib/tests/test_mathtext.py b/lib/matplotlib/tests/test_mathtext.py
@@ -280,3 +280,14 @@ def test_single_minus_sign():
 def test_spaces(fig_test, fig_ref):
     fig_test.subplots().set_title(r"$1\,2\>3\ 4$")
     fig_ref.subplots().set_title(r"$1\/2\:3~4$")
+
+
+def test_math_to_image(tmpdir):
+    mathtext.math_to_image('$x^2$', str(tmpdir.join('example.png')))
+    mathtext.math_to_image('$x^2$', io.BytesIO())
+
+
+def test_mathtext_to_png(tmpdir):
+    mt = mathtext.MathTextParser('bitmap')
+    mt.to_png(str(tmpdir.join('example.png')), '$x^2$')
+    mt.to_png(io.BytesIO(), '$x^2$')
