Merge pull request #4260 from nulano/imagegrab_xcb

ImageGrab.grab() for Linux with XCB

Tests/test_imagegrab.py

@@ -2,13 +2,15 @@ import subprocess
 import sys
 
 import pytest
+from PIL import Image, ImageGrab
 
 from .helper import assert_image
 
-try:
-    from PIL import ImageGrab
 
 class TestImageGrab:
+    @pytest.mark.skipif(
+        sys.platform not in ("win32", "darwin"), reason="requires Windows or macOS"
+    )
     def test_grab(self):
         for im in [
             ImageGrab.grab(),
@@ -20,13 +22,40 @@ try:
         im = ImageGrab.grab(bbox=(10, 20, 50, 80))
         assert_image(im, im.mode, (40, 60))
 
+    @pytest.mark.skipif(not Image.core.HAVE_XCB, reason="requires XCB")
+    def test_grab_x11(self):
+        try:
+            if sys.platform not in ("win32", "darwin"):
+                im = ImageGrab.grab()
+                assert_image(im, im.mode, im.size)
+
+            im2 = ImageGrab.grab(xdisplay="")
+            assert_image(im2, im2.mode, im2.size)
+        except IOError as e:
+            pytest.skip(str(e))
+
+    @pytest.mark.skipif(Image.core.HAVE_XCB, reason="tests missing XCB")
+    def test_grab_no_xcb(self):
+        if sys.platform not in ("win32", "darwin"):
+            with pytest.raises(IOError) as e:
+                ImageGrab.grab()
+            assert str(e.value).startswith("Pillow was built without XCB support")
+
+        with pytest.raises(IOError) as e:
+            ImageGrab.grab(xdisplay="")
+        assert str(e.value).startswith("Pillow was built without XCB support")
+
+    @pytest.mark.skipif(not Image.core.HAVE_XCB, reason="requires XCB")
+    def test_grab_invalid_xdisplay(self):
+        with pytest.raises(IOError) as e:
+            ImageGrab.grab(xdisplay="error.test:0.0")
+        assert str(e.value).startswith("X connection failed")
+
     def test_grabclipboard(self):
         if sys.platform == "darwin":
             subprocess.call(["screencapture", "-cx"])
-            else:
+        elif sys.platform == "win32":
-                p = subprocess.Popen(
+            p = subprocess.Popen(["powershell", "-command", "-"], stdin=subprocess.PIPE)
-                    ["powershell", "-command", "-"], stdin=subprocess.PIPE
-                )
             p.stdin.write(
                 b"""[Reflection.Assembly]::LoadWithPartialName("System.Drawing")
 [Reflection.Assembly]::LoadWithPartialName("System.Windows.Forms")
@@ -34,35 +63,11 @@ $bmp = New-Object Drawing.Bitmap 200, 200
 [Windows.Forms.Clipboard]::SetImage($bmp)"""
             )
             p.communicate()
+        else:
+            with pytest.raises(NotImplementedError) as e:
+                ImageGrab.grabclipboard()
+            assert str(e.value) == "ImageGrab.grabclipboard() is macOS and Windows only"
+            return
 
         im = ImageGrab.grabclipboard()
         assert_image(im, im.mode, im.size)
-
-
-except ImportError:
-
-    class TestImageGrab:
-        @pytest.mark.skip(reason="ImageGrab ImportError")
-        def test_skip(self):
-            pass
-
-
-class TestImageGrabImport:
-    def test_import(self):
-        # Arrange
-        exception = None
-
-        # Act
-        try:
-            from PIL import ImageGrab
-
-            ImageGrab.__name__  # dummy to prevent Pyflakes warning
-        except Exception as e:
-            exception = e
-
-        # Assert
-        if sys.platform in ["win32", "darwin"]:
-            assert exception is None
-        else:
-            assert isinstance(exception, ImportError)
-            assert str(exception) == "ImageGrab is macOS and Windows only"

docs/reference/ImageGrab.rst

@@ -11,13 +11,13 @@ or the clipboard to a PIL image memory.
 
 .. versionadded:: 1.1.3
 
-.. py:function:: PIL.ImageGrab.grab(bbox=None, include_layered_windows=False, all_screens=False)
+.. py:function:: PIL.ImageGrab.grab(bbox=None, include_layered_windows=False, all_screens=False, xdisplay=None)
 
     Take a snapshot of the screen. The pixels inside the bounding box are
     returned as an "RGB" image on Windows or "RGBA" on macOS.
     If the bounding box is omitted, the entire screen is copied.
 
-    .. versionadded:: 1.1.3 (Windows), 3.0.0 (macOS)
+    .. versionadded:: 1.1.3 (Windows), 3.0.0 (macOS), 7.1.0 (Linux (X11))
 
     :param bbox: What region to copy. Default is the entire screen.
                  Note that on Windows OS, the top-left point may be negative if ``all_screens=True`` is used.
@@ -27,6 +27,11 @@ or the clipboard to a PIL image memory.
     :param all_screens: Capture all monitors. Windows OS only.
 
         .. versionadded:: 6.2.0
+
+    :param xdisplay: X11 Display address. Pass ``None`` to grab the default system screen.
+                     Pass ``""`` to grab the default X11 screen on Windows or macOS.
+
+        .. versionadded:: 7.1.0
     :return: An image
 
 .. py:function:: PIL.ImageGrab.grabclipboard()

setup.py

@@ -286,6 +286,7 @@ class pil_build_ext(build_ext):
             "webpmux",
             "jpeg2000",
             "imagequant",
+            "xcb",
         ]
 
         required = {"jpeg", "zlib"}
