Add a callback to save() method, in case where 'ouput' already exists

BoboTiG · BoboTiG · commit ce3cd0a1ae77 · 2014-03-18T11:56:21.000+01:00
diff --git a/README.rst b/README.rst
@@ -23,7 +23,7 @@ Passing the `--debug` argument will make it more verbose.
 
 
 Instance the good class
-========================
+=======================
 
 You can determine automatically which class to use::
 
@@ -50,8 +50,8 @@ When initialising an instance of MSS, you can enable debug output::
     mss = mss_class(debug=True)
 
 
-save(output='screenshot', screen=-1)
----------------------------------
+save(output='screenshot', screen=-1, callback=lambda *x: True)
+--------------------------------------------------------------
 
 For each monitor, grab a screen shot and save it to a file.
 
@@ -61,18 +61,23 @@ Parameters::
     screen - integer - grab one screen shot of all monitors (screen=-1)
                        grab one screen shot by monitor (screen=0)
                        grab the screen shot of the monitor N (screen=N)
+    callback - function - in case where output already exists, call
+                          the defined callback function with output
+                          as parameter. If it returns True, then
+                          continue; else ignores the monitor and
+                          switches to ne next.
 
 This is a generator which returns created files::
 
-    'screenshot-0.png',
     'screenshot-1.png',
+    'screenshot-2.png',
     ...,
     'screenshot-N.png'
     or
     'screenshot-full.png'
 
 
-Example
+Examples
 ========
 
 Then, it is quite simple::
@@ -91,9 +96,23 @@ Then, it is quite simple::
     for filename in mss.save(screen=-1):
         print('File: "{}" created.'.format(filename))
 
+    # Example with a callback
+    def on_exists(fname):
+        ''' Callback example when we try to overwrite an existing screen shot. '''
+
+        from os import rename
+        newfile = fname + '.old'
+        print('Renaming "{}" to "{}"'.format(fname, newfile))
+        rename(fname, newfile)
+        return True
+
+    # Screen shot of the monitor 1, with callback
+    for filename in mss.save(output='monitor-1', screen=1, callback=on_exists):
+        print('File: "{}" created.'.format(filename))
+
 
 Bonus
-======
+=====
 
 Just for fun ...
 Show us your screen shot with all monitors in one file, we will update the gallery.
diff --git a/mss.py b/mss.py
@@ -31,6 +31,7 @@
 
 
 from struct import pack
+from os.path import isfile
 from platform import system
 import sys
 import zlib
@@ -45,7 +46,7 @@ class ScreenshotError(Exception):
     from LaunchServices import kUTTypePNG
 elif system() == 'Linux':
     from os import environ
-    from os.path import expanduser, isfile
+    from os.path import expanduser
     import xml.etree.ElementTree as ET
     from ctypes.util import find_library
     from ctypes import byref, cast, cdll, POINTER, Structure, \
@@ -202,14 +203,19 @@ def get_pixels(self, monitor_infos):
         '''
         return NotImplemented
 
-    def save(self, output='screenshot', screen=0):
+    def save(self, output='screenshot', screen=0, callback=lambda *x: True):
         ''' For each monitor, grab a screen shot and save it to a file.
 
             Parameters:
              - output - string - the output filename without extension
              - screen - int - grab one screen shot of all monitors (screen=-1)
                               grab one screen shot by monitor (screen=0)
                               grab the screen shot of the monitor N (screen=N)
+             - callback - function - in case where output already exists, call
+                                     the defined callback function with output
+                                     as parameter. If it returns True, then
+                                     continue; else ignores the monitor and
+                                     switches to ne next.
 
             This is a generator which returns created files:
                 'screenshot-1.png',
@@ -235,6 +241,8 @@ def save(self, output='screenshot', screen=0):
                 filename = '{}.png'.format(output)
 
             if screen <= 0 or (screen > 0 and i+1 == screen):
+                if isfile(filename) and not callback(filename):
+                    continue
                 self.debug('save', 'filename', filename)
                 self.save_img(data=self.get_pixels(monitor),
                               width=monitor[b'width'],
@@ -763,6 +771,21 @@ def timer(msg):
     with timer('All in one  '):
         for filename in mss.save(screen=-1):
             print('        File: {}'.format(filename))
+
+    # Example with a callback
+    def on_exists(fname):
+        ''' Callback example when we try to overwrite an existing screen shot. '''
+
+        from os import rename
+        newfile = fname + '.old'
+        print('        Renaming {} to {}'.format(fname, newfile))
+        rename(fname, newfile)
+        return True
+
+    # Screen shot of the monitor 1, with callback
+    with timer('Monitor 1   '):
+        for filename in mss.save(output='monitor-1', screen=1, callback=on_exists):
+            print('        File: {}'.format(filename))
     return 0
 
 
