1
2
3
4
5 import math
6
7 try:
8 import ImageChops
9 import ImageGrab
10 import ImageDraw
11 HAVE_PIL = True
12 except:
13 try:
14 from PIL import ImageChops
15 from PIL import ImageGrab
16 from PIL import ImageDraw
17 HAVE_PIL = True
18 except:
19 HAVE_PIL = False
20
22 """Get screenshots."""
23
25 """Draw a black rectangle.
26 @param img: PIL Image object
27 @param xy: Coordinates as refined in PIL rectangle() doc
28 @return: Image with black rectangle
29 """
30 dr = ImageDraw.Draw(img)
31 dr.rectangle(xy, fill="black", outline="black")
32 return img
33
35 """Is Python Image Library installed?
36 @return: installed status.
37 """
38 return HAVE_PIL
39
40 - def equal(self, img1, img2, skip_area=None):
41 """Compares two screenshots using Root-Mean-Square Difference (RMS).
42 @param img1: screenshot to compare.
43 @param img2: screenshot to compare.
44 @return: equal status.
45 """
46 if not HAVE_PIL:
47 return None
48
49
50
51
52
53
54 if skip_area:
55
56 img1 = img1.copy()
57 img2 = img2.copy()
58
59 for img in (img1, img2):
60 self._draw_rectangle(img, skip_area)
61
62
63
64
65 diff = ImageChops.difference(img1, img2)
66 h = diff.histogram()
67 sq = (value * ((idx % 256)**2) for idx, value in enumerate(h))
68 sum_of_squares = sum(sq)
69 rms = math.sqrt(sum_of_squares/float(img1.size[0] * img1.size[1]))
70
71
72 return rms < 8
73
75 """Take a screenshot.
76 @return: screenshot or None.
77 """
78 if not HAVE_PIL:
79 return None
80
81 return ImageGrab.grab()
82