@@ -681,6 +682,12 @@ class pil_build_ext(build_ext):
                 ):
                     feature.webpmux = "libwebpmux"
 
+        if feature.want("xcb"):
+            _dbg("Looking for xcb")
+            if _find_include_file(self, "xcb/xcb.h"):
+                if _find_library_file(self, "xcb"):
+                    feature.xcb = "xcb"
+
         for f in feature:
             if not getattr(feature, f) and feature.require(f):
                 if f in ("jpeg", "zlib"):
@@ -715,6 +722,9 @@ class pil_build_ext(build_ext):
         if feature.tiff:
             libs.append(feature.tiff)
             defs.append(("HAVE_LIBTIFF", None))
+        if feature.xcb:
+            libs.append(feature.xcb)
+            defs.append(("HAVE_XCB", None))
         if sys.platform == "win32":
             libs.extend(["kernel32", "user32", "gdi32"])
         if struct.unpack("h", b"\0\1")[0] == 1:
@@ -813,6 +823,7 @@ class pil_build_ext(build_ext):
             (feature.lcms, "LITTLECMS2"),
             (feature.webp, "WEBP"),
             (feature.webpmux, "WEBPMUX"),
+            (feature.xcb, "XCB (X protocol)"),
         ]
 
         all = 1

src/PIL/ImageGrab.py

@@ -15,18 +15,18 @@
 # See the README file for information on usage and redistribution.
 #
 
-import os
-import subprocess
 import sys
-import tempfile
 
 from . import Image
 
-if sys.platform not in ["win32", "darwin"]:
+if sys.platform == "darwin":
-    raise ImportError("ImageGrab is macOS and Windows only")
+    import os
+    import tempfile
+    import subprocess
 
 
-def grab(bbox=None, include_layered_windows=False, all_screens=False):
+def grab(bbox=None, include_layered_windows=False, all_screens=False, xdisplay=None):
+    if xdisplay is None:
         if sys.platform == "darwin":
             fh, filepath = tempfile.mkstemp(".png")
             os.close(fh)
@@ -38,8 +38,11 @@ def grab(bbox=None, include_layered_windows=False, all_screens=False):
                 im_cropped = im.crop(bbox)
                 im.close()
                 return im_cropped
