Merge pull request #4615 from nulano/clipboard-png

Author: hugovk

Tests/test_imagegrab.py

@@ -1,10 +1,11 @@
+import os
 import subprocess
 import sys
 
 import pytest
 from PIL import Image, ImageGrab
 
-from .helper import assert_image
+from .helper import assert_image, assert_image_equal_tofile
 
 
 class TestImageGrab:
@@ -71,3 +72,27 @@ def test_grabclipboard(self):
 
         im = ImageGrab.grabclipboard()
         assert_image(im, im.mode, im.size)
+
+    @pytest.mark.skipif(sys.platform != "win32", reason="Windows only")
+    def test_grabclipboard_file(self):
+        p = subprocess.Popen(["powershell", "-command", "-"], stdin=subprocess.PIPE)
+        p.stdin.write(rb'Set-Clipboard -Path "Tests\images\hopper.gif"')
+        p.communicate()
+
+        im = ImageGrab.grabclipboard()
+        assert len(im) == 1
+        assert os.path.samefile(im[0], "Tests/images/hopper.gif")
+
+    @pytest.mark.skipif(sys.platform != "win32", reason="Windows only")
+    def test_grabclipboard_png(self):
+        p = subprocess.Popen(["powershell", "-command", "-"], stdin=subprocess.PIPE)
+        p.stdin.write(
+            rb"""$bytes = [System.IO.File]::ReadAllBytes("Tests\images\hopper.png")
+$ms = new-object System.IO.MemoryStream(, $bytes)
+[Reflection.Assembly]::LoadWithPartialName("System.Windows.Forms")
+[Windows.Forms.Clipboard]::SetData("PNG", $ms)"""
+        )
+        p.communicate()
+
+        im = ImageGrab.grabclipboard()
+        assert_image_equal_tofile(im, "Tests/images/hopper.png")

src/PIL/ImageGrab.py

@@ -93,12 +93,28 @@ def grabclipboard():
         os.unlink(filepath)
         return im
     elif sys.platform == "win32":
-        data = Image.core.grabclipboard_win32()
+        fmt, data = Image.core.grabclipboard_win32()
+        if fmt == "file":  # CF_HDROP
+            import struct
+
+            o = struct.unpack_from("I", data)[0]
+            if data[16] != 0:
+                files = data[o:].decode("utf-16le").split("\0")
+            else:
+                files = data[o:].decode("mbcs").split("\0")
+            return files[: files.index("")]
         if isinstance(data, bytes):
-            from . import BmpImagePlugin
             import io
 
-            return BmpImagePlugin.DibImageFile(io.BytesIO(data))
-        return data
+            data = io.BytesIO(data)
+            if fmt == "png":
+                from . import PngImagePlugin
+
+                return PngImagePlugin.PngImageFile(data)
+            elif fmt == "DIB":
+                from . import BmpImagePlugin
+
+                return BmpImagePlugin.DibImageFile(data)
+        return None
     else:
         raise NotImplementedError("ImageGrab.grabclipboard() is macOS and Windows only")

src/display.c

@@ -502,33 +502,45 @@ PyObject*
 PyImaging_GrabClipboardWin32(PyObject* self, PyObject* args)
 {
     int clip;
-    HANDLE handle;
+    HANDLE handle = NULL;
     int size;
     void* data;
     PyObject* result;
+    UINT format;
+    UINT formats[] = { CF_DIB, CF_DIBV5, CF_HDROP, RegisterClipboardFormatA("PNG"), 0 };
+    LPCSTR format_names[] = { "DIB", "DIB", "file", "png", NULL };
 
-    clip = OpenClipboard(NULL);
-    /* FIXME: check error status */
+    if (!OpenClipboard(NULL)) {
+        PyErr_SetString(PyExc_OSError, "failed to open clipboard");
+        return NULL;
+    }
+
+    // find best format as set by clipboard owner
+    format = 0;
+    while (!handle && (format = EnumClipboardFormats(format))) {
+        for (UINT i = 0; formats[i] != 0; i++) {
+            if (format == formats[i]) {
+                handle = GetClipboardData(format);
+                format = i;
+                break;
+            }
+        }
+    }
 
-    handle = GetClipboardData(CF_DIB);
     if (!handle) {
-        /* FIXME: add CF_HDROP support to allow cut-and-paste from
-           the explorer */
         CloseClipboard();
-        Py_INCREF(Py_None);
-        return Py_None;
+        return Py_BuildValue("zO", NULL, Py_None);
     }
 
-    size = GlobalSize(handle);
     data = GlobalLock(handle);
+    size = GlobalSize(handle);
 
     result = PyBytes_FromStringAndSize(data, size);
 
     GlobalUnlock(handle);
-
     CloseClipboard();
 
-    return result;
+    return Py_BuildValue("zN", format_names[format], result);
 }
 
 /* -------------------------------------------------------------------- */