display.c

/*
 * The Python Imaging Library.
 *
 * display support (and other windows-related stuff)
 *
 * History:
 * 1996-05-13 fl  Windows DIB support
 * 1996-05-21 fl  Added palette stuff
 * 1996-05-28 fl  Added display_mode stuff
 * 1997-09-21 fl  Added draw primitive
 * 2001-09-17 fl  Added ImagingGrabScreen (from _grabscreen.c)
 * 2002-05-12 fl  Added ImagingListWindows
 * 2002-11-19 fl  Added clipboard support
 * 2002-11-25 fl  Added GetDC/ReleaseDC helpers
 * 2003-05-21 fl  Added create window support (including window callback)
 * 2003-09-05 fl  Added fromstring/tostring methods
 * 2009-03-14 fl  Added WMF support (from pilwmf)
 *
 * Copyright (c) 1997-2003 by Secret Labs AB.
 * Copyright (c) 1996-1997 by Fredrik Lundh.
 *
 * See the README file for information on usage and redistribution.
 */

#define PY_SSIZE_T_CLEAN
#include "Python.h"

#include "libImaging/Imaging.h"

/* -------------------------------------------------------------------- */
/* Windows DIB support */

#ifdef _WIN32

#include "libImaging/ImDib.h"

#if SIZEOF_VOID_P == 8
#define F_HANDLE "K"
#else
#define F_HANDLE "k"
#endif

typedef struct {
    PyObject_HEAD
    ImagingDIB dib;
} ImagingDisplayObject;

static PyTypeObject ImagingDisplayType;

static ImagingDisplayObject*
_new(const char* mode, int xsize, int ysize)
{
    ImagingDisplayObject *display;

    if (PyType_Ready(&ImagingDisplayType) < 0) {
        return NULL;
    }

    display = PyObject_New(ImagingDisplayObject, &ImagingDisplayType);
    if (display == NULL) {
        return NULL;
    }

    display->dib = ImagingNewDIB(mode, xsize, ysize);
    if (!display->dib) {
        Py_DECREF(display);
        return NULL;
    }

    return display;
}

static void
_delete(ImagingDisplayObject* display)
{
    if (display->dib) {
        ImagingDeleteDIB(display->dib);
    }
    PyObject_Del(display);
}

static PyObject*
_expose(ImagingDisplayObject* display, PyObject* args)
{
    HDC hdc;
    if (!PyArg_ParseTuple(args, F_HANDLE, &hdc)) {
        return NULL;
    }

    ImagingExposeDIB(display->dib, hdc);

    Py_INCREF(Py_None);
    return Py_None;
}

static PyObject*
_draw(ImagingDisplayObject* display, PyObject* args)
{
    HDC hdc;
    int dst[4];
    int src[4];
    if (!PyArg_ParseTuple(args, F_HANDLE "(iiii)(iiii)", &hdc,
                          dst+0, dst+1, dst+2, dst+3,
                          src+0, src+1, src+2, src+3)) {
        return NULL;
    }

    ImagingDrawDIB(display->dib, hdc, dst, src);

    Py_INCREF(Py_None);
    return Py_None;
}

extern Imaging PyImaging_AsImaging(PyObject *op);

static PyObject*
_paste(ImagingDisplayObject* display, PyObject* args)
{
    Imaging im;

    PyObject* op;
    int xy[4];
    xy[0] = xy[1] = xy[2] = xy[3] = 0;
    if (!PyArg_ParseTuple(args, "O|(iiii)", &op, xy+0, xy+1, xy+2, xy+3)) {
        return NULL;
    }
    im = PyImaging_AsImaging(op);
    if (!im) {
        return NULL;
    }

    if (xy[2] <= xy[0]) {
        xy[2] = xy[0] + im->xsize;
    }
    if (xy[3] <= xy[1]) {
        xy[3] = xy[1] + im->ysize;
    }

    ImagingPasteDIB(display->dib, im, xy);

    Py_INCREF(Py_None);
    return Py_None;
}

static PyObject*
_query_palette(ImagingDisplayObject* display, PyObject* args)
{
    HDC hdc;
    int status;

    if (!PyArg_ParseTuple(args, F_HANDLE, &hdc)) {
        return NULL;
    }

    status = ImagingQueryPaletteDIB(display->dib, hdc);

    return Py_BuildValue("i", status);
}

