"""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 colorclass.windows import WINDOWS_CODES
from tests.conftest import PROJECT_ROOT
STARTF_USEFILLATTRIBUTE = 0x00000010
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, white_bg=False):
"""Constructor.
:param bool maximize: Start process in new console window, maximized.
:param bool white_bg: New console window will be black text on white background.
: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)
if white_bg:
self.dwFlags |= STARTF_USEFILLATTRIBUTE
self.dwFillAttribute = WINDOWS_CODES['hibgwhite'] | WINDOWS_CODES['black']
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, white_bg=False):
"""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.
:param bool white_bg: New console window will be black text on white background.
"""
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, white_bg=white_bg)
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 written 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)
