Skip to content

Commit

Permalink
Allow handle to be supplied for ImageGrab.grab() on Windows
Browse files Browse the repository at this point in the history
  • Loading branch information
radarhere committed Nov 3, 2024
1 parent 5771f0e commit 356c3c1
Show file tree
Hide file tree
Showing 4 changed files with 51 additions and 7 deletions.
5 changes: 5 additions & 0 deletions Tests/test_imagegrab.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(handle=-1)

def test_grabclipboard(self) -> None:
if sys.platform == "darwin":
subprocess.call(["screencapture", "-cx"])
Expand Down
5 changes: 5 additions & 0 deletions docs/reference/ImageGrab.rst
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down
11 changes: 10 additions & 1 deletion src/PIL/ImageGrab.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
handle: int | ImageWin.HWND | None = None,
) -> Image.Image:
im: Image.Image
if xdisplay is None:
Expand All @@ -51,8 +56,12 @@ def grab(
return im_resized
return im
elif sys.platform == "win32":
if handle is not None:
all_screens = -1
offset, size, data = Image.core.grabscreen_win32(
include_layered_windows, all_screens
include_layered_windows,
all_screens,
int(handle) if handle is not None else 0,
)
im = Image.frombytes(
"RGB",
Expand Down
37 changes: 31 additions & 6 deletions src/display.c
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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

View check run for this annotation

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

View check run for this annotation

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);
Expand All @@ -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

View check run for this annotation

Codecov / codecov/patch

src/display.c#L388

Added line #L388 was not covered by tests
}

bitmap = CreateCompatibleBitmap(screen, width, height);
if (!bitmap) {
goto error;
Expand Down Expand Up @@ -420,7 +441,11 @@ PyImaging_GrabScreenWin32(PyObject *self, PyObject *args) {
PyErr_SetString(PyExc_OSError, "screen grab failed");

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

Check warning on line 445 in src/display.c

View check run for this annotation

Codecov / codecov/patch

src/display.c#L445

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

Check warning on line 447 in src/display.c

View check run for this annotation

Codecov / codecov/patch

src/display.c#L447

Added line #L447 was not covered by tests
}

return NULL;
}
Expand Down

0 comments on commit 356c3c1

Please sign in to comment.