static PyObject*
_getdc(ImagingDisplayObject* display, PyObject* args)
{
    HWND window;
    HDC dc;

    if (!PyArg_ParseTuple(args, F_HANDLE, &window)) {
        return NULL;
    }

    dc = GetDC(window);
    if (!dc) {
        PyErr_SetString(PyExc_OSError, "cannot create dc");
        return NULL;
    }

    return Py_BuildValue(F_HANDLE, dc);
}

static PyObject*
_releasedc(ImagingDisplayObject* display, PyObject* args)
{
    HWND window;
    HDC dc;

    if (!PyArg_ParseTuple(args, F_HANDLE F_HANDLE, &window, &dc)) {
        return NULL;
    }

    ReleaseDC(window, dc);

    Py_INCREF(Py_None);
    return Py_None;
}

static PyObject*
_frombytes(ImagingDisplayObject* display, PyObject* args)
{
    char* ptr;
    Py_ssize_t bytes;

    if (!PyArg_ParseTuple(args, "y#:frombytes", &ptr, &bytes)) {
        return NULL;
    }

    if (display->dib->ysize * display->dib->linesize != bytes) {
        PyErr_SetString(PyExc_ValueError, "wrong size");
        return NULL;
    }

    memcpy(display->dib->bits, ptr, bytes);

    Py_INCREF(Py_None);
    return Py_None;
}

static PyObject*
_tobytes(ImagingDisplayObject* display, PyObject* args)
{
    if (!PyArg_ParseTuple(args, ":tobytes")) {
        return NULL;
    }

    return PyBytes_FromStringAndSize(
        display->dib->bits, display->dib->ysize * display->dib->linesize
        );
}

static struct PyMethodDef methods[] = {
    {"draw", (PyCFunction)_draw, 1},
    {"expose", (PyCFunction)_expose, 1},
    {"paste", (PyCFunction)_paste, 1},
    {"query_palette", (PyCFunction)_query_palette, 1},
    {"getdc", (PyCFunction)_getdc, 1},
    {"releasedc", (PyCFunction)_releasedc, 1},
    {"frombytes", (PyCFunction)_frombytes, 1},
    {"tobytes", (PyCFunction)_tobytes, 1},
    {NULL, NULL} /* sentinel */
};

static PyObject*
_getattr_mode(ImagingDisplayObject* self, void* closure)
{
    return Py_BuildValue("s", self->dib->mode);
}

static PyObject*
_getattr_size(ImagingDisplayObject* self, void* closure)
{
    return Py_BuildValue("ii", self->dib->xsize, self->dib->ysize);
}

static struct PyGetSetDef getsetters[] = {
    { "mode",   (getter) _getattr_mode },
    { "size",   (getter) _getattr_size },
    { NULL }
};

static PyTypeObject ImagingDisplayType = {
    PyVarObject_HEAD_INIT(NULL, 0)
    "ImagingDisplay",            /*tp_name*/
    sizeof(ImagingDisplayObject),/*tp_size*/
    0,                           /*tp_itemsize*/
    /* methods */
    (destructor)_delete,        /*tp_dealloc*/
    0,                          /*tp_print*/
    0,                          /*tp_getattr*/
    0,                          /*tp_setattr*/
    0,                          /*tp_compare*/
    0,                          /*tp_repr*/
    0,                          /*tp_as_number */
    0,                          /*tp_as_sequence */
    0,                          /*tp_as_mapping */
    0,                          /*tp_hash*/
    0,                          /*tp_call*/
    0,                          /*tp_str*/
    0,                          /*tp_getattro*/
    0,                          /*tp_setattro*/
    0,                          /*tp_as_buffer*/
    Py_TPFLAGS_DEFAULT,         /*tp_flags*/
    0,                          /*tp_doc*/
    0,                          /*tp_traverse*/
    0,                          /*tp_clear*/
    0,                          /*tp_richcompare*/
    0,                          /*tp_weaklistoffset*/
    0,                          /*tp_iter*/
    0,                          /*tp_iternext*/
    methods,                    /*tp_methods*/
    0,                          /*tp_members*/
    getsetters,                 /*tp_getset*/
};

PyObject*
PyImaging_DisplayWin32(PyObject* self, PyObject* args)
{
    ImagingDisplayObject* display;
    char *mode;
    int xsize, ysize;

    if (!PyArg_ParseTuple(args, "s(ii)", &mode, &xsize, &ysize)) {
        return NULL;
    }

    display = _new(mode, xsize, ysize);
    if (display == NULL) {
        return NULL;
    }

    return (PyObject*) display;
}

