diff --git a/Tests/test_imagegrab.py b/Tests/test_imagegrab.py index 9768eb6ce..790847734 100644 --- a/Tests/test_imagegrab.py +++ b/Tests/test_imagegrab.py @@ -2,67 +2,72 @@ import subprocess 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), - ]: +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)) + + @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) - im = ImageGrab.grab(bbox=(10, 20, 50, 80)) - assert_image(im, im.mode, (40, 60)) + im2 = ImageGrab.grab(xdisplay="") + assert_image(im2, im2.mode, im2.size) + except IOError as e: + pytest.skip(str(e)) - 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") + @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"]) + 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() - - 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 + ) + 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) diff --git a/docs/reference/ImageGrab.rst b/docs/reference/ImageGrab.rst index e94e21cb9..ddd5bbbb5 100644 --- a/docs/reference/ImageGrab.rst +++ b/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() diff --git a/setup.py b/setup.py index 20dae91fa..3e1a812b6 100755 --- a/setup.py +++ b/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 diff --git a/src/PIL/ImageGrab.py b/src/PIL/ImageGrab.py index 059a71fe7..66e2e8560 100644 --- a/src/PIL/ImageGrab.py +++ b/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") diff --git a/src/PIL/features.py b/src/PIL/features.py index 0a1d5d611..ac06c0f71 100644 --- a/src/PIL/features.py +++ b/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) diff --git a/src/_imaging.c b/src/_imaging.c index f40b19e4d..0c3d766f3 100644 --- a/src/_imaging.c +++ b/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}, - {"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}, @@ -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; diff --git a/src/display.c b/src/display.c index 4c2faf9e0..21869b26e 100644 --- a/src/display.c +++ b/src/display.c @@ -815,3 +815,88 @@ error: } #endif /* _WIN32 */ + +/* -------------------------------------------------------------------- */ +/* X11 support */ + +#ifdef HAVE_XCB +#include + +/* -------------------------------------------------------------------- */ +/* 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 */