Skip to content

Commit

Permalink
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Allow window to be supplied for ImageGrab.grab() on Windows
Browse files Browse the repository at this point in the history
radarhere committed Nov 4, 2024

Verified

This commit was created on GitHub.com and signed with GitHub’s verified signature. The key has expired.
1 parent 5771f0e commit 607acbf
Showing 4 changed files with 56 additions and 8 deletions.
5 changes: 5 additions & 0 deletions Tests/test_imagegrab.py
Original file line number Diff line number Diff line change
@@ -57,6 +57,11 @@ def test_grab_invalid_xdisplay(self) -> None:
ImageGrab.grab(xdisplay="error.test:0.0")
assert str(e.value).startswith("X connection failed")

@pytest.mark.skipif(sys.platform != "win32", reason="Windows only")
def test_grab_invalid_handle(self) -> None:
with pytest.raises(OSError):
ImageGrab.grab(window=-1)

def test_grabclipboard(self) -> None:
if sys.platform == "darwin":
subprocess.call(["screencapture", "-cx"])
5 changes: 5 additions & 0 deletions docs/reference/ImageGrab.rst
Original file line number Diff line number Diff line change
@@ -39,6 +39,11 @@ or the clipboard to a PIL image memory.
You can check X11 support using :py:func:`PIL.features.check_feature` with ``feature="xcb"``.

.. versionadded:: 7.1.0

:param handle:
HWND, to capture a single window. Windows only.

.. versionadded:: 11.1.0
:return: An image

.. py:function:: grabclipboard()
11 changes: 10 additions & 1 deletion src/PIL/ImageGrab.py
Original file line number Diff line number Diff line change
@@ -22,15 +22,20 @@
import subprocess
import sys
import tempfile
from typing import TYPE_CHECKING

from . import Image

if TYPE_CHECKING:
from . import ImageWin


def grab(
bbox: tuple[int, int, int, int] | None = None,
include_layered_windows: bool = False,
all_screens: bool = False,
xdisplay: str | None = None,
window: int | ImageWin.HWND | None = None,
) -> Image.Image:
im: Image.Image
if xdisplay is None:
@@ -51,8 +56,12 @@ def grab(
return im_resized
return im
elif sys.platform == "win32":
if window is not None:
all_screens = -1
offset, size, data = Image.core.grabscreen_win32(
include_layered_windows, all_screens
include_layered_windows,
all_screens,
int(window) if window is not None else 0,
)
im = Image.frombytes(
"RGB",
43 changes: 36 additions & 7 deletions src/display.c
Original file line number Diff line number Diff line change
@@ -320,25 +320,36 @@ typedef HANDLE(__stdcall *Func_SetThreadDpiAwarenessContext)(HANDLE);

PyObject *
PyImaging_GrabScreenWin32(PyObject *self, PyObject *args) {
int x = 0, y = 0, width, height;
int includeLayeredWindows = 0, all_screens = 0;
int x = 0, y = 0, width = -1, height;
int includeLayeredWindows = 0, screens = 0;
HBITMAP bitmap;
BITMAPCOREHEADER core;
HDC screen, screen_copy;
HWND wnd;
DWORD rop;
PyObject *buffer;
HANDLE dpiAwareness;
HMODULE user32;
Func_SetThreadDpiAwarenessContext SetThreadDpiAwarenessContext_function;

if (!PyArg_ParseTuple(args, "|ii", &includeLayeredWindows, &all_screens)) {
if (!PyArg_ParseTuple(
args, "|ii" F_HANDLE, &includeLayeredWindows, &screens, &wnd
)) {
return NULL;
}

/* step 1: create a memory DC large enough to hold the
entire screen */

screen = CreateDC("DISPLAY", NULL, NULL, NULL);
if (screens == -1) {
screen = GetDC(wnd);
if (screen == NULL) {
PyErr_SetString(PyExc_OSError, "unable to get device context for handle");
return NULL;
}
} else {
screen = CreateDC("DISPLAY", NULL, NULL, NULL);
}
screen_copy = CreateCompatibleDC(screen);

// added in Windows 10 (1607)
@@ -351,11 +362,17 @@ PyImaging_GrabScreenWin32(PyObject *self, PyObject *args) {
dpiAwareness = SetThreadDpiAwarenessContext_function((HANDLE)-3);
}

if (all_screens) {
if (screens == 1) {
x = GetSystemMetrics(SM_XVIRTUALSCREEN);
y = GetSystemMetrics(SM_YVIRTUALSCREEN);
width = GetSystemMetrics(SM_CXVIRTUALSCREEN);
height = GetSystemMetrics(SM_CYVIRTUALSCREEN);
} else if (screens == -1) {
RECT rect;

Check warning on line 371 in src/display.c

Codecov / codecov/patch

src/display.c#L371

Added line #L371 was not covered by tests
if (GetClientRect(wnd, &rect)) {
width = rect.right;
height = rect.bottom;

Check warning on line 374 in src/display.c

Codecov / codecov/patch

src/display.c#L373-L374

Added lines #L373 - L374 were not covered by tests
}
} else {
width = GetDeviceCaps(screen, HORZRES);
height = GetDeviceCaps(screen, VERTRES);
@@ -367,6 +384,10 @@ PyImaging_GrabScreenWin32(PyObject *self, PyObject *args) {

FreeLibrary(user32);

if (width == -1) {
goto error;

Check warning on line 388 in src/display.c

Codecov / codecov/patch

src/display.c#L388

Added line #L388 was not covered by tests
}

bitmap = CreateCompatibleBitmap(screen, width, height);
if (!bitmap) {
goto error;
@@ -412,15 +433,23 @@ PyImaging_GrabScreenWin32(PyObject *self, PyObject *args) {

DeleteObject(bitmap);
DeleteDC(screen_copy);
DeleteDC(screen);
if (screens == -1) {
ReleaseDC(wnd, screen);

Check warning on line 437 in src/display.c

Codecov / codecov/patch

src/display.c#L437

Added line #L437 was not covered by tests
} else {
DeleteDC(screen);
}

return Py_BuildValue("(ii)(ii)N", x, y, width, height, buffer);

error:
PyErr_SetString(PyExc_OSError, "screen grab failed");

DeleteDC(screen_copy);
DeleteDC(screen);
if (screens == -1) {
ReleaseDC(wnd, screen);

Check warning on line 449 in src/display.c

Codecov / codecov/patch

src/display.c#L449

Added line #L449 was not covered by tests
} else {
DeleteDC(screen);

Check warning on line 451 in src/display.c

Codecov / codecov/patch

src/display.c#L451

Added line #L451 was not covered by tests
}

return NULL;
}

0 comments on commit 607acbf

Please sign in to comment.