Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Allow HWND to be passed to ImageGrab.grab() on Windows #8516

Open
wants to merge 11 commits into
base: main
Choose a base branch
from
7 changes: 7 additions & 0 deletions Tests/test_imagegrab.py
Original file line number Diff line number Diff line change
Expand Up @@ -57,6 +57,13 @@ 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, match="unable to get device context for handle"):
ImageGrab.grab(window=-1)
with pytest.raises(OSError, match="screen grab failed"):
ImageGrab.grab(window=0)

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,
window: 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 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",
Expand Down
60 changes: 51 additions & 9 deletions src/display.c
Original file line number Diff line number Diff line change
Expand Up @@ -316,29 +316,42 @@
/* -------------------------------------------------------------------- */
/* Windows screen grabber */

typedef HANDLE(__stdcall *Func_GetWindowDpiAwarenessContext)(HANDLE);
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_GetWindowDpiAwarenessContext GetWindowDpiAwarenessContext_function;
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 @@ -347,15 +360,32 @@
SetThreadDpiAwarenessContext_function = (Func_SetThreadDpiAwarenessContext
)GetProcAddress(user32, "SetThreadDpiAwarenessContext");
if (SetThreadDpiAwarenessContext_function != NULL) {
// DPI_AWARENESS_CONTEXT_PER_MONITOR_AWARE = ((DPI_CONTEXT_HANDLE)-3)
dpiAwareness = SetThreadDpiAwarenessContext_function((HANDLE)-3);
if (screens == -1) {
GetWindowDpiAwarenessContext_function = (Func_GetWindowDpiAwarenessContext
)GetProcAddress(user32, "GetWindowDpiAwarenessContext");
DPI_AWARENESS_CONTEXT dpiAwarenessContext =
GetWindowDpiAwarenessContext_function(wnd);
if (dpiAwarenessContext != NULL) {
radarhere marked this conversation as resolved.
Show resolved Hide resolved
dpiAwareness =
SetThreadDpiAwarenessContext_function(dpiAwarenessContext);

Check warning on line 370 in src/display.c

View check run for this annotation

Codecov / codecov/patch

src/display.c#L369-L370

Added lines #L369 - L370 were not covered by tests
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It should probably never happen that SetThreadDpiAwarenessContext is found but GetWindowDpiAwarenessContext is not.

However, if it somehow does happen, we would then use dpiAwareness before its initialized after measuring the window size. So we should ideally either fallback to PER_MONITOR_AWARE context if GetWindowDpiAwarenessContext is not found, or skip the reset a few lines later.

}
} else {
// DPI_AWARENESS_CONTEXT_PER_MONITOR_AWARE = ((DPI_CONTEXT_HANDLE)-3)
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;
if (GetClientRect(wnd, &rect)) {
width = rect.right;
height = rect.bottom;

Check warning on line 387 in src/display.c

View check run for this annotation

Codecov / codecov/patch

src/display.c#L386-L387

Added lines #L386 - L387 were not covered by tests
}
} else {
width = GetDeviceCaps(screen, HORZRES);
height = GetDeviceCaps(screen, VERTRES);
Expand All @@ -367,6 +397,10 @@

FreeLibrary(user32);

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

bitmap = CreateCompatibleBitmap(screen, width, height);
if (!bitmap) {
goto error;
Expand Down Expand Up @@ -412,15 +446,23 @@

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

Check warning on line 450 in src/display.c

View check run for this annotation

Codecov / codecov/patch

src/display.c#L450

Added line #L450 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);
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

You added this for the error: branch but not for the success branch.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ok, I've updated the commit to add it or success as well.

} else {
DeleteDC(screen);

Check warning on line 464 in src/display.c

View check run for this annotation

Codecov / codecov/patch

src/display.c#L464

Added line #L464 was not covered by tests
}

return NULL;
}
Expand Down
Loading