Skip to content

Commit 810f5e9

Browse files
author
BoboTiG
committed
full functional on Mac OSX and code review
1 parent ab3fc19 commit 810f5e9

2 files changed

Lines changed: 107 additions & 97 deletions

File tree

‎README.rst‎

Lines changed: 28 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -78,31 +78,34 @@ Then, it is quite simple::
7878

7979
mss = mss_class()
8080

81-
# One screen shot per monitor
82-
for filename in mss.save():
83-
print('File: "{}" created.'.format(filename))
84-
85-
# Screen shot of the monitor 1
86-
for filename in mss.save(output='monitor-1', screen=1):
87-
print('File: "{}" created.'.format(filename))
88-
89-
# A shot to grab them all :)
90-
for filename in mss.save(screen=-1):
91-
print('File: "{}" created.'.format(filename))
92-
93-
# Example with a callback
94-
def on_exists(fname):
95-
''' Callback example when we try to overwrite an existing screen shot. '''
96-
97-
from os import rename
98-
newfile = fname + '.old'
99-
print('Renaming "{}" to "{}"'.format(fname, newfile))
100-
rename(fname, newfile)
101-
return True
102-
103-
# Screen shot of the monitor 1, with callback
104-
for filename in mss.save(output='monitor-1', screen=1, callback=on_exists):
105-
print('File: "{}" created.'.format(filename))
81+
try:
82+
# One screen shot per monitor
83+
for filename in mss.save():
84+
print('File: "{}" created.'.format(filename))
85+
86+
# Screen shot of the monitor 1
87+
for filename in mss.save(output='monitor-%d.png', screen=1):
88+
print('File: "{}" created.'.format(filename))
89+
90+
# A shot to grab them all :)
91+
for filename in mss.save(output='full-screenshot.png', screen=-1):
92+
print('File: "{}" created.'.format(filename))
93+
94+
# Example with a callback
95+
def on_exists(fname):
96+
''' Callback example when we try to overwrite an existing screen shot. '''
97+
98+
from os import rename
99+
newfile = fname + '.old'
100+
print('Renaming "{}" to "{}"'.format(fname, newfile))
101+
rename(fname, newfile)
102+
return True
103+
104+
# Screen shot of the monitor 1, with callback
105+
for fname in mss.save(output='mon-%d.png', screen=1, callback=on_exists):
106+
print('File: "{}" created.'.format(fname))
107+
except ScreenshotError as ex:
108+
print(ex)
106109

107110

108111
Bonus

‎mss.py‎

Lines changed: 79 additions & 72 deletions
Original file line numberDiff line numberDiff line change
@@ -31,8 +31,8 @@
3131

3232

3333
from struct import pack
34-
from os.path import isfile
3534
from platform import system
35+
import os.path
3636
import sys
3737
import zlib
3838

@@ -229,24 +229,22 @@ def save(self,
229229
self.debug('save', 'output', output)
230230

231231
# Monitors screen shots!
232+
dir_ = os.path.dirname(os.path.abspath(__file__))
232233
for i, monitor in enumerate(self.enum_display_monitors(screen)):
233234
self.debug('save', 'monitor', monitor)
234-
if screen == -1:
235-
filename = '{}-full.png'.format(output)
236-
elif screen == 0:
237-
filename = '{}-{}.png'.format(output, i+1)
238-
else:
239-
filename = '{}.png'.format(output)
240-
241235
if screen <= 0 or (screen > 0 and i+1 == screen):
242-
if isfile(filename) and not callback(filename):
236+
fname = output
237+
if '%d' in output:
238+
fname = output.replace('%d', str(i+1))
239+
fname = os.path.join(dir_, os.path.basename(fname))
240+
self.debug('save', 'fname', fname)
241+
if os.path.isfile(fname) and not callback(fname):
243242
continue
244-
self.debug('save', 'filename', filename)
245243
self.save_img(data=self.get_pixels(monitor),
246244
width=monitor[b'width'],
247245
height=monitor[b'height'],
248-
output=filename)
249-
yield filename
246+
output=fname)
247+
yield fname
250248

