Skip to content

Commit dc22c73

Browse files
author
Mickaël S
committed
- Mac: remove custom to_png(), use parent.to_png()
- Mac: get_pixels() returns the good type - Define self.width and self.height
1 parent 4ffd271 commit dc22c73

File tree

4 files changed

+46
-47
lines changed

4 files changed

+46
-47
lines changed

‎mss/base.py‎

Lines changed: 10 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,8 @@ class MSSBase(object):
1616

1717
monitors = []
1818
image = None
19+
width = 0
20+
height = 0
1921

2022
def __enter__(self):
2123
''' For the cool call `with MSS() as mss:`. '''
@@ -49,6 +51,8 @@ def enum_display_monitors(self, force=False):
4951
def get_pixels(self, monitor):
5052
''' Retrieve screen pixels for a given monitor.
5153
54+
This method has to define self.width and self.height.
55+
5256
`monitor` is a dict with:
5357
{
5458
'left': the x-coordinate of the upper-left corner,
@@ -90,10 +94,7 @@ def save(self, mon=0, output='monitor-%d.png', callback=lambda *x: True):
9094
if '%d' in output:
9195
fname = output.replace('%d', str(i + 1))
9296
callback(fname)
93-
self.to_png(data=self.get_pixels(monitor),
94-
width=monitor[b'width'],
95-
height=monitor[b'height'],
96-
output=fname)
97+
self.to_png(self.get_pixels(monitor), fname)
9798
yield fname
9899
else:
99100
# A screenshot of all monitors together or
@@ -108,13 +109,10 @@ def save(self, mon=0, output='monitor-%d.png', callback=lambda *x: True):
108109
if '%d' in output:
109110
output = output.replace('%d', str(mon_number))
110111
callback(output)
111-
self.to_png(data=self.get_pixels(monitor),
112-
width=monitor[b'width'],
113-
height=monitor[b'height'],
114-
output=output)
112+
self.to_png(self.get_pixels(monitor), output)
115113
yield output
116114

117-
def to_png(self, data, width, height, output):
115+
def to_png(self, data, output):
118116
''' Dump data to the image file.
119117
Pure python PNG implementation.
120118
http://inaps.org/journal/comment-fonctionne-le-png
@@ -123,17 +121,17 @@ def to_png(self, data, width, height, output):
123121
# pylint: disable=no-self-use
124122

125123
p__ = pack
126-
line = width * 3
124+
line = self.width * 3
127125
png_filter = p__(b'>B', 0)
128126
scanlines = b''.join(
129127
[png_filter + data[y * line:y * line + line]
130-
for y in range(height)])
128+
for y in range(self.height)])
131129

132130
magic = p__(b'>8B', 137, 80, 78, 71, 13, 10, 26, 10)
133131

134132
# Header: size, marker, data, CRC32
135133
ihdr = [b'', b'IHDR', b'', b'']
136-
ihdr[2] = p__(b'>2I5B', width, height, 8, 2, 0, 0, 0)
134+
ihdr[2] = p__(b'>2I5B', self.width, self.height, 8, 2, 0, 0, 0)
137135
ihdr[3] = p__(b'>I', crc32(b''.join(ihdr[1:3])) & 0xffffffff)
138136
ihdr[0] = p__(b'>I', len(ihdr[2]))
139137

‎mss/darwin.py‎

Lines changed: 16 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,8 @@
1111
CGImageDestinationAddImage, CGImageDestinationCreateWithURL,
1212
CGImageDestinationFinalize, CGRect, CGRectInfinite, CGRectStandardize,
1313
CGWindowListCreateImage, kCGNullWindowID, kCGWindowImageDefault,
14-
kCGWindowListOptionOnScreenOnly)
14+
kCGWindowListOptionOnScreenOnly, CGDataProviderCopyData,
15+
CGImageGetDataProvider, CGImageGetWidth, CGImageGetHeight)
1516

