wclip.py

from collections import OrderedDict
from ctypes import sizeof
from ctypes.wintypes import (
    DWORD,
    LONG,
    WORD,
    )

import pygame
from win32clipboard import (
    OpenClipboard, 
    EnumClipboardFormats,
    GetClipboardData,
    CloseClipboard,
    CF_DIB,
    )

class _clipboard:
    def __enter__(self):
        OpenClipboard()

    def __exit__(self, type, value, traceback):
        CloseClipboard()
clipboard = _clipboard()

class Bounds:
    def __init__(self, start, end):
        self.start = start
        self.end = end

    def __len__(self):
        return self.end - self.start + 1

class Header:
    fields = OrderedDict((
        ("biSize", DWORD),
        ("biWidth", LONG),
        ("biHeight", LONG),
        ("biPlanes", WORD),
        ("biBitCount", WORD),
        ("biCompression", DWORD),
        ("biSizeImage", DWORD),
        ("biXPelsPerMeter", LONG),
        ("biYPelsPerMeter", LONG),
        ("biClrUsed", DWORD),
        ("biClrImportant", DWORD),
        ("bitmask_r", DWORD),
        ("bitmask_g", DWORD),
        ("bitmask_b", DWORD),
        ))

    def __init__(self, **kwargs):
        self.biSize = None
        self.biWidth = None
        self.biHeight = None
        self.biPlanes = None
        self.biBitCount = None
        self.biCompression = None
        self.biSizeImage = None
        self.biXPelsPerMeter = None
        self.biYPelsPerMeter = None
        self.biClrUsed = None
        self.biClrImportant = None
        self.bitmask_r = None
        self.bitmask_g = None
        self.bitmask_b = None
        self.__dict__.update(kwargs)

def _get_available_clipboard_formats():
    """requires `with clipboard`"""
    formats = []
    format = EnumClipboardFormats()
    while format != 0:
        formats.append(format)
        format = EnumClipboardFormats(format)
    return formats

def _get_data():
    """requres `with clipboard`"""
    if CF_DIB not in _get_available_clipboard_formats():
        raise TypeError("Clipboard contents are not a screenshot")

    data = GetClipboardData(CF_DIB)

    header = {}
    offset = 0
    for field, datatype in Header.fields.items():
        bounds = _type_offset_to_bounds(datatype, offset)
        header[field] = _byteslice_to_int(data, bounds)
        offset += len(bounds)
        if field == "biClrImportant" and header["biCompression"] != 3:
            #Not a bitfield image, those 12 bytes are for nerds
            break

    return Header(**header), data[offset:]

def _byteslice_to_int(bytestring, bounds):
    bytes = (bytestring[i] for i in range(bounds.start, bounds.end + 1))
    total = 0
    for i, byte in enumerate(bytes):
        total += byte << (8 * i)

    return total

def _type_offset_to_bounds(datatype, offset):
    size = sizeof(datatype)
    end = (size - 1) + offset
    return Bounds(offset, end)

def get_clipboard_screenshot():
    with clipboard:
        header, data = _get_data()

    w = header.biWidth
    h = header.biHeight
    #Data is encoded as BGRX, which isn't something Pygame can handle, so we
    # reverse it to make it XRGB, then take the first byte off and stick it on
    # the end to make RGBX.
    #Bizarrely, some of the X bytes aren't 0xff; the margin down the side of
    # the line numbers in VS, and the description box in the properties window,
    # for instance, are 0x0.
    im = pygame.image.fromstring(data[::-1][1:] + b"\xff", (w, h), "RGBX")
    #Though now the image is the other wrong way up, so give it a flip.
    im = pygame.transform.flip(im, True, False)

    return im