407 lines
9.3 KiB
Python
407 lines
9.3 KiB
Python
"""
|
|
A wrapper around the Direct Rendering Manager (DRM) library, which itself is a
|
|
wrapper around the Direct Rendering Interface (DRI) between the kernel and
|
|
userland.
|
|
|
|
Since we are masochists, we use ctypes instead of cffi to load libdrm and
|
|
access several symbols within it. We use Python's file descriptor and mmap
|
|
wrappers.
|
|
|
|
At some point in the future, cffi could be used, for approximately the same
|
|
cost in lines of code.
|
|
"""
|
|
|
|
from ctypes import *
|
|
import mmap
|
|
import os
|
|
|
|
from PIL import Image
|
|
|
|
|
|
class DrmVersion(Structure):
|
|
"""
|
|
The version of a DRM node.
|
|
"""
|
|
|
|
_fields_ = [
|
|
("version_major", c_int),
|
|
("version_minor", c_int),
|
|
("version_patchlevel", c_int),
|
|
("name_len", c_int),
|
|
("name", c_char_p),
|
|
("date_len", c_int),
|
|
("date", c_char_p),
|
|
("desc_len", c_int),
|
|
("desc", c_char_p),
|
|
]
|
|
|
|
_l = None
|
|
|
|
def __repr__(self):
|
|
return "%s %d.%d.%d (%s) (%s)" % (
|
|
self.name,
|
|
self.version_major,
|
|
self.version_minor,
|
|
self.version_patchlevel,
|
|
self.desc,
|
|
self.date,
|
|
)
|
|
|
|
def __del__(self):
|
|
if self._l:
|
|
self._l.drmFreeVersion(self)
|
|
|
|
|
|
class DrmModeResources(Structure):
|
|
"""
|
|
Resources associated with setting modes on a DRM node.
|
|
"""
|
|
|
|
_fields_ = [
|
|
("count_fbs", c_int),
|
|
("fbs", POINTER(c_uint)),
|
|
("count_crtcs", c_int),
|
|
("crtcs", POINTER(c_uint)),
|
|
("count_connectors", c_int),
|
|
("connectors", POINTER(c_uint)),
|
|
("count_encoders", c_int),
|
|
("encoders", POINTER(c_uint)),
|
|
("min_width", c_int),
|
|
("max_width", c_int),
|
|
("min_height", c_int),
|
|
("max_height", c_int),
|
|
]
|
|
|
|
_fd = None
|
|
_l = None
|
|
|
|
def __repr__(self):
|
|
return "<DRM mode resources>"
|
|
|
|
def __del__(self):
|
|
if self._l:
|
|
self._l.drmModeFreeResources(self)
|
|
|
|
def getValidCrtc(self):
|
|
for i in xrange(0, self.count_crtcs):
|
|
crtc_id = self.crtcs[i]
|
|
crtc = self._l.drmModeGetCrtc(self._fd, crtc_id).contents
|
|
if crtc.mode_valid:
|
|
return crtc
|
|
return None
|
|
|
|
def getCrtc(self, crtc_id=None):
|
|
"""
|
|
Obtain the CRTC at a given index.
|
|
|
|
@param crtc_id: The CRTC to get.
|
|
"""
|
|
crtc = None
|
|
if crtc_id:
|
|
crtc = self._l.drmModeGetCrtc(self._fd, crtc_id).contents
|
|
else:
|
|
crtc = self.getValidCrtc()
|
|
crtc._fd = self._fd
|
|
crtc._l = self._l
|
|
return crtc
|
|
|
|
|
|
class DrmModeCrtc(Structure):
|
|
"""
|
|
A DRM modesetting CRTC.
|
|
"""
|
|
|
|
_fields_ = [
|
|
("crtc_id", c_uint),
|
|
("buffer_id", c_uint),
|
|
("x", c_uint),
|
|
("y", c_uint),
|
|
("width", c_uint),
|
|
("height", c_uint),
|
|
("mode_valid", c_int),
|
|
# XXX incomplete struct!
|
|
]
|
|
|
|
_fd = None
|
|
_l = None
|
|
|
|
def __repr__(self):
|
|
return "<CRTC (%d)>" % self.crtc_id
|
|
|
|
def __del__(self):
|
|
if self._l:
|
|
self._l.drmModeFreeCrtc(self)
|
|
|
|
def hasFb(self):
|
|
"""
|
|
Whether this CRTC has an associated framebuffer.
|
|
"""
|
|
|
|
return self.buffer_id != 0
|
|
|
|
def fb(self):
|
|
"""
|
|
Obtain the framebuffer, if one is associated.
|
|
"""
|
|
|
|
if self.hasFb():
|
|
fb = self._l.drmModeGetFB(self._fd, self.buffer_id).contents
|
|
fb._fd = self._fd
|
|
fb._l = self._l
|
|
return fb
|
|
else:
|
|
raise RuntimeError("CRTC %d doesn't have a framebuffer!" %
|
|
self.crtc_id)
|
|
|
|
|
|
class drm_mode_map_dumb(Structure):
|
|
"""
|
|
Request a mapping of a modesetting buffer.
|
|
|
|
The map will be "dumb;" it will be accessible via mmap() but very slow.
|
|
"""
|
|
|
|
_fields_ = [
|
|
("handle", c_uint),
|
|
("pad", c_uint),
|
|
("offset", c_ulonglong),
|
|
]
|
|
|
|
|
|
# This constant is not defined in any one header; it is the pieced-together
|
|
# incantation for the ioctl that performs dumb mappings. I would love for this
|
|
# to not have to be here, but it can't be imported from any header easily.
|
|
DRM_IOCTL_MODE_MAP_DUMB = 0xc01064b3
|
|
|
|
|
|
class DrmModeFB(Structure):
|
|
"""
|
|
A DRM modesetting framebuffer.
|
|
"""
|
|
|
|
_fields_ = [
|
|
("fb_id", c_uint),
|
|
("width", c_uint),
|
|
("height", c_uint),
|
|
("pitch", c_uint),
|
|
("bpp", c_uint),
|
|
("depth", c_uint),
|
|
("handle", c_uint),
|
|
]
|
|
|
|
_l = None
|
|
_map = None
|
|
|
|
def __repr__(self):
|
|
s = "<Framebuffer (%dx%d (pitch %d bytes), %d bits/pixel, depth %d)"
|
|
vitals = s % (
|
|
self.width,
|
|
self.height,
|
|
self.pitch,
|
|
self.bpp,
|
|
self.depth,
|
|
)
|
|
if self._map:
|
|
tail = " (mapped)>"
|
|
else:
|
|
tail = ">"
|
|
return vitals + tail
|
|
|
|
def __del__(self):
|
|
if self._l:
|
|
self._l.drmModeFreeFB(self)
|
|
|
|
def map(self):
|
|
"""
|
|
Map the framebuffer.
|
|
"""
|
|
|
|
if self._map:
|
|
return
|
|
|
|
mapDumb = drm_mode_map_dumb()
|
|
mapDumb.handle = self.handle
|
|
|
|
rv = self._l.drmIoctl(self._fd, DRM_IOCTL_MODE_MAP_DUMB,
|
|
pointer(mapDumb))
|
|
if rv:
|
|
raise IOError(rv, os.strerror(rv))
|
|
|
|
size = self.pitch * self.height
|
|
|
|
# mmap.mmap() has a totally different order of arguments in Python
|
|
# compared to C; check the documentation before altering this
|
|
# incantation.
|
|
self._map = mmap.mmap(self._fd, size, flags=mmap.MAP_SHARED,
|
|
prot=mmap.PROT_READ, offset=mapDumb.offset)
|
|
|
|
def unmap(self):
|
|
"""
|
|
Unmap the framebuffer.
|
|
"""
|
|
|
|
if self._map:
|
|
self._map.close()
|
|
self._map = None
|
|
|
|
|
|
def loadDRM():
|
|
"""
|
|
Load a handle to libdrm.
|
|
|
|
In addition to loading, this function also configures the argument and
|
|
return types of functions.
|
|
"""
|
|
|
|
l = cdll.LoadLibrary("libdrm.so")
|
|
|
|
l.drmGetVersion.argtypes = [c_int]
|
|
l.drmGetVersion.restype = POINTER(DrmVersion)
|
|
|
|
l.drmFreeVersion.argtypes = [POINTER(DrmVersion)]
|
|
l.drmFreeVersion.restype = None
|
|
|
|
l.drmModeGetResources.argtypes = [c_int]
|
|
l.drmModeGetResources.restype = POINTER(DrmModeResources)
|
|
|
|
l.drmModeFreeResources.argtypes = [POINTER(DrmModeResources)]
|
|
l.drmModeFreeResources.restype = None
|
|
|
|
l.drmModeGetCrtc.argtypes = [c_int, c_uint]
|
|
l.drmModeGetCrtc.restype = POINTER(DrmModeCrtc)
|
|
|
|
l.drmModeFreeCrtc.argtypes = [POINTER(DrmModeCrtc)]
|
|
l.drmModeFreeCrtc.restype = None
|
|
|
|
l.drmModeGetFB.argtypes = [c_int, c_uint]
|
|
l.drmModeGetFB.restype = POINTER(DrmModeFB)
|
|
|
|
l.drmModeFreeFB.argtypes = [POINTER(DrmModeFB)]
|
|
l.drmModeFreeFB.restype = None
|
|
|
|
l.drmIoctl.argtypes = [c_int, c_ulong, c_voidp]
|
|
l.drmIoctl.restype = c_int
|
|
|
|
return l
|
|
|
|
|
|
class DRM(object):
|
|
"""
|
|
A DRM node.
|
|
"""
|
|
|
|
def __init__(self, library, fd):
|
|
self._l = library
|
|
self._fd = fd
|
|
|
|
def __repr__(self):
|
|
return "<DRM (FD %d)>" % self._fd
|
|
|
|
@classmethod
|
|
def fromHandle(cls, handle):
|
|
"""
|
|
Create a node from a file handle.
|
|
|
|
@param handle: A file-like object backed by a file descriptor.
|
|
"""
|
|
|
|
self = cls(loadDRM(), handle.fileno())
|
|
# We must keep the handle alive, and we cannot trust the caller to
|
|
# keep it alive for us.
|
|
self._handle = handle
|
|
return self
|
|
|
|
def version(self):
|
|
"""
|
|
Obtain the version.
|
|
"""
|
|
|
|
v = self._l.drmGetVersion(self._fd).contents
|
|
v._l = self._l
|
|
return v
|
|
|
|
def resources(self):
|
|
"""
|
|
Obtain the modesetting resources.
|
|
"""
|
|
|
|
resources_ptr = self._l.drmModeGetResources(self._fd)
|
|
if resources_ptr:
|
|
r = resources_ptr.contents
|
|
r._fd = self._fd
|
|
r._l = self._l
|
|
return r
|
|
|
|
return None
|
|
|
|
|
|
def drmFromPath(path):
|
|
"""
|
|
Given a DRM node path, open the corresponding node.
|
|
|
|
@param path: The path of the minor node to open.
|
|
"""
|
|
|
|
handle = open(path)
|
|
return DRM.fromHandle(handle)
|
|
|
|
|
|
def _bgrx24(i):
|
|
b = ord(next(i))
|
|
g = ord(next(i))
|
|
r = ord(next(i))
|
|
next(i)
|
|
return r, g, b
|
|
|
|
|
|
def _screenshot(image, fb):
|
|
fb.map()
|
|
m = fb._map
|
|
lineLength = fb.width * fb.bpp // 8
|
|
pitch = fb.pitch
|
|
pixels = []
|
|
|
|
if fb.depth == 24:
|
|
unformat = _bgrx24
|
|
else:
|
|
raise RuntimeError("Couldn't unformat FB: %r" % fb)
|
|
|
|
for y in range(fb.height):
|
|
offset = y * pitch
|
|
m.seek(offset)
|
|
channels = m.read(lineLength)
|
|
ichannels = iter(channels)
|
|
for x in range(fb.width):
|
|
rgb = unformat(ichannels)
|
|
image.putpixel((x, y), rgb)
|
|
|
|
fb.unmap()
|
|
|
|
return pixels
|
|
|
|
|
|
_drm = None
|
|
|
|
def crtcScreenshot(crtc_id=None):
|
|
"""
|
|
Take a screenshot, returning an image object.
|
|
|
|
@param crtc_id: The CRTC to screenshot.
|
|
"""
|
|
|
|
global _drm
|
|
if not _drm:
|
|
paths = ["/dev/dri/" + n for n in os.listdir("/dev/dri")]
|
|
for p in paths:
|
|
d = drmFromPath(p)
|
|
if d.resources():
|
|
_drm = d
|
|
break
|
|
|
|
if _drm:
|
|
fb = _drm.resources().getCrtc(crtc_id).fb()
|
|
image = Image.new("RGB", (fb.width, fb.height))
|
|
pixels = _screenshot(image, fb)
|
|
return image
|
|
|
|
raise RuntimeError("Couldn't screenshot with DRM devices")
|