1 #!/usr/bin/env python 2 # Copyright (c) 2011 The Chromium Authors. All rights reserved. 3 # Use of this source code is governed by a BSD-style license that can be 4 # found in the LICENSE file. 5 6 """SiteCompare module for invoking, locating, and manipulating windows. 7 8 This module is a catch-all wrapper for operating system UI functionality 9 that doesn't belong in other modules. It contains functions for finding 10 particular windows, scraping their contents, and invoking processes to 11 create them. 12 """ 13 14 import os 15 import string 16 import time 17 18 import PIL.ImageGrab 19 import pywintypes 20 import win32event 21 import win32gui 22 import win32process 23 24 25 def FindChildWindows(hwnd, path): 26 """Find a set of windows through a path specification. 27 28 Args: 29 hwnd: Handle of the parent window 30 path: Path to the window to find. Has the following form: 31 "foo/bar/baz|foobar/|foobarbaz" 32 The slashes specify the "path" to the child window. 33 The text is the window class, a pipe (if present) is a title. 34 * is a wildcard and will find all child windows at that level 35 36 Returns: 37 A list of the windows that were found 38 """ 39 windows_to_check = [hwnd] 40 41 # The strategy will be to take windows_to_check and use it 42 # to find a list of windows that match the next specification 43 # in the path, then repeat with the list of found windows as the 44 # new list of windows to check 45 for segment in path.split("/"): 46 windows_found = [] 47 check_values = segment.split("|") 48 49 # check_values is now a list with the first element being 50 # the window class, the second being the window caption. 51 # If the class is absent (or wildcarded) set it to None 52 if check_values[0] == "*" or not check_values[0]: check_values[0] = None 53 54 # If the window caption is also absent, force it to None as well 55 if len(check_values) == 1: check_values.append(None) 56 57 # Loop through the list of windows to check 58 for window_check in windows_to_check: 59 window_found = None 60 while window_found != 0: # lint complains, but 0 != None 61 if window_found is None: window_found = 0 62 try: 63 # Look for the next sibling (or first sibling if window_found is 0) 64 # of window_check with the specified caption and/or class 65 window_found = win32gui.FindWindowEx( 66 window_check, window_found, check_values[0], check_values[1]) 67 except pywintypes.error, e: 68 # FindWindowEx() raises error 2 if not found 69 if e[0] == 2: 70 window_found = 0 71 else: 72 raise e 73 74 # If FindWindowEx struck gold, add to our list of windows found 75 if window_found: windows_found.append(window_found) 76 77 # The windows we found become the windows to check for the next segment 78 windows_to_check = windows_found 79 80 return windows_found 81 82 83 def FindChildWindow(hwnd, path): 84 """Find a window through a path specification. 85 86 This method is a simple wrapper for FindChildWindows() for the 87 case (the majority case) where you expect to find a single window 88 89 Args: 90 hwnd: Handle of the parent window 91 path: Path to the window to find. See FindChildWindows() 92 93 Returns: 94 The window that was found 95 """ 96 return FindChildWindows(hwnd, path)[0] 97 98 99 def ScrapeWindow(hwnd, rect=None): 100 """Scrape a visible window and return its contents as a bitmap. 101 102 Args: 103 hwnd: handle of the window to scrape 104 rect: rectangle to scrape in client coords, defaults to the whole thing 105 If specified, it's a 4-tuple of (left, top, right, bottom) 106 107 Returns: 108 An Image containing the scraped data 109 """ 110 # Activate the window 111 SetForegroundWindow(hwnd) 112 113 # If no rectangle was specified, use the fill client rectangle 114 if not rect: rect = win32gui.GetClientRect(hwnd) 115 116 upper_left = win32gui.ClientToScreen(hwnd, (rect[0], rect[1])) 117 lower_right = win32gui.ClientToScreen(hwnd, (rect[2], rect[3])) 118 rect = upper_left+lower_right 119 120 return PIL.ImageGrab.grab(rect) 121 122 123 def SetForegroundWindow(hwnd): 124 """Bring a window to the foreground.""" 125 win32gui.SetForegroundWindow(hwnd) 126 127 128 def InvokeAndWait(path, cmdline="", timeout=10, tick=1.): 129 """Invoke an application and wait for it to bring up a window. 130 131 Args: 132 path: full path to the executable to invoke 133 cmdline: command line to pass to executable 134 timeout: how long (in seconds) to wait before giving up 135 tick: length of time to wait between checks 136 137 Returns: 138 A tuple of handles to the process and the application's window, 139 or (None, None) if it timed out waiting for the process 140 """ 141 142 def EnumWindowProc(hwnd, ret): 143 """Internal enumeration func, checks for visibility and proper PID.""" 144 if win32gui.IsWindowVisible(hwnd): # don't bother even checking hidden wnds 145 pid = win32process.GetWindowThreadProcessId(hwnd)[1] 146 if pid == ret[0]: 147 ret[1] = hwnd 148 return 0 # 0 means stop enumeration 149 return 1 # 1 means continue enumeration 150 151 # We don't need to change anything about the startupinfo structure 152 # (the default is quite sufficient) but we need to create it just the 153 # same. 154 sinfo = win32process.STARTUPINFO() 155 156 proc = win32process.CreateProcess( 157 path, # path to new process's executable 158 cmdline, # application's command line 159 None, # process security attributes (default) 160 None, # thread security attributes (default) 161 False, # inherit parent's handles 162 0, # creation flags 163 None, # environment variables 164 None, # directory 165 sinfo) # default startup info 166 167 # Create process returns (prochandle, pid, threadhandle, tid). At 168 # some point we may care about the other members, but for now, all 169 # we're after is the pid 170 pid = proc[2] 171 172 # Enumeration APIs can take an arbitrary integer, usually a pointer, 173 # to be passed to the enumeration function. We'll pass a pointer to 174 # a structure containing the PID we're looking for, and an empty out 175 # parameter to hold the found window ID 176 ret = [pid, None] 177 178 tries_until_timeout = timeout/tick 179 num_tries = 0 180 181 # Enumerate top-level windows, look for one with our PID 182 while num_tries < tries_until_timeout and ret[1] is None: 183 try: 184 win32gui.EnumWindows(EnumWindowProc, ret) 185 except pywintypes.error, e: 186 # error 0 isn't an error, it just meant the enumeration was 187 # terminated early 188 if e[0]: raise e 189 190 time.sleep(tick) 191 num_tries += 1 192 193 # TODO(jhaas): Should we throw an exception if we timeout? Or is returning 194 # a window ID of None sufficient? 195 return (proc[0], ret[1]) 196 197 198 def WaitForProcessExit(proc, timeout=None): 199 """Waits for a given process to terminate. 200 201 Args: 202 proc: handle to process 203 timeout: timeout (in seconds). None = wait indefinitely 204 205 Returns: 206 True if process ended, False if timed out 207 """ 208 if timeout is None: 209 timeout = win32event.INFINITE 210 else: 211 # convert sec to msec 212 timeout *= 1000 213 214 return (win32event.WaitForSingleObject(proc, timeout) == 215 win32event.WAIT_OBJECT_0) 216 217 218 def WaitForThrobber(hwnd, rect=None, timeout=20, tick=0.1, done=10): 219 """Wait for a browser's "throbber" (loading animation) to complete. 220 221 Args: 222 hwnd: window containing the throbber 223 rect: rectangle of the throbber, in client coords. If None, whole window 224 timeout: if the throbber is still throbbing after this long, give up 225 tick: how often to check the throbber 226 done: how long the throbber must be unmoving to be considered done 227 228 Returns: 229 Number of seconds waited, -1 if timed out 230 """ 231 if not rect: rect = win32gui.GetClientRect(hwnd) 232 233 # last_throbber will hold the results of the preceding scrape; 234 # we'll compare it against the current scrape to see if we're throbbing 235 last_throbber = ScrapeWindow(hwnd, rect) 236 start_clock = time.clock() 237 timeout_clock = start_clock + timeout 238 last_changed_clock = start_clock; 239 240 while time.clock() < timeout_clock: 241 time.sleep(tick) 242 243 current_throbber = ScrapeWindow(hwnd, rect) 244 if current_throbber.tostring() != last_throbber.tostring(): 245 last_throbber = current_throbber 246 last_changed_clock = time.clock() 247 else: 248 if time.clock() - last_changed_clock > done: 249 return last_changed_clock - start_clock 250 251 return -1 252 253 254 def MoveAndSizeWindow(wnd, position=None, size=None, child=None): 255 """Moves and/or resizes a window. 256 257 Repositions and resizes a window. If a child window is provided, 258 the parent window is resized so the child window has the given size 259 260 Args: 261 wnd: handle of the frame window 262 position: new location for the frame window 263 size: new size for the frame window (or the child window) 264 child: handle of the child window 265 266 Returns: 267 None 268 """ 269 rect = win32gui.GetWindowRect(wnd) 270 271 if position is None: position = (rect[0], rect[1]) 272 if size is None: 273 size = (rect[2]-rect[0], rect[3]-rect[1]) 274 elif child is not None: 275 child_rect = win32gui.GetWindowRect(child) 276 slop = (rect[2]-rect[0]-child_rect[2]+child_rect[0], 277 rect[3]-rect[1]-child_rect[3]+child_rect[1]) 278 size = (size[0]+slop[0], size[1]+slop[1]) 279 280 win32gui.MoveWindow(wnd, # window to move 281 position[0], # new x coord 282 position[1], # new y coord 283 size[0], # new width 284 size[1], # new height 285 True) # repaint? 286 287 288 def EndProcess(proc, code=0): 289 """Ends a process. 290 291 Wraps the OS TerminateProcess call for platform-independence 292 293 Args: 294 proc: process ID 295 code: process exit code 296 297 Returns: 298 None 299 """ 300 win32process.TerminateProcess(proc, code) 301 302 303 def URLtoFilename(url, path=None, extension=None): 304 """Converts a URL to a filename, given a path. 305 306 This in theory could cause collisions if two URLs differ only 307 in unprintable characters (eg. http://www.foo.com/?bar and 308 http://www.foo.com/:bar. In practice this shouldn't be a problem. 309 310 Args: 311 url: The URL to convert 312 path: path to the directory to store the file 313 extension: string to append to filename 314 315 Returns: 316 filename 317 """ 318 trans = string.maketrans(r'\/:*?"<>|', '_________') 319 320 if path is None: path = "" 321 if extension is None: extension = "" 322 if len(path) > 0 and path[-1] != '\\': path += '\\' 323 url = url.translate(trans) 324 return "%s%s%s" % (path, url, extension) 325 326 327 def PreparePath(path): 328 """Ensures that a given path exists, making subdirectories if necessary. 329 330 Args: 331 path: fully-qualified path of directory to ensure exists 332 333 Returns: 334 None 335 """ 336 try: 337 os.makedirs(path) 338 except OSError, e: 339 if e[0] != 17: raise e # error 17: path already exists 340 341 342 def main(): 343 PreparePath(r"c:\sitecompare\scrapes\ie7") 344 # We're being invoked rather than imported. Let's do some tests 345 346 # Hardcode IE's location for the purpose of this test 347 (proc, wnd) = InvokeAndWait( 348 r"c:\program files\internet explorer\iexplore.exe") 349 350 # Find the browser pane in the IE window 351 browser = FindChildWindow( 352 wnd, "TabWindowClass/Shell DocObject View/Internet Explorer_Server") 353 354 # Move and size the window 355 MoveAndSizeWindow(wnd, (0, 0), (1024, 768), browser) 356 357 # Take a screenshot 358 i = ScrapeWindow(browser) 359 360 i.show() 361 362 EndProcess(proc, 0) 363 364 365 if __name__ == "__main__": 366 sys.exit(main()) 367
