xcb screengrab support

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:

src/PIL/ImageGrab.py

@@ -22,11 +22,18 @@ import tempfile
 
 from . import Image
 
-if sys.platform not in ["win32", "darwin"]:
-    raise ImportError("ImageGrab is macOS and Windows only")
+if sys.platform == "win32":
+    pass
+elif sys.platform == "darwin":
+    import os
+    import tempfile
+    import subprocess
+elif not Image.core.HAVE_XCB:
+    raise ImportError("ImageGrab requires Windows, macOS, or the XCB library")
 
 
-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,7 +45,8 @@ def grab(bbox=None, include_layered_windows=False, all_screens=False):
                 im_cropped = im.crop(bbox)
                 im.close()
                 return im_cropped
-    else:
+            return im
+        elif sys.platform == "win32":
             offset, size, data = Image.core.grabscreen(include_layered_windows, all_screens)
             im = Image.frombytes(
                 "RGB",
@@ -55,6 +63,22 @@ 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 OSError("XCB support not included")
+    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,8 +105,8 @@ 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

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

@@ -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;

src/display.c

@@ -815,3 +815,82 @@ 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)) {
+        xcb_disconnect(connection);
+        PyErr_SetString(PyExc_IOError, "X connection failed");
+        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) {
+        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) {
+        free(error);
+        xcb_disconnect(connection);
+        PyErr_SetString(PyExc_IOError, "X get_image failed");
+        return NULL;
+    }
+
+    /* store data in Python buffer */
+
+    buffer = PyBytes_FromStringAndSize((char*)xcb_get_image_data(reply),
+                                       xcb_get_image_data_length(reply));
+
+    free(reply);
+    xcb_disconnect(connection);
+
+    if (!buffer)
+        return NULL;
+
+    return Py_BuildValue("(ii)N", width, height, buffer);
+}
+
+#endif /* HAVE_XCB */