Fix writing to all manner of StringIO-like objects in PS backend.

Author: mdboom

lib/matplotlib/__init__.py

@@ -115,11 +115,16 @@
 def compare_versions(a, b):
     "return True if a is greater than or equal to b"
     if a:
+        if six.PY3:
+            if isinstance(a, bytes):
+                a = a.decode('ascii')
+            if isinstance(b, bytes):
+                b = b.decode('ascii')
         a = distutils.version.LooseVersion(a)
         b = distutils.version.LooseVersion(b)
-        if a>=b: return True
-        else: return False
-    else: return False
+        return a >= b
+    else:
+        return False
 
 try:
     import pyparsing

lib/matplotlib/backends/backend_ps.py

@@ -6,6 +6,7 @@
                         unicode_literals)
 
 import six
+from six.moves import StringIO
 
 import glob, math, os, shutil, sys, time
 def _fn_name(): return sys._getframe(1).f_code.co_name
@@ -24,7 +25,7 @@ def _fn_name(): return sys._getframe(1).f_code.co_name
      FigureManagerBase, FigureCanvasBase
 
 from matplotlib.cbook import is_string_like, get_realpath_and_stat, \
-    is_writable_file_like, maxdict
+    is_writable_file_like, maxdict, file_requires_unicode
 from matplotlib.mlab import quad2cubic
 from matplotlib.figure import Figure
 
@@ -1085,7 +1086,7 @@ def write(self, *kl, **kwargs):
 
             self._pswriter = NullWriter()
         else:
-            self._pswriter = six.moves.cStringIO()
+            self._pswriter = io.StringIO()
 
 
         # mixed mode rendering
@@ -1105,8 +1106,6 @@ def write(self, *kl, **kwargs):
         self.figure.set_edgecolor(origedgecolor)
 
         def print_figure_impl():
-            fh = io.TextIOWrapper(raw_fh, encoding="ascii")
-
             # write the PostScript headers
             if isEPSF: print("%!PS-Adobe-3.0 EPSF-3.0", file=fh)
             else: print("%!PS-Adobe-3.0", file=fh)
@@ -1157,7 +1156,7 @@ def print_figure_impl():
                             fh.flush()
                             convert_ttf_to_ps(
                                 font_filename.encode(sys.getfilesystemencoding()),
-                                raw_fh, fonttype, glyph_ids)
+                                fh, fonttype, glyph_ids)
             print("end", file=fh)
             print("%%EndProlog", file=fh)
 
@@ -1181,22 +1180,31 @@ def print_figure_impl():
             if not isEPSF: print("%%EOF", file=fh)
             fh.flush()
 
-            if six.PY3:
-                fh.detach()
-
         if rcParams['ps.usedistiller']:
             # We are going to use an external program to process the output.
             # Write to a temporary file.
             fd, tmpfile = mkstemp()
-            with io.open(fd, 'wb') as raw_fh:
+            with io.open(fd, 'w', encoding='ascii') as fh:
                 print_figure_impl()
         else:
             # Write directly to outfile.
             if passed_in_file_object:
-                raw_fh = outfile
+                requires_unicode = file_requires_unicode(outfile)
+
+                if (not requires_unicode and
+                    (six.PY3 or not isinstance(outfile, StringIO))):
+                    fh = io.TextIOWrapper(outfile, encoding="ascii")
+                    # Prevent the io.TextIOWrapper from closing the
+                    # underlying file
+                    def do_nothing():
+                        pass
+                    fh.close = do_nothing
+                else:
+                    fh = outfile
+
                 print_figure_impl()
             else:
-                with io.open(outfile, 'wb') as raw_fh:
+                with io.open(outfile, 'w', encoding='ascii') as fh:
                     print_figure_impl()
 
         if rcParams['ps.usedistiller']:
@@ -1206,8 +1214,12 @@ def print_figure_impl():
                 xpdf_distill(tmpfile, isEPSF, ptype=papertype, bbox=bbox)
 
             if passed_in_file_object:
-                with io.open(tmpfile, 'rb') as fh:
-                    outfile.write(fh.read())
+                if file_requires_unicode(outfile):
+                    with io.open(tmpfile, 'rb') as fh:
+                        outfile.write(fh.read().decode('ascii'))
+                else:
+                    with io.open(tmpfile, 'rb') as fh:
+                        outfile.write(fh.read())
             else:
                 with io.open(outfile, 'w') as fh:
                     pass
