summaryrefslogtreecommitdiffstats
path: root/tests/screenshot.py
blob: 6ccb5930edd6afe12ea1fc7a7403f798f027b152 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
"""Take screenshots and search for subimages in images."""

import ctypes
import os
import random
import struct
import subprocess
import time

try:
    from itertools import izip
except ImportError:
    izip = zip  # Py3

from tests import PROJECT_ROOT

STARTF_USESHOWWINDOW = getattr(subprocess, 'STARTF_USESHOWWINDOW', 1)
STILL_ACTIVE = 259
SW_MAXIMIZE = 3


class StartupInfo(ctypes.Structure):
    """STARTUPINFO structure."""

    _fields_ = [
        ('cb', ctypes.c_ulong),
        ('lpReserved', ctypes.c_char_p),
        ('lpDesktop', ctypes.c_char_p),
        ('lpTitle', ctypes.c_char_p),
        ('dwX', ctypes.c_ulong),
        ('dwY', ctypes.c_ulong),
        ('dwXSize', ctypes.c_ulong),
        ('dwYSize', ctypes.c_ulong),
        ('dwXCountChars', ctypes.c_ulong),
        ('dwYCountChars', ctypes.c_ulong),
        ('dwFillAttribute', ctypes.c_ulong),
        ('dwFlags', ctypes.c_ulong),
        ('wShowWindow', ctypes.c_ushort),
        ('cbReserved2', ctypes.c_ushort),
        ('lpReserved2', ctypes.c_char_p),
        ('hStdInput', ctypes.c_ulong),
        ('hStdOutput', ctypes.c_ulong),
        ('hStdError', ctypes.c_ulong),
    ]

    def __init__(self, maximize=False, title=None):
        """Constructor.

        :param bool maximize: Start process in new console window, maximized.
        :param bytes title: Set new window title to this instead of exe path.
        """
        super(StartupInfo, self).__init__()
        self.cb = ctypes.sizeof(self)
        if maximize:
            self.dwFlags |= STARTF_USESHOWWINDOW
            self.wShowWindow = SW_MAXIMIZE
        if title:
            self.lpTitle = ctypes.c_char_p(title)


class ProcessInfo(ctypes.Structure):
    """PROCESS_INFORMATION structure."""

    _fields_ = [
        ('hProcess', ctypes.c_void_p),
        ('hThread', ctypes.c_void_p),
        ('dwProcessId', ctypes.c_ulong),
        ('dwThreadId', ctypes.c_ulong),
    ]


class RunNewConsole(object):
    """Run the command in a new console window. Windows only. Use in a with statement.

    subprocess sucks and really limits your access to the win32 API. Its implementation is half-assed. Using this so
    that STARTUPINFO.lpTitle actually works and STARTUPINFO.dwFillAttribute produce the expected result.
    """

    def __init__(self, command, maximized=False, title=None):
        """Constructor.

        :param iter command: Command to run.
        :param bool maximized: Start process in new console window, maximized.
        :param bytes title: Set new window title to this. Needed by user32.FindWindow.
        """
        if title is None:
            title = 'pytest-{0}-{1}'.format(os.getpid(), random.randint(1000, 9999)).encode('ascii')
        self.startup_info = StartupInfo(maximize=maximized, title=title)
        self.process_info = ProcessInfo()
        self.command_str = subprocess.list2cmdline(command).encode('ascii')
        self._handles = list()
        self._kernel32 = ctypes.LibraryLoader(ctypes.WinDLL).kernel32
        self._kernel32.GetExitCodeProcess.argtypes = [ctypes.c_void_p, ctypes.POINTER(ctypes.c_ulong)]
        self._kernel32.GetExitCodeProcess.restype = ctypes.c_long

    def __del__(self):
        """Close win32 handles."""
        while self._handles:
            try:
                self._kernel32.CloseHandle(self._handles.pop(0))  # .pop() is thread safe.
            except IndexError:
                break

    def __enter__(self):
        """Entering the `with` block. Runs the process."""
        if not self._kernel32.CreateProcessA(
            None,  # lpApplicationName
            self.command_str,  # lpCommandLine
            None,  # lpProcessAttributes
            None,  # lpThreadAttributes
            False,  # bInheritHandles
            subprocess.CREATE_NEW_CONSOLE,  # dwCreationFlags
            None,  # lpEnvironment
            str(PROJECT_ROOT).encode('ascii'),  # lpCurrentDirectory
            ctypes.byref(self.startup_info),  # lpStartupInfo
            ctypes.byref(self.process_info)  # lpProcessInformation
        ):
            raise ctypes.WinError()

        # Add handles added by the OS.
        self._handles.append(self.process_info.hProcess)
        self._handles.append(self.process_info.hThread)

        # Get hWnd.
        self.hwnd = 0
        for _ in range(int(5 / 0.1)):
            # Takes time for console window to initialize.
            self.hwnd = ctypes.windll.user32.FindWindowA(None, self.startup_info.lpTitle)
            if self.hwnd:
                break
            time.sleep(0.1)
        assert self.hwnd

        # Return generator that yields window size/position.
        return self._iter_pos()

    def __exit__(self, *_):
        """Cleanup."""
        try:
            # Verify process exited 0.
            status = ctypes.c_ulong(STILL_ACTIVE)
            while status.value == STILL_ACTIVE:
                time.sleep(0.1)
                if not self._kernel32.GetExitCodeProcess(self.process_info.hProcess, ctypes.byref(status)):
                    raise ctypes.WinError()
            assert status.value == 0
        finally:
            # Close handles.
            self.__del__()

    def _iter_pos(self):
        """Yield new console window's current position and dimensions.

        :return: Yields region the new window is in (left, upper, right, lower).
        :rtype: tuple
        """
        rect = ctypes.create_string_buffer(16)  # To be written to by GetWindowRect. RECT structure.
        while ctypes.windll.user32.GetWindowRect(self.hwnd, rect):
            left, top, right, bottom = struct.unpack('llll', rect.raw)
            width, height = right - left, bottom - top
            assert width > 1
            assert height > 1
            yield left, top, right, bottom
        raise StopIteration