PyObject*
PyImaging_DisplayModeWin32(PyObject* self, PyObject* args)
{
    char *mode;
    int size[2];

    mode = ImagingGetModeDIB(size);

    return Py_BuildValue("s(ii)", mode, size[0], size[1]);
}

/* -------------------------------------------------------------------- */
/* Windows screen grabber */

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;
    HBITMAP bitmap;
    BITMAPCOREHEADER core;
    HDC screen, screen_copy;
    DWORD rop;
    PyObject* buffer;
    HANDLE dpiAwareness;
    HMODULE user32;
    Func_SetThreadDpiAwarenessContext SetThreadDpiAwarenessContext_function;

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

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

    screen = CreateDC("DISPLAY", NULL, NULL, NULL);
    screen_copy = CreateCompatibleDC(screen);

    // added in Windows 10 (1607)
    // loaded dynamically to avoid link errors
    user32 = LoadLibraryA("User32.dll");
    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 (all_screens) {
        x = GetSystemMetrics(SM_XVIRTUALSCREEN);
        y = GetSystemMetrics(SM_YVIRTUALSCREEN);
        width = GetSystemMetrics(SM_CXVIRTUALSCREEN);
        height = GetSystemMetrics(SM_CYVIRTUALSCREEN);
    } else {
        width = GetDeviceCaps(screen, HORZRES);
        height = GetDeviceCaps(screen, VERTRES);
    }

    if (SetThreadDpiAwarenessContext_function != NULL) {
        SetThreadDpiAwarenessContext_function(dpiAwareness);
    }

    FreeLibrary(user32);

    bitmap = CreateCompatibleBitmap(screen, width, height);
    if (!bitmap) {
        goto error;
    }

    if (!SelectObject(screen_copy, bitmap)) {
        goto error;
    }

    /* step 2: copy bits into memory DC bitmap */

    rop = SRCCOPY;
    if (includeLayeredWindows) {
        rop |= CAPTUREBLT;
    }
    if (!BitBlt(screen_copy, 0, 0, width, height, screen, x, y, rop)) {
        goto error;
    }

    /* step 3: extract bits from bitmap */

    buffer = PyBytes_FromStringAndSize(NULL, height * ((width*3 + 3) & -4));
    if (!buffer) {
        return NULL;
    }

    core.bcSize = sizeof(core);
    core.bcWidth = width;
    core.bcHeight = height;
    core.bcPlanes = 1;
    core.bcBitCount = 24;
    if (!GetDIBits(screen_copy, bitmap, 0, height, PyBytes_AS_STRING(buffer),
                   (BITMAPINFO*) &core, DIB_RGB_COLORS)) {
        goto error;
    }

    DeleteObject(bitmap);
    DeleteDC(screen_copy);
    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);

    return NULL;
}

static BOOL CALLBACK list_windows_callback(HWND hwnd, LPARAM lParam)
{
    PyObject* window_list = (PyObject*) lParam;
    PyObject* item;
    PyObject* title;
    RECT inner, outer;
    int title_size;
    int status;

    /* get window title */
    title_size = GetWindowTextLength(hwnd);
    if (title_size > 0) {
        title = PyUnicode_FromStringAndSize(NULL, title_size);
        if (title) {
            GetWindowTextW(hwnd, PyUnicode_AS_UNICODE(title), title_size+1);
        }
    } else {
        title = PyUnicode_FromString("");
    }
    if (!title) {
        return 0;
    }

    /* get bounding boxes */
    GetClientRect(hwnd, &inner);
    GetWindowRect(hwnd, &outer);

    item = Py_BuildValue(
        F_HANDLE "N(iiii)(iiii)", hwnd, title,
        inner.left, inner.top, inner.right, inner.bottom,
        outer.left, outer.top, outer.right, outer.bottom
        );
    if (!item) {
        return 0;
    }

    status = PyList_Append(window_list, item);

    Py_DECREF(item);

    if (status < 0) {
        return 0;
    }

    return 1;
}

PyObject*
PyImaging_ListWindowsWin32(PyObject* self, PyObject* args)
{
    PyObject* window_list;

    window_list = PyList_New(0);
    if (!window_list) {
        return NULL;
    }

    EnumWindows(list_windows_callback, (LPARAM) window_list);

    if (PyErr_Occurred()) {
        Py_DECREF(window_list);
        return NULL;
    }

    return window_list;
}

/* -------------------------------------------------------------------- */
/* Windows clipboard grabber */