1617
from .base import MSSBase
1718
from .exception import ScreenshotError
@@ -58,30 +59,27 @@ def enum_display_monitors(self, force=False):
5859
return self.monitors
5960

6061
def get_pixels(self, monitor):
61-
''' Retrieve all pixels from a monitor. Pixels have to be RGB.
62-
'''
62+
''' Retrieve all pixels from a monitor. Pixels have to be RGB. '''
6363

6464
width, height = monitor[b'width'], monitor[b'height']
6565
left, top = monitor[b'left'], monitor[b'top']
6666
rect = CGRect((left, top), (width, height))
6767
options = kCGWindowListOptionOnScreenOnly
6868
winid = kCGNullWindowID
6969
default = kCGWindowImageDefault
70-
self.image = CGWindowListCreateImage(rect, options, winid, default)
71-
if not self.image:
72-
raise ScreenshotError('CGWindowListCreateImage() failed.')
73-
return self.image
7470

75-
def to_png(self, data, width, height, output):
76-
''' Use of internal tools, faster and less code to write :) '''
71+
image_ref = CGWindowListCreateImage(rect, options, winid, default)
72+
if not image_ref:
73+
raise ScreenshotError('CGWindowListCreateImage() failed.')
7774

78-
url = NSURL.fileURLWithPath_(output)
79-
dest = CGImageDestinationCreateWithURL(url, 'public.png', 1, None)
80-
if not dest:
81-
err = 'CGImageDestinationCreateWithURL() failed.'
82-
raise ScreenshotError(err)
75+
self.width = CGImageGetWidth(image_ref)
76+
self.height = CGImageGetHeight(image_ref)
77+
image_data = CGDataProviderCopyData(CGImageGetDataProvider(image_ref))
8378

84-
CGImageDestinationAddImage(dest, data, None)
85-
if CGImageDestinationFinalize(dest):
86-
return True
87-
raise ScreenshotError('CGImageDestinationFinalize() failed.')
79+
# Replace pixels values: BGRA to RGB.
80+
image_data = bytearray(image_data)
81+
image = bytearray(self.height * self.width * 3)
82+
image[0::3], image[1::3], image[2::3] = \
83+
image_data[2::4], image_data[1::4], image_data[0::4]
84+
self.image = bytes(image)
85+
return self.image

‎mss/linux.py‎

Lines changed: 9 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -244,23 +244,25 @@ def enum_display_monitors(self, force=False):
244244
def get_pixels(self, monitor):
245245
''' Retrieve all pixels from a monitor. Pixels have to be RGB. '''
246246

247-
width, height = monitor[b'width'], monitor[b'height']
247+
self.width = monitor[b'width']
248+
self.height = monitor[b'height']
248249
left, top = monitor[b'left'], monitor[b'top']
249250
zpixmap = 2
250251
allplanes = self.xlib.XAllPlanes()
251252

252253
# Fix for XGetImage:
253254
# expected LP_Display instance instead of LP_XWindowAttributes
254255
root = cast(self.root, POINTER(Display))
255-
ximage = self.xlib.XGetImage(self.display, root, left, top, width,
256-
height, allplanes, zpixmap)
256+
ximage = self.xlib.XGetImage(self.display, root, left, top,
257+
self.width, self.height, allplanes,
258+
zpixmap)
257259
if not ximage:
258260
raise ScreenshotError('xlib.XGetImage() failed.')
259261

260262
if not self.use_mss:
261263
self.get_pixels_slow(ximage)
262264
else:
263-
self.image = create_string_buffer(height * width * 3)
265+
self.image = create_string_buffer(self.height * self.width * 3)
264266
ret = self.mss.GetXImagePixels(ximage, self.image)
265267
if not ret:
266268
self.xlib.XDestroyImage(ximage)
@@ -288,14 +290,14 @@ def pix(pixel, _resultats={}, p__=pack):
288290
p__(b'<B', pixel & bmask)
289291
return _resultats[pixel]
290292