@@ -1224,7 +1236,12 @@ def _print_figure_tex(self, outfile, format, dpi, facecolor, edgecolor,
         package. These files are processed to yield the final ps or eps file.
         """
         isEPSF = format == 'eps'
-        title = outfile
+        if is_string_like(outfile):
+            title = outfile
+        elif is_writable_file_like(outfile):
+            title = None
+        else:
+            raise ValueError("outfile must be a path or a file-like object")
 
         self.figure.dpi = 72 # ignore the dpi kwarg
         width, height = self.figure.get_size_inches()
@@ -1252,7 +1269,7 @@ def write(self, *kl, **kwargs):
 
             self._pswriter = NullWriter()
         else:
-            self._pswriter = six.moves.cStringIO()
+            self._pswriter = io.StringIO()
 
 
         # mixed mode rendering
@@ -1273,11 +1290,7 @@ def write(self, *kl, **kwargs):
 
         # write to a temp file, we'll move it to outfile when done
         fd, tmpfile = mkstemp()
-        if six.PY3:
-            fh = io.open(fd, 'w', encoding='ascii')
-        else:
-            fh = io.open(fd, 'wb')
-        with fh:
+        with io.open(fd, 'w', encoding='ascii') as fh:
             # write the Encapsulated PostScript headers
             print("%!PS-Adobe-3.0 EPSF-3.0", file=fh)
             if title: print("%%Title: "+title, file=fh)
@@ -1357,17 +1370,13 @@ def write(self, *kl, **kwargs):
             else: gs_distill(tmpfile, isEPSF, ptype=papertype, bbox=bbox,
                              rotated=psfrag_rotated)
 
-        is_file = False
-        if six.PY3:
-            if isinstance(outfile, io.IOBase):
-                is_file = True
-        else:
-            if isinstance(outfile, file):
-                is_file = True
-
-        if is_file:
-            with io.open(tmpfile, 'rb') as fh:
-                outfile.write(fh.read())
+        if is_writable_file_like(outfile):
+            if file_requires_unicode(outfile):
+                with io.open(tmpfile, 'rb') as fh:
+                    outfile.write(fh.read().decode('ascii'))
+            else:
+                with io.open(tmpfile, 'rb') as fh:
+                    outfile.write(fh.read())
         else:
             with io.open(outfile, 'wb') as fh:
                 pass

lib/matplotlib/cbook.py

@@ -693,6 +693,19 @@ def is_writable_file_like(obj):
     return hasattr(obj, 'write') and six.callable(obj.write)
 
 
+def file_requires_unicode(x):
+    """
+    Returns `True` if the given writable file-like object requires Unicode
+    to be written to it.
+    """
+    try:
+        x.write(b'')
+    except TypeError:
+        return True
+    else:
+        return False
+
+
 def is_scalar(obj):
     'return true if *obj* is not string like and is not iterable'
     return not is_string_like(obj) and not iterable(obj)

lib/matplotlib/tests/test_backend_ps.py

@@ -1,17 +1,81 @@
+# -*- coding: utf-8 -*-
+
 from __future__ import (absolute_import, division, print_function,
                         unicode_literals)
 
+import io
+import re
+
 import six
+from nose.plugins.skip import SkipTest
 
+import matplotlib
 import matplotlib.pyplot as plt
 from matplotlib.testing.decorators import cleanup
 
 
-@cleanup
-def test_savefig_to_stringio():
+def _test_savefig_to_stringio(format='ps'):
+    buffers = [
+        six.moves.StringIO(),
+        io.StringIO(),
+        io.BytesIO()]
+
     plt.figure()
     plt.plot([0, 1], [0, 1])
-    plt.savefig(six.moves.StringIO(), format="ps")
+    plt.title("Déjà vu")
+    for buffer in buffers:
+        plt.savefig(buffer, format=format)
+
+    values = [x.getvalue() for x in buffers]
+
+    if six.PY3:
+        values = [
+            values[0].encode('ascii'),
+            values[1].encode('ascii'),
+            values[2]]
+
+    # Remove comments from the output.  This includes things that
+    # could change from run to run, such as the time.
+    values = [re.sub(b'%%.*?\n', b'', x) for x in values]
+
+    assert values[0] == values[1]
+    assert values[1] == values[2]
+
+
+@cleanup
+def test_savefig_to_stringio():
+    _test_savefig_to_stringio()
+
+
+@cleanup
+def test_savefig_to_stringio_with_distiller():
+    matplotlib.rcParams['ps.usedistiller'] = 'ghostscript'
+    _test_savefig_to_stringio()
+
+
+@cleanup
+def test_savefig_to_stringio_with_usetex():
+    if not matplotlib.checkdep_tex():
+        raise SkipTest("This test requires a TeX installation")
+
+    matplotlib.rcParams['text.latex.unicode'] = True
+    matplotlib.rcParams['text.usetex'] = True
+    _test_savefig_to_stringio()
+
+
+@cleanup
+def test_savefig_to_stringio_eps():
+    _test_savefig_to_stringio(format='eps')
+
+
+@cleanup
+def test_savefig_to_stringio_with_usetex_eps():
+    if not matplotlib.checkdep_tex():
+        raise SkipTest("This test requires a TeX installation")
+
+    matplotlib.rcParams['text.latex.unicode'] = True
+    matplotlib.rcParams['text.usetex'] = True
+    _test_savefig_to_stringio(format='eps')
 
 
 if __name__ == '__main__':

lib/matplotlib/texmanager.py

@@ -413,7 +413,8 @@ def make_dvi(self, tex, fontsize):
                 raise RuntimeError(
                     ('LaTeX was not able to process the following '
                      'string:\n%s\nHere is the full report generated by '
-                     'LaTeX: \n\n' % repr(tex)) + report)
+                     'LaTeX: \n\n' % repr(tex.encode('unicode_escape')) +
+                     report))
             else:
                 mpl.verbose.report(report, 'debug')
             for fname in glob.glob(basefile + '*'):

src/_ttconv.cpp

@@ -49,9 +49,15 @@ class PythonFileWriter : public TTStreamWriter
         if (_write_method)
         {
             #if PY3K
-            result = PyObject_CallFunction(_write_method, (char *)"y", a);
-            #else
             result = PyObject_CallFunction(_write_method, (char *)"s", a);
+            #else
+            PyObject* decoded = NULL;
+            decoded = PyUnicode_FromString(a);
+            if (decoded == NULL) {
+                throw PythonExceptionOccurred();
+            }
+            result = PyObject_CallFunction(_write_method, "O", decoded);
+            Py_DECREF(decoded);
             #endif
             if (! result)
             {