PyObject*
PyImaging_GrabClipboardWin32(PyObject* self, PyObject* args)
{
    int clip;
    HANDLE handle = NULL;
    int size;
    void* data;
    PyObject* result;
    UINT format;
    UINT formats[] = { CF_DIB, CF_DIBV5, CF_HDROP, RegisterClipboardFormatA("PNG"), 0 };
    LPCSTR format_names[] = { "DIB", "DIB", "file", "png", NULL };

    if (!OpenClipboard(NULL)) {
        PyErr_SetString(PyExc_OSError, "failed to open clipboard");
        return NULL;
    }

    // find best format as set by clipboard owner
    format = 0;
    while (!handle && (format = EnumClipboardFormats(format))) {
        for (UINT i = 0; formats[i] != 0; i++) {
            if (format == formats[i]) {
                handle = GetClipboardData(format);
                format = i;
                break;
            }
        }
    }

    if (!handle) {
        CloseClipboard();
        return Py_BuildValue("zO", NULL, Py_None);
    }

    data = GlobalLock(handle);
    size = GlobalSize(handle);

    result = PyBytes_FromStringAndSize(data, size);

    GlobalUnlock(handle);
    CloseClipboard();

    return Py_BuildValue("zN", format_names[format], result);
}

/* -------------------------------------------------------------------- */
/* Windows class */

#ifndef WM_MOUSEWHEEL
#define WM_MOUSEWHEEL 522
#endif

static int mainloop = 0;

static void
callback_error(const char* handler)
{
    PyObject* sys_stderr;

    sys_stderr = PySys_GetObject("stderr");

    if (sys_stderr) {
        PyFile_WriteString("*** ImageWin: error in ", sys_stderr);
        PyFile_WriteString((char*) handler, sys_stderr);
        PyFile_WriteString(":\n", sys_stderr);
    }

    PyErr_Print();
    PyErr_Clear();
}

static LRESULT CALLBACK
windowCallback(HWND wnd, UINT message, WPARAM wParam, LPARAM lParam)
{
    PAINTSTRUCT ps;
    PyObject* callback = NULL;
    PyObject* result;
    PyThreadState* threadstate;
    PyThreadState* current_threadstate;
    HDC dc;
    RECT rect;
    LRESULT status = 0;

    /* set up threadstate for messages that calls back into python */
    switch (message) {
    case WM_CREATE:
        mainloop++;
        break;
    case WM_DESTROY:
        mainloop--;
        /* fall through... */
    case WM_PAINT:
    case WM_SIZE:
        callback = (PyObject*) GetWindowLongPtr(wnd, 0);
        if (callback) {
            threadstate = (PyThreadState*)
                GetWindowLongPtr(wnd, sizeof(PyObject*));
            current_threadstate = PyThreadState_Swap(NULL);
            PyEval_RestoreThread(threadstate);
        } else {
            return DefWindowProc(wnd, message, wParam, lParam);
        }
    }

    /* process message */
    switch (message) {

    case WM_PAINT:
        /* redraw (part of) window.  this generates a WCK-style
           damage/clear/repair cascade */
        BeginPaint(wnd, &ps);
        dc = GetDC(wnd);
        GetWindowRect(wnd, &rect); /* in screen coordinates */

        result = PyObject_CallFunction(
            callback, "siiii", "damage",
            ps.rcPaint.left, ps.rcPaint.top,
            ps.rcPaint.right, ps.rcPaint.bottom
            );
        if (result) {
            Py_DECREF(result);
        } else {
            callback_error("window damage callback");
        }

        result = PyObject_CallFunction(
            callback, "s" F_HANDLE "iiii", "clear", dc,
            0, 0, rect.right-rect.left, rect.bottom-rect.top
            );
        if (result) {
            Py_DECREF(result);
        } else {
            callback_error("window clear callback");
        }

        result = PyObject_CallFunction(
            callback, "s" F_HANDLE "iiii", "repair", dc,
            0, 0, rect.right-rect.left, rect.bottom-rect.top
            );
        if (result) {
            Py_DECREF(result);
        } else {
            callback_error("window repair callback");
        }

        ReleaseDC(wnd, dc);
        EndPaint(wnd, &ps);
        break;

    case WM_SIZE:
        /* resize window */
        result = PyObject_CallFunction(
            callback, "sii", "resize", LOWORD(lParam), HIWORD(lParam)
            );
        if (result) {
            InvalidateRect(wnd, NULL, 1);
            Py_DECREF(result);
        } else {
            callback_error("window resize callback");
        }
        break;

    case WM_DESTROY:
        /* destroy window */
        result = PyObject_CallFunction(callback, "s", "destroy");
        if (result) {
            Py_DECREF(result);
        } else {
            callback_error("window destroy callback");
        }
        Py_DECREF(callback);
        break;

    default:
        status = DefWindowProc(wnd, message, wParam, lParam);
    }

    if (callback) {
        /* restore thread state */
        PyEval_SaveThread();
        PyThreadState_Swap(threadstate);
    }

    return status;
}