291-
width = ximage.contents.width
292-
height = ximage.contents.height
293+
self.width = ximage.contents.width
294+
self.height = ximage.contents.height
293295
rmask = ximage.contents.red_mask
294296
bmask = ximage.contents.blue_mask
295297
gmask = ximage.contents.green_mask
296298
get_pix = self.xlib.XGetPixel
297299
pixels = [pix(get_pix(ximage, x, y))
298-
for y in range(height) for x in range(width)]
300+
for y in range(self.height) for x in range(self.width)]
299301
self.image = b''.join(pixels)
300302
return self.image
301303

‎mss/windows.py‎

Lines changed: 11 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -117,7 +117,8 @@ def get_pixels(self, monitor):
117117
Thanks to http://stackoverflow.com/a/3688682
118118
'''
119119

120-
width, height = monitor[b'width'], monitor[b'height']
120+
self.width = monitor[b'width']
121+
self.height = monitor[b'height']
121122
left, top = monitor[b'left'], monitor[b'top']
122123
srcdc = None
123124
memdc = None
@@ -126,24 +127,24 @@ def get_pixels(self, monitor):
126127
try:
127128
bmi = BITMAPINFO()
128129
bmi.bmiHeader.biSize = sizeof(BITMAPINFOHEADER)
129-
bmi.bmiHeader.biWidth = width
130-
bmi.bmiHeader.biHeight = -height # Why minux? See [1]
130+
bmi.bmiHeader.biWidth = self.width
131+
bmi.bmiHeader.biHeight = -self.height # Why minux? See [1]
131132
bmi.bmiHeader.biPlanes = 1 # Always 1
132133
bmi.bmiHeader.biBitCount = 32 # See [2]
133134
bmi.bmiHeader.biCompression = 0 # 0 = BI_RGB (no compression)
134135
bmi.bmiHeader.biClrUsed = 0 # See [3]
135136
bmi.bmiHeader.biClrImportant = 0 # See [3]
136137

137-
image_data = create_string_buffer(height * width * 4) # See [2]
138+
image_data = create_string_buffer(self.height * self.width * 4) # See [2]
138139
srcdc = windll.user32.GetWindowDC(0)
139140
memdc = windll.gdi32.CreateCompatibleDC(srcdc)
140141
bmp = windll.gdi32.CreateCompatibleBitmap(srcdc, width, height)
141142
windll.gdi32.SelectObject(memdc, bmp)
142-
windll.gdi32.BitBlt(memdc, 0, 0, width, height, srcdc, left, top,
143-
0xCC0020) # 0xCC0020 = SRCCOPY
144-
bits = windll.gdi32.GetDIBits(memdc, bmp, 0, height, image_data,
145-
bmi, 0) # 0 = DIB_RGB_COLORS
146-
if bits != height:
143+
windll.gdi32.BitBlt(memdc, 0, 0, self.width, self.height, srcdc,
144+
left, top, 0xCC0020) # 0xCC0020 = SRCCOPY
145+
bits = windll.gdi32.GetDIBits(memdc, bmp, 0, self.height,
146+
image_data, bmi, 0) # 0 = DIB_RGB_COLORS
147+
if bits != self.height:
147148
raise ScreenshotError('gdi32.GetDIBits() failed.')
148149
finally:
149150
# Clean up
@@ -155,7 +156,7 @@ def get_pixels(self, monitor):
155156
windll.gdi32.DeleteObject(bmp)
156157

157158
# Replace pixels values: BGRX to RGB. See [2].
158-
image = bytearray(height * width * 3)
159+
image = bytearray(self.height * self.width * 3)
159160
image[0::3], image[1::3], image[2::3] = \
160161
image_data[2::4], image_data[1::4], image_data[0::4]
161162
self.image = bytes(image)

0 commit comments

Comments
 (0)