1from pyvirtualdisplay.display import Display 2from PIL import Image 3from PIL import ImageChops 4import logging 5import pyscreenshot 6import time 7 8 9log = logging.getLogger(__name__) 10 11 12# class DisplayError(Exception): 13# pass 14 15class DisplayTimeoutError(Exception): 16 pass 17 18 19class SmartDisplay(Display): 20 pyscreenshot_backend = None 21 pyscreenshot_childprocess = True 22 23 def autocrop(self, im): 24 '''Crop borders off an image. 25 26 :param im: Source image. 27 :param bgcolor: Background color, using either a color tuple or a color name (1.1.4 only). 28 :return: An image without borders, or None if there's no actual content in the image. 29 ''' 30 if im.mode != "RGB": 31 im = im.convert("RGB") 32 bg = Image.new("RGB", im.size, self.bgcolor) 33 diff = ImageChops.difference(im, bg) 34 bbox = diff.getbbox() 35 if bbox: 36 return im.crop(bbox) 37 return None # no contents 38 39 def grab(self, autocrop=True): 40 try: 41 # first try newer pyscreenshot version 42 img = pyscreenshot.grab( 43 childprocess=self.pyscreenshot_childprocess, 44 backend=self.pyscreenshot_backend, 45 ) 46 except TypeError: 47 # try older pyscreenshot version 48 img = pyscreenshot.grab() 49 50 if autocrop: 51 img = self.autocrop(img) 52 return img 53 54 def waitgrab(self, timeout=60, autocrop=True, cb_imgcheck=None): 55 '''start process and create screenshot. 56 Repeat screenshot until it is not empty and 57 cb_imgcheck callback function returns True 58 for current screenshot. 59 60 :param autocrop: True -> crop screenshot 61 :param timeout: int 62 :param cb_imgcheck: None or callback for testing img, 63 True = accept img, 64 False = reject img 65 ''' 66 t = 0 67 sleep_time = 0.3 # for fast windows 68 repeat_time = 1 69 while 1: 70 log.debug('sleeping %s secs' % str(sleep_time)) 71 time.sleep(sleep_time) 72 t += sleep_time 73 img = self.grab(autocrop=autocrop) 74 if img: 75 if not cb_imgcheck: 76 break 77 if cb_imgcheck(img): 78 break 79 sleep_time = repeat_time 80 repeat_time += 1 # progressive 81 if t > timeout: 82 msg = 'Timeout! elapsed time:%s timeout:%s ' % (t, timeout) 83 raise DisplayTimeoutError(msg) 84 break 85 86 log.debug('screenshot is empty, next try..') 87 assert img 88# if not img: 89# log.debug('screenshot is empty!') 90 return img 91