diff options
Diffstat (limited to '')
| -rw-r--r-- | tests/screenshot.py | 106 |
1 files changed, 62 insertions, 44 deletions
diff --git a/tests/screenshot.py b/tests/screenshot.py index 6ccb593..0786d48 100644 --- a/tests/screenshot.py +++ b/tests/screenshot.py @@ -14,7 +14,7 @@ except ImportError: from tests import PROJECT_ROOT -STARTF_USESHOWWINDOW = getattr(subprocess, 'STARTF_USESHOWWINDOW', 1) +STARTF_USESHOWWINDOW = getattr(subprocess, "STARTF_USESHOWWINDOW", 1) STILL_ACTIVE = 259 SW_MAXIMIZE = 3 @@ -23,24 +23,24 @@ 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), + ("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): @@ -49,7 +49,7 @@ class StartupInfo(ctypes.Structure): :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__() + super().__init__() self.cb = ctypes.sizeof(self) if maximize: self.dwFlags |= STARTF_USESHOWWINDOW @@ -62,14 +62,14 @@ 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), + ("hProcess", ctypes.c_void_p), + ("hThread", ctypes.c_void_p), + ("dwProcessId", ctypes.c_ulong), + ("dwThreadId", ctypes.c_ulong), ] -class RunNewConsole(object): +class RunNewConsole: """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 @@ -84,20 +84,27 @@ class RunNewConsole(object): :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') + title = "pytest-{}-{}".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.command_str = subprocess.list2cmdline(command).encode("ascii") + self._handles = [] self._kernel32 = ctypes.LibraryLoader(ctypes.WinDLL).kernel32 - self._kernel32.GetExitCodeProcess.argtypes = [ctypes.c_void_p, ctypes.POINTER(ctypes.c_ulong)] + 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. + self._kernel32.CloseHandle( + self._handles.pop(0) + ) # .pop() is thread safe. except IndexError: break @@ -111,9 +118,9 @@ class RunNewConsole(object): False, # bInheritHandles subprocess.CREATE_NEW_CONSOLE, # dwCreationFlags None, # lpEnvironment - str(PROJECT_ROOT).encode('ascii'), # lpCurrentDirectory + str(PROJECT_ROOT).encode("ascii"), # lpCurrentDirectory ctypes.byref(self.startup_info), # lpStartupInfo - ctypes.byref(self.process_info) # lpProcessInformation + ctypes.byref(self.process_info), # lpProcessInformation ): raise ctypes.WinError() @@ -125,7 +132,9 @@ class RunNewConsole(object): 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) + self.hwnd = ctypes.windll.user32.FindWindowA( + None, self.startup_info.lpTitle + ) if self.hwnd: break time.sleep(0.1) @@ -141,7 +150,9 @@ class RunNewConsole(object): 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)): + if not self._kernel32.GetExitCodeProcess( + self.process_info.hProcess, ctypes.byref(status) + ): raise ctypes.WinError() assert status.value == 0 finally: @@ -154,9 +165,11 @@ class RunNewConsole(object): :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. + 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) + left, top, right, bottom = struct.unpack("llll", rect.raw) width, height = right - left, bottom - top assert width > 1 assert height > 1 @@ -179,8 +192,7 @@ def iter_rows(pil_image): :rtype: tuple """ iterator = izip(*(iter(pil_image.getdata()),) * pil_image.width) - for row in iterator: - yield row + yield from iterator def get_most_interesting_row(pil_image): @@ -224,11 +236,13 @@ def count_subimages(screenshot, subimg): 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: + 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: + 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 @@ -248,11 +262,12 @@ def try_candidates(screenshot, subimg_candidates, expected_count): :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: + with rgba_s.convert(mode="RGB") as subimg: # Make sure subimage isn't too large. assert subimg.width < 256 assert subimg.height < 256 @@ -277,14 +292,17 @@ def screenshot_until_match(save_to, timeout, subimg_candidates, expected_count, :param iter gen: Generator yielding window position and size to crop screenshot to. """ from PIL import ImageGrab - assert save_to.endswith('.png') + + 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) + 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 |