-    else:
+            return im
-        offset, size, data = Image.core.grabscreen(include_layered_windows, all_screens)
+        elif sys.platform == "win32":
+            offset, size, data = Image.core.grabscreen_win32(
+                include_layered_windows, all_screens
+            )
             im = Image.frombytes(
                 "RGB",
                 size,
@@ -55,6 +58,14 @@ def grab(bbox=None, include_layered_windows=False, all_screens=False):
                 left, top, right, bottom = bbox
                 im = im.crop((left - x0, top - y0, right - x0, bottom - y0))
             return im
+    # use xdisplay=None for default display on non-win32/macOS systems
+    if not Image.core.HAVE_XCB:
+        raise IOError("Pillow was built without XCB support")
+    size, data = Image.core.grabscreen_x11(xdisplay)
+    im = Image.frombytes("RGB", size, data, "raw", "BGRX", size[0] * 4, 1)
+    if bbox:
+        im = im.crop(bbox)
+    return im
 
 
 def grabclipboard():
@@ -81,11 +92,13 @@ def grabclipboard():
             im.load()
         os.unlink(filepath)
         return im
-    else:
+    elif sys.platform == "win32":
-        data = Image.core.grabclipboard()
+        data = Image.core.grabclipboard_win32()
         if isinstance(data, bytes):
             from . import BmpImagePlugin
             import io
 
             return BmpImagePlugin.DibImageFile(io.BytesIO(data))
         return data
+    else:
+        raise NotImplementedError("ImageGrab.grabclipboard() is macOS and Windows only")

src/PIL/features.py

@@ -56,6 +56,7 @@ features = {
     "raqm": ("PIL._imagingft", "HAVE_RAQM"),
     "libjpeg_turbo": ("PIL._imaging", "HAVE_LIBJPEGTURBO"),
     "libimagequant": ("PIL._imaging", "HAVE_LIBIMAGEQUANT"),
+    "xcb": ("PIL._imaging", "HAVE_XCB"),
 }
 
 
@@ -132,6 +133,7 @@ def pilinfo(out=None, supported_formats=True):
         ("libtiff", "LIBTIFF"),
         ("raqm", "RAQM (Bidirectional Text)"),
         ("libimagequant", "LIBIMAGEQUANT (Quantization method)"),
+        ("xcb", "XCB (X protocol)"),
     ]:
         if check(name):
             print("---", feature, "support ok", file=out)

src/_imaging.c

@@ -3817,6 +3817,9 @@ extern PyObject* PyImaging_ListWindowsWin32(PyObject* self, PyObject* args);
 extern PyObject* PyImaging_EventLoopWin32(PyObject* self, PyObject* args);
 extern PyObject* PyImaging_DrawWmf(PyObject* self, PyObject* args);
 #endif
+#ifdef HAVE_XCB
+extern PyObject* PyImaging_GrabScreenX11(PyObject* self, PyObject* args);
+#endif
 
 /* Experimental path stuff (in path.c) */
 extern PyObject* PyPath_Create(ImagingObject* self, PyObject* args);
