daleathan · March 27, 2016 20:24 · Mar 27, 2016
diff --git a/wclip.py b/wclip.py
@@ -0,0 +1,126 @@
+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