PyObject*
PyImaging_CreateWindowWin32(PyObject* self, PyObject* args)
{
    HWND wnd;
    WNDCLASS windowClass;

    char* title;
    PyObject* callback;
    int width = 0, height = 0;
    if (!PyArg_ParseTuple(args, "sO|ii", &title, &callback, &width, &height)) {
        return NULL;
    }

    if (width <= 0) {
        width = CW_USEDEFAULT;
    }
    if (height <= 0) {
        height = CW_USEDEFAULT;
    }

    /* register toplevel window class */
    windowClass.style = CS_CLASSDC;
    windowClass.cbClsExtra = 0;
    windowClass.cbWndExtra = sizeof(PyObject*) + sizeof(PyThreadState*);
    windowClass.hInstance = GetModuleHandle(NULL);
    /* windowClass.hbrBackground = (HBRUSH) (COLOR_BTNFACE + 1); */
    windowClass.hbrBackground = NULL;
    windowClass.lpszMenuName = NULL;
    windowClass.lpszClassName = "pilWindow";
    windowClass.lpfnWndProc = windowCallback;
    windowClass.hIcon = LoadIcon(GetModuleHandle(NULL), MAKEINTRESOURCE(1));
    windowClass.hCursor = LoadCursor(NULL, IDC_ARROW); /* CROSS? */

    RegisterClass(&windowClass); /* FIXME: check return status */

    wnd = CreateWindowEx(
        0, windowClass.lpszClassName, title,
        WS_OVERLAPPEDWINDOW,
        CW_USEDEFAULT, CW_USEDEFAULT, width, height,
        HWND_DESKTOP, NULL, NULL, NULL
        );

    if (!wnd) {
        PyErr_SetString(PyExc_OSError, "failed to create window");
        return NULL;
    }

    /* register window callback */
    Py_INCREF(callback);
    SetWindowLongPtr(wnd, 0, (LONG_PTR) callback);
    SetWindowLongPtr(wnd, sizeof(callback), (LONG_PTR) PyThreadState_Get());

    Py_BEGIN_ALLOW_THREADS
    ShowWindow(wnd, SW_SHOWNORMAL);
    SetForegroundWindow(wnd); /* to make sure it's visible */
    Py_END_ALLOW_THREADS

    return Py_BuildValue(F_HANDLE, wnd);
}

PyObject*
PyImaging_EventLoopWin32(PyObject* self, PyObject* args)
{
    MSG msg;

    Py_BEGIN_ALLOW_THREADS
    while (mainloop && GetMessage(&msg, NULL, 0, 0)) {
        TranslateMessage(&msg);
        DispatchMessage(&msg);
    }
    Py_END_ALLOW_THREADS

    Py_INCREF(Py_None);
    return Py_None;
}

/* -------------------------------------------------------------------- */
/* windows WMF renderer */

#define GET32(p,o) ((DWORD*)(p+o))[0]