def iter_rows(pil_image):
    """Yield tuple of pixels for each row in the image.

    itertools.izip in Python 2.x and zip in Python 3.x are writen in C. Much faster than anything else I've found
    written in pure Python.

    From:
    http://stackoverflow.com/questions/1624883/alternative-way-to-split-a-list-into-groups-of-n/1625023#1625023

    :param PIL.Image.Image pil_image: Image to read from.

    :return: Yields rows.
    :rtype: tuple
    """
    iterator = izip(*(iter(pil_image.getdata()),) * pil_image.width)
    for row in iterator:
        yield row


def get_most_interesting_row(pil_image):
    """Look for a row in the image that has the most unique pixels.

    :param PIL.Image.Image pil_image: Image to read from.

    :return: Row (tuple of pixel tuples), row as a set, first pixel tuple, y offset from top.
    :rtype: tuple
    """
    final = (None, set(), None, None)  # row, row_set, first_pixel, y_pos
    for y_pos, row in enumerate(iter_rows(pil_image)):
        row_set = set(row)
        if len(row_set) > len(final[1]):
            final = row, row_set, row[0], y_pos
        if len(row_set) == pil_image.width:
            break  # Can't get bigger.
    return final


def count_subimages(screenshot, subimg):
    """Check how often subimg appears in the screenshot image.

    :param PIL.Image.Image screenshot: Screen shot to search through.
    :param PIL.Image.Image subimg: Subimage to search for.

    :return: Number of times subimg appears in the screenshot.
    :rtype: int
    """
    # Get row to search for.
    si_pixels = list(subimg.getdata())  # Load entire subimg into memory.
    si_width = subimg.width
    si_height = subimg.height
    si_row, si_row_set, si_pixel, si_y = get_most_interesting_row(subimg)
    occurrences = 0

    # Look for subimg row in screenshot, then crop and compare pixel arrays.
    for y_pos, row in enumerate(iter_rows(screenshot)):
        if si_row_set - set(row):
            continue  # Some pixels not found.
        for x_pos in range(screenshot.width - si_width + 1):
            if row[x_pos] != si_pixel:
                continue  # First pixel does not match.
            if row[x_pos:x_pos + si_width] != si_row:
                continue  # Row does not match.
            # Found match for interesting row of subimg in screenshot.
            y_corrected = y_pos - si_y
            with screenshot.crop((x_pos, y_corrected, x_pos + si_width, y_corrected + si_height)) as cropped:
                if list(cropped.getdata()) == si_pixels:
                    occurrences += 1

    return occurrences


def try_candidates(screenshot, subimg_candidates, expected_count):
    """Call count_subimages() for each subimage candidate until.

    If you get ImportError run "pip install pillow". Only OSX and Windows is supported.

    :param PIL.Image.Image screenshot: Screen shot to search through.
    :param iter subimg_candidates: Subimage paths to look for. List of strings.
    :param int expected_count: Try until any a subimage candidate is found this many times.

    :return: Number of times subimg appears in the screenshot.
    :rtype: int
    """
    from PIL import Image
    count_found = 0

    for subimg_path in subimg_candidates:
        with Image.open(subimg_path) as rgba_s:
            with rgba_s.convert(mode='RGB') as subimg:
                # Make sure subimage isn't too large.
                assert subimg.width < 256
                assert subimg.height < 256

                # Count.
                count_found = count_subimages(screenshot, subimg)
                if count_found == expected_count:
                    break  # No need to try other candidates.

    return count_found


def screenshot_until_match(save_to, timeout, subimg_candidates, expected_count, gen):
    """Take screenshots until one of the 'done' subimages is found. Image is saved when subimage found or at timeout.

    If you get ImportError run "pip install pillow". Only OSX and Windows is supported.

    :param str save_to: Save screenshot to this PNG file path when expected count found or timeout.
    :param int timeout: Give up after these many seconds.
    :param iter subimg_candidates: Subimage paths to look for. List of strings.
    :param int expected_count: Keep trying until any of subimg_candidates is found this many times.
    :param iter gen: Generator yielding window position and size to crop screenshot to.
    """
    from PIL import ImageGrab
    assert save_to.endswith('.png')
    stop_after = time.time() + timeout

    # Take screenshots until subimage is found.
    while True:
        with ImageGrab.grab(next(gen)) as rgba:
            with rgba.convert(mode='RGB') as screenshot:
                count_found = try_candidates(screenshot, subimg_candidates, expected_count)
                if count_found == expected_count or time.time() > stop_after:
                    screenshot.save(save_to)
                    assert count_found == expected_count
                    return
        time.sleep(0.5)