251249
def save_img(self, data, width, height, output):
252250
''' Dump data to the image file.
@@ -257,9 +255,6 @@ def save_img(self, data, width, height, output):
257255

258256
self.debug('save_img')
259257

260-
if output[-4:] != '.png':
261-
output += '.png'
262-
263258
to_take = (width * 3 + 3) & -4
264259
padding = 0 if to_take % 8 == 0 else (to_take % 8) // 2
265260
scanlines = b''.join(
@@ -290,6 +285,9 @@ def save_img(self, data, width, height, output):
290285
fileh.write(
291286
magic + b''.join(ihdr) + b''.join(idat) + b''.join(iend)
292287
)
288+
if not os.path.isfile(output):
289+
msg = 'Impossible to write data to file "{}".'.format(output)
290+
raise ScreenshotError(msg)
293291

294292

295293
class MSSMac(MSS):
@@ -301,7 +299,7 @@ def init(self):
301299
''' Mac OSX initialisations '''
302300
self.debug('init')
303301

304-
def enum_display_monitors(self, screen=-1):
302+
def enum_display_monitors(self, screen=0):
305303
''' Get positions of one or more monitors.
306304
Returns a dict with minimal requirements (see MSS class).
307305
'''
@@ -345,24 +343,28 @@ def get_pixels(self, monitor):
345343
width, height = monitor[b'width'], monitor[b'height']
346344
left, top = monitor[b'left'], monitor[b'top']
347345
rect = CGRect((left, top), (width, height))
348-
self.image = CGWindowListCreateImage(
346+
image = CGWindowListCreateImage(
349347
rect, kCGWindowListOptionOnScreenOnly,
350348
kCGNullWindowID, kCGWindowImageDefault)
351-
return self.image
349+
if not image:
350+
raise ScreenshotError('CGWindowListCreateImage() failed.')
351+
return image
352352

353353
def save_img(self, data, width, height, output):
354-
''' Use my own save_img() method. Because I'm Mac ... '''
354+
''' Use my own save_img() method. Because I'm Mac! '''
355355

356-
self.debug('MSSMac: save_img()')
356+
self.debug('MSSMac: save_img')
357357

358-
dpi = 72
359358
url = NSURL.fileURLWithPath_(output)
360359
dest = CGImageDestinationCreateWithURL(url, kUTTypePNG, 1, None)
361-
properties = {
362-
kCGImagePropertyDPIWidth: dpi,
363-
kCGImagePropertyDPIHeight: dpi,
364-
}
365-
CGImageDestinationAddImage(dest, data, properties)
360+
if not dest:
361+
msg = 'Error while trying to create the image destination ' \
362+
'to "{}".'.format(output)
363+
raise ScreenshotError(msg)
364+
CGImageDestinationAddImage(dest, data, None)
365+
if not CGImageDestinationFinalize(dest):
366+
msg = 'Impossible to write data to file "{}".'.format(output)
367+
raise ScreenshotError(msg)
366368

367369

368370
class MSSLinux(MSS):
@@ -385,7 +387,7 @@ def init(self):
385387

386388
x11 = find_library('X11')
387389
if not x11:
388-
raise ScreenshotError('MSSLinux: no X11 library found.')
390+
raise ScreenshotError('No X11 library found.')
389391

390392
self.xlib = cdll.LoadLibrary(x11)
391393
self.debug('init', 'xlib', self.xlib)
@@ -401,7 +403,7 @@ def init(self):
401403
else:
402404
disp = environ['DISPLAY']
403405
except KeyError:
404-
err = 'MSSLinux: $DISPLAY not set. Stopping to prevent segfault.'
406+
err = '$DISPLAY not set. Stopping to prevent segfault.'
405407
raise ScreenshotError(err)
406408
self.debug('init', '$DISPLAY', disp)
407409

@@ -456,7 +458,7 @@ def _x11_config(self):
456458
self.debug('_x11_config')
457459

458460
monitors = expanduser('~/.config/monitors.xml')
459-
if not isfile(monitors):
461+
if not os.path.isfile(monitors):
460462
self.debug('ERROR', 'MSSLinux: _x11_config() failed.')
461463
return
462464

@@ -493,7 +495,7 @@ def _xfce4_config(self):
493495

494496
path_ = '~/.config/xfce4/xfconf/xfce-perchannel-xml/displays.xml'
495497
monitors = expanduser(path_)
496-
if not isfile(monitors):
498+
if not os.path.isfile(monitors):
497499
self.debug('ERROR', 'MSSLinux: _xfce4_config() failed.')
498500
return
499501

@@ -519,7 +521,7 @@ def _xfce4_config(self):
519521
b'rotation': rotation
520522
})
521523

522-
def enum_display_monitors(self, screen=-1):
524+
def enum_display_monitors(self, screen=0):
523525
''' Get positions of one or more monitors.
524526
Returns a dict with minimal requirements (see MSS class).
525527
'''
@@ -560,10 +562,10 @@ def get_pixels(self, monitor):
560562
# expected LP_Display instance instead of LP_XWindowAttributes
561563
root = cast(self.root, POINTER(Display))
562564

