Merge pull request #1365 from radarhere/psdraw

Fixed PSDraw stdout Python 3 compatibility

PIL/EpsImagePlugin.py

@@ -376,9 +376,10 @@ def _save(im, fp, filename, eps=1):
             pass
 
     base_fp = fp
-    fp = NoCloseStream(fp)
-    if sys.version_info[0] > 2:
-        fp = io.TextIOWrapper(fp, encoding='latin-1')
+    if fp != sys.stdout:
+        fp = NoCloseStream(fp)
+        if sys.version_info[0] > 2:
+            fp = io.TextIOWrapper(fp, encoding='latin-1')
 
     if eps:
         #

PIL/ImageFile.py

@@ -461,6 +461,9 @@ def _save(im, fp, tile, bufsize=0):
     # But, it would need at least the image size in most cases. RawEncode is
     # a tricky case.
     bufsize = max(MAXBLOCK, bufsize, im.size[0] * 4)  # see RawEncode.c
+    if fp == sys.stdout:
+        fp.flush()
+        return
     try:
         fh = fp.fileno()
         fp.flush()

PIL/PSDraw.py

@@ -16,11 +16,12 @@
 #
 
 from PIL import EpsImagePlugin
-
+import sys
 
 ##
 # Simple Postscript graphics interface.
 
+
 class PSDraw(object):
     """
     Sets up printing to the given file. If **file** is omitted,
@@ -29,12 +30,11 @@ class PSDraw(object):
 
     def __init__(self, fp=None):
         if not fp:
-            import sys
             fp = sys.stdout
         self.fp = fp
 
     def _fp_write(self, to_write):
-        if bytes is str:
+        if bytes is str or self.fp == sys.stdout:
             self.fp.write(to_write)
         else:
             self.fp.write(bytes(to_write, 'UTF-8'))
@@ -47,7 +47,7 @@ class PSDraw(object):
                        "/showpage { } def\n"
                        "%%EndComments\n"
                        "%%BeginDocument\n")
-        # self.fp_write(ERROR_PS)  # debugging!
+        # self._fp_write(ERROR_PS)  # debugging!
         self._fp_write(EDROFF_PS)
         self._fp_write(VDI_PS)
         self._fp_write("%%EndProlog\n")

Tests/test_psdraw.py

@@ -1,25 +1,17 @@
 from helper import unittest, PillowTestCase
 
+from PIL import Image, PSDraw
+import os
+import sys
+
 
 class TestPsDraw(PillowTestCase):
 
-    def test_draw_postscript(self):
-
-        # Based on Pillow tutorial, but there is no textsize:
-        # http://pillow.readthedocs.org/en/latest/handbook/tutorial.html
-
-        # Arrange
-        from PIL import Image
-        from PIL import PSDraw
-        tempfile = self.tempfile('temp.ps')
-        fp = open(tempfile, "wb")
-
+    def _create_document(self, ps):
         im = Image.open("Tests/images/hopper.ppm")
         title = "hopper"
         box = (1*72, 2*72, 7*72, 10*72)  # in points
 
-        # Act
-        ps = PSDraw.PSDraw(fp)
         ps.begin_document(title)
 
         # draw diagonal lines in a cross
@@ -35,14 +27,43 @@ class TestPsDraw(PillowTestCase):
         ps.text((3*72, 4*72), title)
 
         ps.end_document()
+
+    def test_draw_postscript(self):
+
+        # Based on Pillow tutorial, but there is no textsize:
+        # http://pillow.readthedocs.org/en/latest/handbook/tutorial.html
+
+        # Arrange
+        tempfile = self.tempfile('temp.ps')
+        fp = open(tempfile, "wb")
+
+        # Act
+        ps = PSDraw.PSDraw(fp)
+        self._create_document(ps)
         fp.close()
 
         # Assert
         # Check non-zero file was created
-        import os
         self.assertTrue(os.path.isfile(tempfile))
         self.assertGreater(os.path.getsize(tempfile), 0)
 
+    def test_stdout(self):
+        # Temporarily redirect stdout
+        try:
+            from cStringIO import StringIO
+        except ImportError:
+            from io import StringIO
+        old_stdout = sys.stdout
+        sys.stdout = mystdout = StringIO()
+
+        ps = PSDraw.PSDraw()
+        self._create_document(ps)
+
+        # Reset stdout
+        sys.stdout = old_stdout
+
+        self.assertNotEqual(mystdout.getvalue(), "")
+
 
 if __name__ == '__main__':
     unittest.main()