@@ -3889,13 +3892,16 @@ static PyMethodDef functions[] = {
 #ifdef _WIN32
     {"display", (PyCFunction)PyImaging_DisplayWin32, 1},
     {"display_mode", (PyCFunction)PyImaging_DisplayModeWin32, 1},
-    {"grabscreen", (PyCFunction)PyImaging_GrabScreenWin32, 1},
+    {"grabscreen_win32", (PyCFunction)PyImaging_GrabScreenWin32, 1},
-    {"grabclipboard", (PyCFunction)PyImaging_GrabClipboardWin32, 1},
+    {"grabclipboard_win32", (PyCFunction)PyImaging_GrabClipboardWin32, 1},
     {"createwindow", (PyCFunction)PyImaging_CreateWindowWin32, 1},
     {"eventloop", (PyCFunction)PyImaging_EventLoopWin32, 1},
     {"listwindows", (PyCFunction)PyImaging_ListWindowsWin32, 1},
     {"drawwmf", (PyCFunction)PyImaging_DrawWmf, 1},
 #endif
+#ifdef HAVE_XCB
+    {"grabscreen_x11", (PyCFunction)PyImaging_GrabScreenX11, 1},
+#endif
 
     /* Utilities */
     {"getcodecstatus", (PyCFunction)_getcodecstatus, 1},
@@ -4015,6 +4021,12 @@ setup_module(PyObject* m) {
   }
 #endif
 
+#ifdef HAVE_XCB
+    PyModule_AddObject(m, "HAVE_XCB", Py_True);
+#else
+    PyModule_AddObject(m, "HAVE_XCB", Py_False);
+#endif
+
     PyDict_SetItemString(d, "PILLOW_VERSION", PyUnicode_FromString(version));
 
     return 0;

src/display.c

@@ -815,3 +815,88 @@ error:
 }
 
 #endif /* _WIN32 */
+
+/* -------------------------------------------------------------------- */
+/* X11 support	*/
+
+#ifdef HAVE_XCB
+#include <xcb/xcb.h>
+
+/* -------------------------------------------------------------------- */
+/* X11 screen grabber */
+
+PyObject*
+PyImaging_GrabScreenX11(PyObject* self, PyObject* args)
+{
+    int width, height;
+    char* display_name;
+    xcb_connection_t* connection;
+    int screen_number;
+    xcb_screen_iterator_t iter;
+    xcb_screen_t* screen = NULL;
+    xcb_get_image_reply_t* reply;
+    xcb_generic_error_t* error;
+    PyObject* buffer = NULL;
+
+    if (!PyArg_ParseTuple(args, "|z", &display_name))
+        return NULL;
+
+    /* connect to X and get screen data */
+
+    connection = xcb_connect(display_name, &screen_number);
+    if (xcb_connection_has_error(connection)) {
+        PyErr_Format(PyExc_IOError, "X connection failed: error %i", xcb_connection_has_error(connection));
+        xcb_disconnect(connection);
+        return NULL;
+    }
+
+    iter = xcb_setup_roots_iterator(xcb_get_setup(connection));
+    for (; iter.rem; --screen_number, xcb_screen_next(&iter)) {
+        if (screen_number == 0) {
+            screen = iter.data;
+            break;
+        }
+    }
+    if (screen == NULL || screen->root == 0) {
+        // this case is usually caught with "X connection failed: error 6" above
+        xcb_disconnect(connection);
+        PyErr_SetString(PyExc_IOError, "X screen not found");
+        return NULL;
+    }
+
+    width = screen->width_in_pixels;
+    height = screen->height_in_pixels;
+
+    /* get image data */
+
+    reply = xcb_get_image_reply(connection,
+                                xcb_get_image(connection, XCB_IMAGE_FORMAT_Z_PIXMAP, screen->root,
+                                              0, 0, width, height, 0x00ffffff),
+                                &error);
+    if (reply == NULL) {
+        PyErr_Format(PyExc_IOError, "X get_image failed: error %i (%i, %i, %i)",
+                     error->error_code, error->major_code, error->minor_code, error->resource_id);
+        free(error);
+        xcb_disconnect(connection);
+        return NULL;
+    }
+
+    /* store data in Python buffer */
+
+    if (reply->depth == 24) {
+        buffer = PyBytes_FromStringAndSize((char*)xcb_get_image_data(reply),
+                                           xcb_get_image_data_length(reply));
+    } else {
+        PyErr_Format(PyExc_IOError, "unsupported bit depth: %i", reply->depth);
+    }
+
+    free(reply);
+    xcb_disconnect(connection);
+
+    if (!buffer)
+        return NULL;
+
+    return Py_BuildValue("(ii)N", width, height, buffer);
+}
+
+#endif /* HAVE_XCB */