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