367 lines
11 KiB
Python
367 lines
11 KiB
Python
|
#!/usr/bin/env python
|
||
|
# Copyright (c) 2011 The Chromium Authors. All rights reserved.
|
||
|
# Use of this source code is governed by a BSD-style license that can be
|
||
|
# found in the LICENSE file.
|
||
|
|
||
|
"""SiteCompare module for invoking, locating, and manipulating windows.
|
||
|
|
||
|
This module is a catch-all wrapper for operating system UI functionality
|
||
|
that doesn't belong in other modules. It contains functions for finding
|
||
|
particular windows, scraping their contents, and invoking processes to
|
||
|
create them.
|
||
|
"""
|
||
|
|
||
|
import os
|
||
|
import string
|
||
|
import time
|
||
|
|
||
|
import PIL.ImageGrab
|
||
|
import pywintypes
|
||
|
import win32event
|
||
|
import win32gui
|
||
|
import win32process
|
||
|
|
||
|
|
||
|
def FindChildWindows(hwnd, path):
|
||
|
"""Find a set of windows through a path specification.
|
||
|
|
||
|
Args:
|
||
|
hwnd: Handle of the parent window
|
||
|
path: Path to the window to find. Has the following form:
|
||
|
"foo/bar/baz|foobar/|foobarbaz"
|
||
|
The slashes specify the "path" to the child window.
|
||
|
The text is the window class, a pipe (if present) is a title.
|
||
|
* is a wildcard and will find all child windows at that level
|
||
|
|
||
|
Returns:
|
||
|
A list of the windows that were found
|
||
|
"""
|
||
|
windows_to_check = [hwnd]
|
||
|
|
||
|
# The strategy will be to take windows_to_check and use it
|
||
|
# to find a list of windows that match the next specification
|
||
|
# in the path, then repeat with the list of found windows as the
|
||
|
# new list of windows to check
|
||
|
for segment in path.split("/"):
|
||
|
windows_found = []
|
||
|
check_values = segment.split("|")
|
||
|
|
||
|
# check_values is now a list with the first element being
|
||
|
# the window class, the second being the window caption.
|
||
|
# If the class is absent (or wildcarded) set it to None
|
||
|
if check_values[0] == "*" or not check_values[0]: check_values[0] = None
|
||
|
|
||
|
# If the window caption is also absent, force it to None as well
|
||
|
if len(check_values) == 1: check_values.append(None)
|
||
|
|
||
|
# Loop through the list of windows to check
|
||
|
for window_check in windows_to_check:
|
||
|
window_found = None
|
||
|
while window_found != 0: # lint complains, but 0 != None
|
||
|
if window_found is None: window_found = 0
|
||
|
try:
|
||
|
# Look for the next sibling (or first sibling if window_found is 0)
|
||
|
# of window_check with the specified caption and/or class
|
||
|
window_found = win32gui.FindWindowEx(
|
||
|
window_check, window_found, check_values[0], check_values[1])
|
||
|
except pywintypes.error, e:
|
||
|
# FindWindowEx() raises error 2 if not found
|
||
|
if e[0] == 2:
|
||
|
window_found = 0
|
||
|
else:
|
||
|
raise e
|
||
|
|
||
|
# If FindWindowEx struck gold, add to our list of windows found
|
||
|
if window_found: windows_found.append(window_found)
|
||
|
|
||
|
# The windows we found become the windows to check for the next segment
|
||
|
windows_to_check = windows_found
|
||
|
|
||
|
return windows_found
|
||
|
|
||
|
|
||
|
def FindChildWindow(hwnd, path):
|
||
|
"""Find a window through a path specification.
|
||
|
|
||
|
This method is a simple wrapper for FindChildWindows() for the
|
||
|
case (the majority case) where you expect to find a single window
|
||
|
|
||
|
Args:
|
||
|
hwnd: Handle of the parent window
|
||
|
path: Path to the window to find. See FindChildWindows()
|
||
|
|
||
|
Returns:
|
||
|
The window that was found
|
||
|
"""
|
||
|
return FindChildWindows(hwnd, path)[0]
|
||
|
|
||
|
|
||
|
def ScrapeWindow(hwnd, rect=None):
|
||
|
"""Scrape a visible window and return its contents as a bitmap.
|
||
|
|
||
|
Args:
|
||
|
hwnd: handle of the window to scrape
|
||
|
rect: rectangle to scrape in client coords, defaults to the whole thing
|
||
|
If specified, it's a 4-tuple of (left, top, right, bottom)
|
||
|
|
||
|
Returns:
|
||
|
An Image containing the scraped data
|
||
|
"""
|
||
|
# Activate the window
|
||
|
SetForegroundWindow(hwnd)
|
||
|
|
||
|
# If no rectangle was specified, use the fill client rectangle
|
||
|
if not rect: rect = win32gui.GetClientRect(hwnd)
|
||
|
|
||
|
upper_left = win32gui.ClientToScreen(hwnd, (rect[0], rect[1]))
|
||
|
lower_right = win32gui.ClientToScreen(hwnd, (rect[2], rect[3]))
|
||
|
rect = upper_left+lower_right
|
||
|
|
||
|
return PIL.ImageGrab.grab(rect)
|
||
|
|
||
|
|
||
|
def SetForegroundWindow(hwnd):
|
||
|
"""Bring a window to the foreground."""
|
||
|
win32gui.SetForegroundWindow(hwnd)
|
||
|
|
||
|
|
||
|
def InvokeAndWait(path, cmdline="", timeout=10, tick=1.):
|
||
|
"""Invoke an application and wait for it to bring up a window.
|
||
|
|
||
|
Args:
|
||
|
path: full path to the executable to invoke
|
||
|
cmdline: command line to pass to executable
|
||
|
timeout: how long (in seconds) to wait before giving up
|
||
|
tick: length of time to wait between checks
|
||
|
|
||
|
Returns:
|
||
|
A tuple of handles to the process and the application's window,
|
||
|
or (None, None) if it timed out waiting for the process
|
||
|
"""
|
||
|
|
||
|
def EnumWindowProc(hwnd, ret):
|
||
|
"""Internal enumeration func, checks for visibility and proper PID."""
|
||
|
if win32gui.IsWindowVisible(hwnd): # don't bother even checking hidden wnds
|
||
|
pid = win32process.GetWindowThreadProcessId(hwnd)[1]
|
||
|
if pid == ret[0]:
|
||
|
ret[1] = hwnd
|
||
|
return 0 # 0 means stop enumeration
|
||
|
return 1 # 1 means continue enumeration
|
||
|
|
||
|
# We don't need to change anything about the startupinfo structure
|
||
|
# (the default is quite sufficient) but we need to create it just the
|
||
|
# same.
|
||
|
sinfo = win32process.STARTUPINFO()
|
||
|
|
||
|
proc = win32process.CreateProcess(
|
||
|
path, # path to new process's executable
|
||
|
cmdline, # application's command line
|
||
|
None, # process security attributes (default)
|
||
|
None, # thread security attributes (default)
|
||
|
False, # inherit parent's handles
|
||
|
0, # creation flags
|
||
|
None, # environment variables
|
||
|
None, # directory
|
||
|
sinfo) # default startup info
|
||
|
|
||
|
# Create process returns (prochandle, pid, threadhandle, tid). At
|
||
|
# some point we may care about the other members, but for now, all
|
||
|
# we're after is the pid
|
||
|
pid = proc[2]
|
||
|
|
||
|
# Enumeration APIs can take an arbitrary integer, usually a pointer,
|
||
|
# to be passed to the enumeration function. We'll pass a pointer to
|
||
|
# a structure containing the PID we're looking for, and an empty out
|
||
|
# parameter to hold the found window ID
|
||
|
ret = [pid, None]
|
||
|
|
||
|
tries_until_timeout = timeout/tick
|
||
|
num_tries = 0
|
||
|
|
||
|
# Enumerate top-level windows, look for one with our PID
|
||
|
while num_tries < tries_until_timeout and ret[1] is None:
|
||
|
try:
|
||
|
win32gui.EnumWindows(EnumWindowProc, ret)
|
||
|
except pywintypes.error, e:
|
||
|
# error 0 isn't an error, it just meant the enumeration was
|
||
|
# terminated early
|
||
|
if e[0]: raise e
|
||
|
|
||
|
time.sleep(tick)
|
||
|
num_tries += 1
|
||
|
|
||
|
# TODO(jhaas): Should we throw an exception if we timeout? Or is returning
|
||
|
# a window ID of None sufficient?
|
||
|
return (proc[0], ret[1])
|
||
|
|
||
|
|
||
|
def WaitForProcessExit(proc, timeout=None):
|
||
|
"""Waits for a given process to terminate.
|
||
|
|
||
|
Args:
|
||
|
proc: handle to process
|
||
|
timeout: timeout (in seconds). None = wait indefinitely
|
||
|
|
||
|
Returns:
|
||
|
True if process ended, False if timed out
|
||
|
"""
|
||
|
if timeout is None:
|
||
|
timeout = win32event.INFINITE
|
||
|
else:
|
||
|
# convert sec to msec
|
||
|
timeout *= 1000
|
||
|
|
||
|
return (win32event.WaitForSingleObject(proc, timeout) ==
|
||
|
win32event.WAIT_OBJECT_0)
|
||
|
|
||
|
|
||
|
def WaitForThrobber(hwnd, rect=None, timeout=20, tick=0.1, done=10):
|
||
|
"""Wait for a browser's "throbber" (loading animation) to complete.
|
||
|
|
||
|
Args:
|
||
|
hwnd: window containing the throbber
|
||
|
rect: rectangle of the throbber, in client coords. If None, whole window
|
||
|
timeout: if the throbber is still throbbing after this long, give up
|
||
|
tick: how often to check the throbber
|
||
|
done: how long the throbber must be unmoving to be considered done
|
||
|
|
||
|
Returns:
|
||
|
Number of seconds waited, -1 if timed out
|
||
|
"""
|
||
|
if not rect: rect = win32gui.GetClientRect(hwnd)
|
||
|
|
||
|
# last_throbber will hold the results of the preceding scrape;
|
||
|
# we'll compare it against the current scrape to see if we're throbbing
|
||
|
last_throbber = ScrapeWindow(hwnd, rect)
|
||
|
start_clock = time.clock()
|
||
|
timeout_clock = start_clock + timeout
|
||
|
last_changed_clock = start_clock;
|
||
|
|
||
|
while time.clock() < timeout_clock:
|
||
|
time.sleep(tick)
|
||
|
|
||
|
current_throbber = ScrapeWindow(hwnd, rect)
|
||
|
if current_throbber.tostring() != last_throbber.tostring():
|
||
|
last_throbber = current_throbber
|
||
|
last_changed_clock = time.clock()
|
||
|
else:
|
||
|
if time.clock() - last_changed_clock > done:
|
||
|
return last_changed_clock - start_clock
|
||
|
|
||
|
return -1
|
||
|
|
||
|
|
||
|
def MoveAndSizeWindow(wnd, position=None, size=None, child=None):
|
||
|
"""Moves and/or resizes a window.
|
||
|
|
||
|
Repositions and resizes a window. If a child window is provided,
|
||
|
the parent window is resized so the child window has the given size
|
||
|
|
||
|
Args:
|
||
|
wnd: handle of the frame window
|
||
|
position: new location for the frame window
|
||
|
size: new size for the frame window (or the child window)
|
||
|
child: handle of the child window
|
||
|
|
||
|
Returns:
|
||
|
None
|
||
|
"""
|
||
|
rect = win32gui.GetWindowRect(wnd)
|
||
|
|
||
|
if position is None: position = (rect[0], rect[1])
|
||
|
if size is None:
|
||
|
size = (rect[2]-rect[0], rect[3]-rect[1])
|
||
|
elif child is not None:
|
||
|
child_rect = win32gui.GetWindowRect(child)
|
||
|
slop = (rect[2]-rect[0]-child_rect[2]+child_rect[0],
|
||
|
rect[3]-rect[1]-child_rect[3]+child_rect[1])
|
||
|
size = (size[0]+slop[0], size[1]+slop[1])
|
||
|
|
||
|
win32gui.MoveWindow(wnd, # window to move
|
||
|
position[0], # new x coord
|
||
|
position[1], # new y coord
|
||
|
size[0], # new width
|
||
|
size[1], # new height
|
||
|
True) # repaint?
|
||
|
|
||
|
|
||
|
def EndProcess(proc, code=0):
|
||
|
"""Ends a process.
|
||
|
|
||
|
Wraps the OS TerminateProcess call for platform-independence
|
||
|
|
||
|
Args:
|
||
|
proc: process ID
|
||
|
code: process exit code
|
||
|
|
||
|
Returns:
|
||
|
None
|
||
|
"""
|
||
|
win32process.TerminateProcess(proc, code)
|
||
|
|
||
|
|
||
|
def URLtoFilename(url, path=None, extension=None):
|
||
|
"""Converts a URL to a filename, given a path.
|
||
|
|
||
|
This in theory could cause collisions if two URLs differ only
|
||
|
in unprintable characters (eg. http://www.foo.com/?bar and
|
||
|
http://www.foo.com/:bar. In practice this shouldn't be a problem.
|
||
|
|
||
|
Args:
|
||
|
url: The URL to convert
|
||
|
path: path to the directory to store the file
|
||
|
extension: string to append to filename
|
||
|
|
||
|
Returns:
|
||
|
filename
|
||
|
"""
|
||
|
trans = string.maketrans(r'\/:*?"<>|', '_________')
|
||
|
|
||
|
if path is None: path = ""
|
||
|
if extension is None: extension = ""
|
||
|
if len(path) > 0 and path[-1] != '\\': path += '\\'
|
||
|
url = url.translate(trans)
|
||
|
return "%s%s%s" % (path, url, extension)
|
||
|
|
||
|
|
||
|
def PreparePath(path):
|
||
|
"""Ensures that a given path exists, making subdirectories if necessary.
|
||
|
|
||
|
Args:
|
||
|
path: fully-qualified path of directory to ensure exists
|
||
|
|
||
|
Returns:
|
||
|
None
|
||
|
"""
|
||
|
try:
|
||
|
os.makedirs(path)
|
||
|
except OSError, e:
|
||
|
if e[0] != 17: raise e # error 17: path already exists
|
||
|
|
||
|
|
||
|
def main():
|
||
|
PreparePath(r"c:\sitecompare\scrapes\ie7")
|
||
|
# We're being invoked rather than imported. Let's do some tests
|
||
|
|
||
|
# Hardcode IE's location for the purpose of this test
|
||
|
(proc, wnd) = InvokeAndWait(
|
||
|
r"c:\program files\internet explorer\iexplore.exe")
|
||
|
|
||
|
# Find the browser pane in the IE window
|
||
|
browser = FindChildWindow(
|
||
|
wnd, "TabWindowClass/Shell DocObject View/Internet Explorer_Server")
|
||
|
|
||
|
# Move and size the window
|
||
|
MoveAndSizeWindow(wnd, (0, 0), (1024, 768), browser)
|
||
|
|
||
|
# Take a screenshot
|
||
|
i = ScrapeWindow(browser)
|
||
|
|
||
|
i.show()
|
||
|
|
||
|
EndProcess(proc, 0)
|
||
|
|
||
|
|
||
|
if __name__ == "__main__":
|
||
|
sys.exit(main())
|