563-
image = self.xlib.XGetImage(self.display, root, left, top, width,
564-
height, allplanes, ZPixmap)
565-
if not image:
566-
raise ScreenshotError('MSSLinux: XGetImage() failed.')
565+
ximage = self.xlib.XGetImage(self.display, root, left, top, width,
566+
height, allplanes, ZPixmap)
567+
if not ximage:
568+
raise ScreenshotError('XGetImage() failed.')
567569

568570
def pix(pixel, _resultats={}):
569571
''' Apply shifts to a pixel to get the RGB values.
@@ -576,12 +578,12 @@ def pix(pixel, _resultats={}):
576578

577579
# http://cgit.freedesktop.org/xorg/lib/libX11/tree/src/ImUtil.c#n444
578580
get_pix = self.xlib.XGetPixel
579-
pixels = [pix(get_pix(image, x, y))
581+
pixels = [pix(get_pix(ximage, x, y))
580582
for y in range(height) for x in range(width)]
581583

582-
self.xlib.XFree(image)
583-
self.image = b''.join(pixels)
584-
return self.image
584+
self.xlib.XFree(ximage)
585+
image = b''.join(pixels)
586+
return image
585587

586588

587589
class MSSWindows(MSS):
@@ -712,7 +714,7 @@ def get_pixels(self, monitor):
712714
windll.gdi32.DeleteObject(bmp)
713715

714716
if bits != height or len(pixels.raw) != buffer_len:
715-
raise ScreenshotError('MSSWindows: GetDIBits() failed.')
717+
raise ScreenshotError('GetDIBits() failed.')
716718

717719
# Note that the origin of the returned image is in the
718720
# bottom-left corner, 32-bit aligned. And it is BGR.
@@ -733,8 +735,8 @@ def _arrange(self, data, width, height):
733735
for x in range(0, width - 2, 3):
734736
scanlines[off+x:off+x+3] = \
735737
b(data[offset+x+2]), b(data[offset+x+1]), b(data[offset+x])
736-
self.image = b''.join(scanlines)
737-
return self.image
738+
image = b''.join(scanlines)
739+
return image
738740

739741

740742
def main(argv=[]):
@@ -758,36 +760,41 @@ def timer(msg):
758760
}
759761
mss = systems[system()](debug='--debug' in argv)
760762

761-
# One screen shot per monitor
762-
with timer('Screen shots'):
763-
for filename in mss.save():
764-
print(' File: {}'.format(filename))
765-
766-
# Screen shot of the monitor 1
767-
with timer('Monitor 1 '):
768-
for filename in mss.save(output='monitor-1', screen=1):
769-
print(' File: {}'.format(filename))
770-
771-
# A shot to grab them all :)
772-
with timer('All in one '):
773-
for filename in mss.save(screen=-1):
774-
print(' File: {}'.format(filename))
775-
776-
# Example with a callback
777-
def on_exists(fname):
778-
''' Callback example when we try to overwrite an existing
779-
screen shot.
780-
'''
781-
from os import rename
782-
newfile = fname + '.old'
783-
print(' Renaming {} to {}'.format(fname, newfile))
784-
rename(fname, newfile)
785-
return True
786-
787-
# Screen shot of the monitor 1, with callback
788-
with timer('Monitor 1 '):
789-
for fname in mss.save(output='mon-1', screen=1, callback=on_exists):
790-
print(' File: {}'.format(fname))
763+
try:
764+
# One screen shot per monitor
765+
with timer('Screen shots'):
766+
for filename in mss.save():
767+
print(' File: {}'.format(filename))
768+
769+
# Screen shot of the monitor 1
770+
with timer('Monitor 1 '):
771+
for filename in mss.save(output='monitor-%d.png', screen=1):
772+
print(' File: {}'.format(filename))
773+
774+
# A shot to grab them all :)
775+
with timer('All in one '):
776+
for filename in mss.save(output='full-screenshot.png', screen=-1):
777+
print(' File: {}'.format(filename))
778+
779+
# Example with a callback
780+
def on_exists(fname):
781+
''' Callback example when we try to overwrite an existing
782+
screen shot.
783+
'''
784+
from os import rename
785+
newfile = fname + '.old'
786+
print(' Renaming {} to {}'.format(fname, newfile))
787+
rename(fname, newfile)
788+
return True
789+
790+
# Screen shot of the monitor 1, with callback
791+
with timer('Monitor 1 '):
792+
for fname in mss.save(output='mon-%d.png',
793+
screen=1, callback=on_exists):
794+
print(' File: {}'.format(fname))
795+
except ScreenshotError as ex:
796+
print(ex)
797+
return 1
791798
return 0
792799

793800

0 commit comments

Comments
 (0)