From f9c74825a6bcbc666d1c68d00267e73be21b5819 Mon Sep 17 00:00:00 2001 From: nulano Date: Fri, 20 Sep 2019 01:14:35 +0200 Subject: [PATCH] xcb screengrab fixes --- Tests/test_imagegrab.py | 97 +++++++++++++++++++++-------------------- setup.py | 1 + src/PIL/ImageGrab.py | 18 +++----- src/display.c | 16 ++++--- 4 files changed, 66 insertions(+), 66 deletions(-) diff --git a/Tests/test_imagegrab.py b/Tests/test_imagegrab.py index 9768eb6ce..63a954c94 100644 --- a/Tests/test_imagegrab.py +++ b/Tests/test_imagegrab.py @@ -2,67 +2,68 @@ 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), - ]: - 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: + @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) - class TestImageGrab: - @pytest.mark.skip(reason="ImageGrab ImportError") - def test_skip(self): - pass + 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 + @pytest.mark.skipif(not Image.core.HAVE_XCB, reason="requires XCB") + def test_grab_invalid_xdisplay(self): exception = None - # Act try: - from PIL import ImageGrab - - ImageGrab.__name__ # dummy to prevent Pyflakes warning + ImageGrab.grab(xdisplay="error.test:0.0") except Exception as e: exception = e - # Assert - if sys.platform in ["win32", "darwin"]: - assert exception is None + assert isinstance(exception, IOError) + assert str(exception).startswith("X connection failed") + + @pytest.mark.skipif( + sys.platform not in ("win32", "darwin"), reason="requires Windows or macOS" + ) + 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" + pytest.skip("ImageGrab.grabclipboard() is macOS and Windows only") + return + + im = ImageGrab.grabclipboard() + assert_image(im, im.mode, im.size) diff --git a/setup.py b/setup.py index b26656552..63ebf046c 100755 --- a/setup.py +++ b/setup.py @@ -823,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 239077b53..9ad91b28d 100644 --- a/src/PIL/ImageGrab.py +++ b/src/PIL/ImageGrab.py @@ -22,12 +22,12 @@ import tempfile from . import Image -if sys.platform == "win32": - pass -elif sys.platform == "darwin": +if sys.platform == "darwin": import os import tempfile import subprocess +elif sys.platform == "win32": + pass elif not Image.core.HAVE_XCB: raise ImportError("ImageGrab requires Windows, macOS, or the XCB library") @@ -65,17 +65,9 @@ def grab(bbox=None, include_layered_windows=False, all_screens=False, xdisplay=N 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") + raise AttributeError("XCB support not present") size, data = Image.core.grabscreen_x11(xdisplay) - im = Image.frombytes( - "RGB", - size, - data, - "raw", - "BGRX", - size[0] * 4, - 1, - ) + im = Image.frombytes("RGB", size, data, "raw", "BGRX", size[0] * 4, 1) if bbox: im = im.crop(bbox) return im diff --git a/src/display.c b/src/display.c index df2b4114f..d3930f5a5 100644 --- a/src/display.c +++ b/src/display.c @@ -845,8 +845,8 @@ PyImaging_GrabScreenX11(PyObject* self, PyObject* args) 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); - PyErr_SetString(PyExc_IOError, "X connection failed"); return NULL; } @@ -857,7 +857,8 @@ PyImaging_GrabScreenX11(PyObject* self, PyObject* args) break; } } - if (screen == NULL) { + 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; @@ -873,16 +874,21 @@ PyImaging_GrabScreenX11(PyObject* self, PyObject* args) 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); - 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)); + if (reply->depth == 24) { + buffer = PyBytes_FromStringAndSize((char*)xcb_get_image_data(reply), + xcb_get_image_data_length(reply)); + } else { + PyErr_Format(PyExc_IOError, "usupported bit depth: %i", reply->depth); + } free(reply); xcb_disconnect(connection);