Merge pull request #8516 from radarhere/imagegrab

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

Tests/test_imagegrab.py

@@ -57,6 +57,13 @@ 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, 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"])

docs/reference/ImageGrab.rst

@@ -9,7 +9,7 @@ or the clipboard to a PIL image memory.
 
 .. versionadded:: 1.1.3
 
-.. py:function:: grab(bbox=None, include_layered_windows=False, all_screens=False, xdisplay=None)
+.. py:function:: grab(bbox=None, include_layered_windows=False, all_screens=False, xdisplay=None, window=None)
 
     Take a snapshot of the screen. The pixels inside the bounding box are returned as
     an "RGBA" on macOS, or an "RGB" image otherwise. If the bounding box is omitted,
@@ -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 window:
+        HWND, to capture a single window. Windows only.
+
+        .. versionadded:: 11.2.0
     :return: An image
 
 .. py:function:: grabclipboard()

docs/releasenotes/11.2.0.rst

@@ -51,6 +51,15 @@ aligned using ``"justify"`` in :py:mod:`~PIL.ImageDraw`::
     draw.multiline_text((0, 0), "Multiline\ntext 1", align="justify")
     draw.multiline_textbbox((0, 0), "Multiline\ntext 2", align="justify")
 
+Specify window in ImageGrab on Windows
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+When using :py:meth:`~PIL.ImageGrab.grab`, a specific window can be selected using the
+HWND::
+
+    from PIL import ImageGrab
+    ImageGrab.grab(window=hwnd)
+
 Check for MozJPEG
 ^^^^^^^^^^^^^^^^^
 

src/PIL/ImageGrab.py

@@ -25,12 +25,17 @@ import tempfile
 
 from . import Image
 
+TYPE_CHECKING = False
+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",

src/display.c

@@ -286,29 +286,42 @@ PyImaging_DisplayModeWin32(PyObject *self, PyObject *args) {
 /* -------------------------------------------------------------------- */
 /* 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;
+    HANDLE dpiAwareness = NULL;
     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)
@@ -317,15 +330,28 @@ PyImaging_GrabScreenWin32(PyObject *self, PyObject *args) {
     SetThreadDpiAwarenessContext_function = (Func_SetThreadDpiAwarenessContext
     )GetProcAddress(user32, "SetThreadDpiAwarenessContext");
     if (SetThreadDpiAwarenessContext_function != NULL) {
+        GetWindowDpiAwarenessContext_function = (Func_GetWindowDpiAwarenessContext
+        )GetProcAddress(user32, "GetWindowDpiAwarenessContext");
+        if (screens == -1 && GetWindowDpiAwarenessContext_function != NULL) {
+            dpiAwareness = GetWindowDpiAwarenessContext_function(wnd);
+        }
         // DPI_AWARENESS_CONTEXT_PER_MONITOR_AWARE = ((DPI_CONTEXT_HANDLE)-3)
-        dpiAwareness = SetThreadDpiAwarenessContext_function((HANDLE)-3);
+        dpiAwareness = SetThreadDpiAwarenessContext_function(
+            dpiAwareness == NULL ? (HANDLE)-3 : dpiAwareness
+        );
     }
 
-    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);
@@ -337,6 +363,10 @@ PyImaging_GrabScreenWin32(PyObject *self, PyObject *args) {
 
     FreeLibrary(user32);
 
+    if (width == -1) {
+        goto error;
+    }
+
     bitmap = CreateCompatibleBitmap(screen, width, height);
     if (!bitmap) {
         goto error;
@@ -382,7 +412,11 @@ PyImaging_GrabScreenWin32(PyObject *self, PyObject *args) {
 
     DeleteObject(bitmap);
     DeleteDC(screen_copy);
-    DeleteDC(screen);
+    if (screens == -1) {
+        ReleaseDC(wnd, screen);
+    } else {
+        DeleteDC(screen);
+    }
 
     return Py_BuildValue("(ii)(ii)N", x, y, width, height, buffer);
 
@@ -390,7 +424,11 @@ error:
     PyErr_SetString(PyExc_OSError, "screen grab failed");
 
     DeleteDC(screen_copy);
-    DeleteDC(screen);
+    if (screens == -1) {
+        ReleaseDC(wnd, screen);
+    } else {
+        DeleteDC(screen);
+    }
 
     return NULL;
 }