PyObject *
PyImaging_DrawWmf(PyObject* self, PyObject* args)
{
    HBITMAP bitmap;
    HENHMETAFILE meta;
    BITMAPCOREHEADER core;
    HDC dc;
    RECT rect;
    PyObject* buffer = NULL;
    char* ptr;

    char* data;
    Py_ssize_t datasize;
    int width, height;
    int x0, y0, x1, y1;
    if (!PyArg_ParseTuple(args, "y#(ii)(iiii):_load", &data, &datasize,
                          &width, &height, &x0, &x1, &y0, &y1)) {
        return NULL;
    }

    /* step 1: copy metafile contents into METAFILE object */

    if (datasize > 22 && GET32(data, 0) == 0x9ac6cdd7) {

        /* placeable windows metafile (22-byte aldus header) */
        meta = SetWinMetaFileBits(datasize-22, data+22, NULL, NULL);

    } else if (datasize > 80 && GET32(data, 0) == 1 &&
               GET32(data, 40) == 0x464d4520) {

        /* enhanced metafile */
        meta = SetEnhMetaFileBits(datasize, data);

    } else {

        /* unknown meta format */
        meta = NULL;

    }

    if (!meta) {
        PyErr_SetString(PyExc_OSError, "cannot load metafile");
        return NULL;
    }

    /* step 2: create bitmap */

    core.bcSize = sizeof(core);
    core.bcWidth = width;
    core.bcHeight = height;
    core.bcPlanes = 1;
    core.bcBitCount = 24;

    dc = CreateCompatibleDC(NULL);

    bitmap = CreateDIBSection(
        dc, (BITMAPINFO*) &core, DIB_RGB_COLORS, &ptr, NULL, 0
        );

    if (!bitmap) {
        PyErr_SetString(PyExc_OSError, "cannot create bitmap");
        goto error;
    }

    if (!SelectObject(dc, bitmap)) {
        PyErr_SetString(PyExc_OSError, "cannot select bitmap");
        goto error;
    }

    /* step 3: render metafile into bitmap */

    rect.left = rect.top = 0;
    rect.right = width;
    rect.bottom = height;

    /* FIXME: make background transparent? configurable? */
    FillRect(dc, &rect, GetStockObject(WHITE_BRUSH));

    if (!PlayEnhMetaFile(dc, meta, &rect)) {
        PyErr_SetString(PyExc_OSError, "cannot render metafile");
        goto error;
    }

    /* step 4: extract bits from bitmap */

    GdiFlush();

    buffer = PyBytes_FromStringAndSize(ptr, height * ((width*3 + 3) & -4));

error:
    DeleteEnhMetaFile(meta);

    if (bitmap) {
        DeleteObject(bitmap);
    }

    DeleteDC(dc);

    return buffer;
}

#endif /* _WIN32 */

/* -------------------------------------------------------------------- */
/* X11 support */

#ifdef HAVE_XCB
#include <xcb/xcb.h>

/* -------------------------------------------------------------------- */
/* X11 screen grabber */

PyObject*
PyImaging_GrabScreenX11(PyObject* self, PyObject* args)
{
    int width, height;
    char* display_name;
    xcb_connection_t* connection;
    int screen_number;
    xcb_screen_iterator_t iter;
    xcb_screen_t* screen = NULL;
    xcb_get_image_reply_t* reply;
    xcb_generic_error_t* error;
    PyObject* buffer = NULL;

    if (!PyArg_ParseTuple(args, "|z", &display_name)) {
        return NULL;
    }

    /* connect to X and get screen data */

    connection = xcb_connect(display_name, &screen_number);
    if (xcb_connection_has_error(connection)) {
        PyErr_Format(PyExc_OSError, "X connection failed: error %i", xcb_connection_has_error(connection));
        xcb_disconnect(connection);
        return NULL;
    }

    iter = xcb_setup_roots_iterator(xcb_get_setup(connection));
    for (; iter.rem; --screen_number, xcb_screen_next(&iter)) {
        if (screen_number == 0) {
            screen = iter.data;
            break;
        }
    }
    if (screen == NULL || screen->root == 0) {
        // this case is usually caught with "X connection failed: error 6" above
        xcb_disconnect(connection);
        PyErr_SetString(PyExc_OSError, "X screen not found");
        return NULL;
    }

    width = screen->width_in_pixels;
    height = screen->height_in_pixels;

    /* get image data */

    reply = xcb_get_image_reply(connection,
                                xcb_get_image(connection, XCB_IMAGE_FORMAT_Z_PIXMAP, screen->root,
                                              0, 0, width, height, 0x00ffffff),
                                &error);
    if (reply == NULL) {
        PyErr_Format(PyExc_OSError, "X get_image failed: error %i (%i, %i, %i)",
                     error->error_code, error->major_code, error->minor_code, error->resource_id);
        free(error);
        xcb_disconnect(connection);
        return NULL;
    }

    /* store data in Python buffer */

    if (reply->depth == 24) {
        buffer = PyBytes_FromStringAndSize((char*)xcb_get_image_data(reply),
                                           xcb_get_image_data_length(reply));
    } else {
        PyErr_Format(PyExc_OSError, "unsupported bit depth: %i", reply->depth);
    }

    free(reply);
    xcb_disconnect(connection);

    if (!buffer) {
        return NULL;
    }

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

#endif /* HAVE_XCB */