Allow window to be supplied for ImageGrab.grab() on Windows

This commit is contained in:
Andrew Murray 2024-11-05 07:05:39 +11:00
parent 5771f0ec37
commit 607acbf95e
4 changed files with 56 additions and 8 deletions

View File

@ -57,6 +57,11 @@ class TestImageGrab:
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"])

View File

@ -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()

View File

@ -22,15 +22,20 @@ import shutil
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",

View File

@ -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 */
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;
if (GetClientRect(wnd, &rect)) {
width = rect.right;
height = rect.bottom;
}
} 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;
}
bitmap = CreateCompatibleBitmap(screen, width, height);
if (!bitmap) {
goto error;
@ -412,7 +433,11 @@ PyImaging_GrabScreenWin32(PyObject *self, PyObject *args) {
DeleteObject(bitmap);
DeleteDC(screen_copy);
if (screens == -1) {
ReleaseDC(wnd, screen);
} else {
DeleteDC(screen);
}
return Py_BuildValue("(ii)(ii)N", x, y, width, height, buffer);
@ -420,7 +445,11 @@ error:
PyErr_SetString(PyExc_OSError, "screen grab failed");
DeleteDC(screen_copy);
if (screens == -1) {
ReleaseDC(wnd, screen);
} else {
DeleteDC(screen);
}
return NULL;
}