ImageGrab.grab() for Linux with XCB

Tests/test_imagegrab.py

@@ -2,67 +2,72 @@
 import sys
 
 import pytest
+from PIL import Image, ImageGrab
 
 from .helper import assert_image
 
-try:
-    from PIL import ImageGrab
 
-    class TestImageGrab:
-        def test_grab(self):
-            for im in [
-                ImageGrab.grab(),
-                ImageGrab.grab(include_layered_windows=True),
-                ImageGrab.grab(all_screens=True),
-            ]:
-                assert_image(im, im.mode, im.size)
-
-            im = ImageGrab.grab(bbox=(10, 20, 50, 80))
-            assert_image(im, im.mode, (40, 60))
-
-        def test_grabclipboard(self):
-            if sys.platform == "darwin":
-                subprocess.call(["screencapture", "-cx"])
-            else:
-                p = subprocess.Popen(
-                    ["powershell", "-command", "-"], stdin=subprocess.PIPE
-                )
-                p.stdin.write(
-                    b"""[Reflection.Assembly]::LoadWithPartialName("System.Drawing")
-[Reflection.Assembly]::LoadWithPartialName("System.Windows.Forms")
-$bmp = New-Object Drawing.Bitmap 200, 200
-[Windows.Forms.Clipboard]::SetImage($bmp)"""
-                )
-                p.communicate()
-
-            im = ImageGrab.grabclipboard()
+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(),
+            ImageGrab.grab(include_layered_windows=True),
+            ImageGrab.grab(all_screens=True),
+        ]:
             assert_image(im, im.mode, im.size)
 
+        im = ImageGrab.grab(bbox=(10, 20, 50, 80))
+        assert_image(im, im.mode, (40, 60))
 
-except ImportError:
-
-    class TestImageGrab:
-        @pytest.mark.skip(reason="ImageGrab ImportError")
-        def test_skip(self):
-            pass
+    @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))
 
-class TestImageGrabImport:
-    def test_import(self):
-        # Arrange
-        exception = None
+    @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")
 
-        # Act
-        try:
-            from PIL import ImageGrab
+        with pytest.raises(IOError) as e:
+            ImageGrab.grab(xdisplay="")
+        assert str(e.value).startswith("Pillow was built without XCB support")
 
-            ImageGrab.__name__  # dummy to prevent Pyflakes warning
-        except Exception as e:
-            exception = e
+    @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")
 
-        # Assert
-        if sys.platform in ["win32", "darwin"]:
-            assert exception is None
+    def test_grabclipboard(self):
+        if sys.platform == "darwin":
+            subprocess.call(["screencapture", "-cx"])
+        elif sys.platform == "win32":
+            p = subprocess.Popen(["powershell", "-command", "-"], stdin=subprocess.PIPE)
+            p.stdin.write(
+                b"""[Reflection.Assembly]::LoadWithPartialName("System.Drawing")
+[Reflection.Assembly]::LoadWithPartialName("System.Windows.Forms")
+$bmp = New-Object Drawing.Bitmap 200, 200
+[Windows.Forms.Clipboard]::SetImage($bmp)"""
+            )
+            p.communicate()
         else:
-            assert isinstance(exception, ImportError)
-            assert str(exception) == "ImageGrab is macOS and Windows only"
+            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)

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 feature:
             "webpmux",
             "jpeg2000",
             "imagequant",
+            "xcb",
         ]
 
         required = {"jpeg", "zlib"}
@@ -681,6 +682,12 @@ def build_extensions(self):
                 ):
                     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 @@ def build_extensions(self):
         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 @@ def summary_report(self, feature):
             (feature.lcms, "LITTLECMS2"),
             (feature.webp, "WEBP"),
             (feature.webpmux, "WEBPMUX"),
+            (feature.xcb, "XCB (X protocol)"),
         ]
 
         all = 1

src/PIL/ImageGrab.py

@@ -15,45 +15,56 @@
 # 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"]:
-    raise ImportError("ImageGrab is macOS and Windows only")
+if sys.platform == "darwin":
+    import os
+    import tempfile
+    import subprocess
 
 
-def grab(bbox=None, include_layered_windows=False, all_screens=False):
-    if sys.platform == "darwin":
-        fh, filepath = tempfile.mkstemp(".png")
-        os.close(fh)
-        subprocess.call(["screencapture", "-x", filepath])
-        im = Image.open(filepath)
-        im.load()
-        os.unlink(filepath)
-        if bbox:
-            im_cropped = im.crop(bbox)
-            im.close()
-            return im_cropped
-    else:
-        offset, size, data = Image.core.grabscreen(include_layered_windows, all_screens)
-        im = Image.frombytes(
-            "RGB",
-            size,
-            data,
-            # RGB, 32-bit line padding, origin lower left corner
-            "raw",
-            "BGR",
-            (size[0] * 3 + 3) & -4,
-            -1,
-        )
-        if bbox:
-            x0, y0 = offset
-            left, top, right, bottom = bbox
-            im = im.crop((left - x0, top - y0, right - x0, bottom - y0))
+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)
+            subprocess.call(["screencapture", "-x", filepath])
+            im = Image.open(filepath)
+            im.load()
+            os.unlink(filepath)
+            if bbox:
+                im_cropped = im.crop(bbox)
+                im.close()
+                return im_cropped
+            return im
+        elif sys.platform == "win32":
+            offset, size, data = Image.core.grabscreen_win32(
+                include_layered_windows, all_screens
+            )
+            im = Image.frombytes(
+                "RGB",
+                size,
+                data,
+                # RGB, 32-bit line padding, origin lower left corner
+                "raw",
+                "BGR",
+                (size[0] * 3 + 3) & -4,
+                -1,
+            )
+            if bbox:
+                x0, y0 = offset
+                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
 
 
@@ -81,11 +92,13 @@ def grabclipboard():
             im.load()
         os.unlink(filepath)
         return im
-    else:
-        data = Image.core.grabclipboard()
+    elif sys.platform == "win32":
+        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 @@ def get_supported_codecs():
     "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

@@ -3781,6 +3781,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);
@@ -3853,13 +3856,16 @@ static PyMethodDef functions[] = {
 #ifdef _WIN32
     {"display", (PyCFunction)PyImaging_DisplayWin32, 1},
     {"display_mode", (PyCFunction)PyImaging_DisplayModeWin32, 1},
-    {"grabscreen", (PyCFunction)PyImaging_GrabScreenWin32, 1},
-    {"grabclipboard", (PyCFunction)PyImaging_GrabClipboardWin32, 1},
+    {"grabscreen_win32", (PyCFunction)PyImaging_GrabScreenWin32, 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},
@@ -3979,6 +3985,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;