Mercurial > ~astiob > upreckon > hgweb

comparison upreckon/win32.py @ 146:d5b6708c1955

Find changesets by keywords (author, files, the commit message), revision number or hash, or revset expression.
Distutils support, reorganization and cleaning up * Removed command-line options -t and -u. * Reorganized code: o all modules are now in package upreckon; o TestCaseNotPassed and its descendants now live in a separate module exceptions; o load_problem now lives in module problem. * Commented out mentions of command-line option -c in --help. * Added a distutils-based setup.py.
author Oleg Oshmyan <chortos@inbox.lv>
date Sat, 28 May 2011 14:24:25 +0100
parents win32.py@98bccf81db4d
children aa343ff41c27
comparison
equal deleted inserted replaced
145:d2c266c8d820 146:d5b6708c1955
1 # Copyright (c) 2010-2011 Chortos-2 <chortos@inbox.lv>
2
3 from __future__ import division, with_statement
4
5 from .compat import *
6 from .exceptions import *
7
8 from ctypes import *
9 from ctypes.wintypes import *
10 from msvcrt import getch as pause
11 from time import clock
12 import os, subprocess, sys, time
13
14 try:
15 from _winreg import *
16 except ImportError:
17 from winreg import *
18
19 # Defaults that may be overwritten by values from _subprocess
20 INFINITE = -1
21 STD_INPUT_HANDLE = -10
22 WAIT_OBJECT_0 = 0
23
24 try:
25 from _subprocess import *
26 except ImportError:
27 pass
28
29 try:
30 from numbers import Integral
31 except ImportError:
32 Integral = int, long
33
34 try:
35 from collections import namedtuple
36 except ImportError:
37 from operator import itemgetter
38 class ProcessTimes(tuple):
39 __slots__ = ()
40 def __new__(cls, creation, exit, kernel, user):
41 return tuple.__new__(cls, (creation, exit, kernel, user))
42 __getnewargs__ = lambda self: tuple(self)
43 creation, exit, kernel, user = map(property, map(itemgetter, range(4)))
44 else:
45 ProcessTimes = namedtuple('ProcessTimes', 'creation exit kernel user')
46
47 __all__ = 'call', 'kill', 'pause', 'clock'
48
49
50 from functools import wraps
51 pathext = [''] + os.environ['PATHEXT'].split(';')
52 @wraps(subprocess.Popen)
53 def Popen(cmdline, *args, **kwargs):
54 try:
55 return subprocess.Popen(cmdline, *args, **kwargs)
56 except WindowsError:
57 for ext in pathext:
58 path = cmdline[0] + ext
59 newcmdline = type(cmdline)((path,)) + cmdline[1:]
60 try:
61 return subprocess.Popen(newcmdline, *args, **kwargs)
62 except WindowsError:
63 pass
64 for branch in HKEY_CURRENT_USER, HKEY_LOCAL_MACHINE:
65 try:
66 path = (R'SOFTWARE\Microsoft\Windows\CurrentVersion'
67 R'\App Paths\%s%s' % (cmdline[0], ext))
68 path = QueryValue(branch, path)
69 break
70 except WindowsError:
71 pass
72 else:
73 continue
74 if path[0] == '"' == path[-1]:
75 path = path[1:-1]
76 newcmdline = type(cmdline)((path,)) + cmdline[1:]
77 try:
78 return subprocess.Popen(newcmdline, *args, **kwargs)
79 except WindowsError:
80 pass
81 # I'd like to transparently re-raise the exception generated
82 # on the very first try, but syntax differences preclude me from
83 # doing so in Python 2 and it can't be done at all in Python 3
84 raise
85
86
87 # Automatically convert _subprocess handle objects into low-level
88 # HANDLEs and replicate their functionality for our own use
89 try:
90 _subprocess_handle = type(GetCurrentProcess())
91 except NameError:
92 _subprocess_handle = Integral
93 class Handle(object):
94 @staticmethod
95 def from_param(handle):
96 if isinstance(handle, (_subprocess_handle, Integral)):
97 return HANDLE(int(handle))
98 elif isinstance(handle, Handle):
99 return HANDLE(handle.handle)
100 elif isinstance(handle, HANDLE):
101 return handle
102 else:
103 raise TypeError('cannot convert %s to a handle' %
104 type(handle).__name__)
105
106 __slots__ = 'handle'
107
108 def __init__(self, handle):
109 if isinstance(handle, Integral):
110 self.handle = handle
111 elif isinstance(handle, HANDLE):
112 self.handle = handle.value
113 elif isinstance(handle, Handle):
114 self.handle = handle.handle
115 elif isinstance(handle, _subprocess_handle):
116 handle = HANDLE(int(handle))
117 flags = DWORD()
118 try:
119 if windll.kernel32.GetHandleInformation(handle, byref(flags)):
120 flags = flags.value
121 else:
122 flags = 0
123 except AttributeError:
124 # Available on NT 3.51 and up, NT line only
125 flags = 0
126 proc = HANDLE(int(GetCurrentProcess()))
127 handle = DuplicateHandle(proc, handle, proc, 0, flags & 1, 2)
128 self.handle = handle.Detach()
129 else:
130 raise TypeError("Handle() argument must be a handle, not '%s'" %
131 type(handle).__name__)
132
133 def __int__(self):
134 return int(self.handle)
135
136 def Detach(self):
137 handle = self.handle
138 self.handle = None
139 return handle
140
141 # This is also __del__, so only locals are accessed
142 def Close(self, _CloseHandle=windll.kernel32.CloseHandle, _HANDLE=HANDLE):
143 if getattr(self, 'handle', None):
144 _CloseHandle(_HANDLE(self.handle))
145 self.handle = None
146 __del__ = Close
147
148 CHAR = c_char
149 INVALID_HANDLE_VALUE = HANDLE(-1).value
150 LPDWORD = POINTER(DWORD)
151 LPFILETIME = POINTER(FILETIME)
152 SIZE_T = ULONG_PTR = WPARAM
153 ULONGLONG = c_ulonglong
154
155 try:
156 unicode
157 except NameError:
158 LPCTSTR = LPCWSTR
159 UNISUFFIX = 'W'
160 else:
161 LPCTSTR = LPCSTR
162 UNISUFFIX = 'A'
163
164
165 prototype = WINFUNCTYPE(BOOL, Handle,
166 LPFILETIME, LPFILETIME, LPFILETIME, LPFILETIME)
167 flags = ((1, 'process'),
168 (2, 'creation'), (2, 'exit'), (2, 'kernel'), (2, 'user'))
169 try:
170 GetProcessTimes = prototype(('GetProcessTimes', windll.kernel32), flags)
171 except AttributeError:
172 # Available on NT 3.5 and up, NT line only
173 GetProcessTimes = None
174 else:
175 def errcheck(result, func, args):
176 if not result: raise WinError()
177 times = ((t.dwHighDateTime << 32 | t.dwLowDateTime) / 10000000
178 for t in args[1:])
179 return ProcessTimes(*times)
180 GetProcessTimes.errcheck = errcheck
181
182
183 class PROCESS_MEMORY_COUNTERS(Structure):
184 _fields_ = (('cb', DWORD),
185 ('PageFaultCount', DWORD),
186 ('PeakWorkingSetSize', SIZE_T),
187 ('WorkingSetSize', SIZE_T),
188 ('QuotaPeakPagedPoolUsage', SIZE_T),
189 ('QuotaPagedPoolUsage', SIZE_T),
190 ('QuotaPeakNonPagedPoolUsage', SIZE_T),
191 ('QuotaNonPagedPoolUsage', SIZE_T),
192 ('PagefileUsage', SIZE_T),
193 ('PeakPagefileUsage', SIZE_T))
194
195 prototype = WINFUNCTYPE(BOOL, Handle, POINTER(PROCESS_MEMORY_COUNTERS), DWORD)
196 flags = ((1, 'process'), (2, 'counters'),
197 (5, 'cb', sizeof(PROCESS_MEMORY_COUNTERS)))
198 try:
199 GetProcessMemoryInfo = prototype(('GetProcessMemoryInfo', windll.psapi),
200 flags)
201 except AttributeError:
202 # Available on NT 4.0 and up, NT line only
203 GetProcessMemoryInfo = None
204 else:
205 def errcheck(result, func, args):
206 if not result: raise WinError()
207 return args
208 GetProcessMemoryInfo.errcheck = errcheck
209
210
211 class _uChar_union(Union):
212 _fields_ = (('UnicodeChar', WCHAR),
213 ('AsciiChar', CHAR))
214
215 class KEY_EVENT_RECORD(Structure):
216 _fields_ = (('bKeyDown', BOOL),
217 ('wRepeatCount', WORD),
218 ('wVirtualKeyCode', WORD),
219 ('wVirtualScanCode', WORD),
220 ('uChar', _uChar_union),
221 ('dwControlKeyState', DWORD))
222
223 RIGHT_ALT_PRESSED = 0x001
224 LEFT_ALT_PRESSED = 0x002
225 RIGHT_CTRL_PRESSED = 0x004
226 LEFT_CTRL_PRESSED = 0x008
227 SHIFT_PRESSED = 0x010
228 NUMLOCK_ON = 0x020
229 SCROLLLOCK_ON = 0x040
230 CAPSLOCK_ON = 0x080
231 ENHANCED_KEY = 0x100
232
233 class _Event_union(Union):
234 _fields_ = ('KeyEvent', KEY_EVENT_RECORD),
235
236 class INPUT_RECORD(Structure):
237 _fields_ = (('EventType', WORD),
238 ('Event', _Event_union))
239
240 KEY_EVENT = 0x01
241 MOUSE_EVENT = 0x02
242 WINDOW_BUFFER_SIZE_EVENT = 0x04
243 MENU_EVENT = 0x08
244 FOCUS_EVENT = 0x10
245
246 prototype = WINFUNCTYPE(BOOL, Handle, POINTER(INPUT_RECORD), DWORD, LPDWORD)
247 flags = (1, 'input'), (2, 'buffer'), (5, 'length', 1), (2, 'number_read')
248 ReadConsoleInput = prototype(('ReadConsoleInputA', windll.kernel32), flags)
249 def errcheck(result, func, args):
250 if not result: raise WinError()
251 return args[1] if args[3] else None
252 ReadConsoleInput.errcheck = errcheck
253
254
255 prototype = WINFUNCTYPE(BOOL, Handle)
256 flags = (1, 'input'),
257 FlushConsoleInputBuffer = prototype(('FlushConsoleInputBuffer',
258 windll.kernel32), flags)
259 def errcheck(result, func, args):
260 if not result: raise WinError()
261 FlushConsoleInputBuffer.errcheck = errcheck
262
263
264 prototype = WINFUNCTYPE(BOOL, Handle, DWORD)
265 flags = (1, 'console'), (1, 'mode')
266 SetConsoleMode = prototype(('SetConsoleMode', windll.kernel32), flags)
267 def errcheck(result, func, args):
268 if not result: raise WinError()
269 SetConsoleMode.errcheck = errcheck
270
271 ENABLE_PROCESSED_INPUT = 0x001
272 ENABLE_LINE_INPUT = 0x002
273 ENABLE_ECHO_INPUT = 0x004
274 ENABLE_WINDOW_INPUT = 0x008
275 ENABLE_MOUSE_INPUT = 0x010
276 ENABLE_INSERT_MODE = 0x020
277 ENABLE_QUICK_EDIT_MODE = 0x040
278 ENABLE_EXTENDED_FLAGS = 0x080
279
280 ENABLE_PROCESSED_OUTPUT = 1
281 ENABLE_WRAP_AT_EOL_OUTPUT = 2
282
283
284 prototype = WINFUNCTYPE(HANDLE, c_void_p, LPCTSTR)
285 flags = (5, 'attributes'), (1, 'name')
286 try:
287 CreateJobObject = prototype(('CreateJobObject'+UNISUFFIX, windll.kernel32),
288 flags)
289 except AttributeError:
290 # Available on 2000 and up, NT line only
291 CreateJobObject = lambda name: None
292 else:
293 def errcheck(result, func, args):
294 if not result: raise WinError()
295 return Handle(result)
296 CreateJobObject.errcheck = errcheck
297
298
299 prototype = WINFUNCTYPE(BOOL, Handle, Handle)
300 flags = (1, 'job'), (1, 'handle')
301 try:
302 AssignProcessToJobObject = prototype(('AssignProcessToJobObject',
303 windll.kernel32), flags)
304 except AttributeError:
305 # Available on 2000 and up, NT line only
306 AssignProcessToJobObject = lambda job, handle: None
307 else:
308 def errcheck(result, func, args):
309 if not result: raise WinError()
310 AssignProcessToJobObject.errcheck = errcheck
311
312
313 class JOBOBJECT_BASIC_LIMIT_INFORMATION(Structure):
314 _fields_ = (('PerProcessUserTimeLimit', LARGE_INTEGER),
315 ('PerJobUserTimeLimit', LARGE_INTEGER),
316 ('LimitFlags', DWORD),
317 ('MinimumWorkingSetSize', SIZE_T),
318 ('MaximumWorkingSetSize', SIZE_T),
319 ('ActiveProcessLimit', DWORD),
320 ('Affinity', ULONG_PTR),
321 ('PriorityClass', DWORD),
322 ('SchedulingClass', DWORD))
323
324 JOB_OBJECT_LIMIT_WORKINGSET = 0x0001
325 JOB_OBJECT_LIMIT_PROCESS_TIME = 0x0002
326 JOB_OBJECT_LIMIT_JOB_TIME = 0x0004
327 JOB_OBJECT_LIMIT_ACTIVE_PROCESS = 0x0008
328 JOB_OBJECT_LIMIT_AFFINITY = 0x0010
329 JOB_OBJECT_LIMIT_PRIORITY_CLASS = 0x0020
330 JOB_OBJECT_LIMIT_PRESERVE_JOB_TIME = 0x0040
331 JOB_OBJECT_LIMIT_SCHEDULING_CLASS = 0x0080
332 JOB_OBJECT_LIMIT_PROCESS_MEMORY = 0x0100
333 JOB_OBJECT_LIMIT_JOB_MEMORY = 0x0200
334 JOB_OBJECT_LIMIT_DIE_ON_UNHANDLED_EXCEPTION = 0x0400
335 JOB_OBJECT_LIMIT_BREAKAWAY_OK = 0x0800
336 JOB_OBJECT_LIMIT_SILENT_BREAKAWAY_OK = 0x1000
337 JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE = 0x2000
338 JOB_OBJECT_LIMIT_SUBSET_AFFINITY = 0x4000
339
340 class IO_COUNTERS(Structure):
341 _fields_ = (('ReadOperationCount', ULONGLONG),
342 ('WriteOperationCount', ULONGLONG),
343 ('OtherOperationCount', ULONGLONG),
344 ('ReadTransferCount', ULONGLONG),
345 ('WriteTransferCount', ULONGLONG),
346 ('OtherTransferCount', ULONGLONG))
347
348 class JOBOBJECT_EXTENDED_LIMIT_INFORMATION(Structure):
349 _fields_ = (('BasicLimitInformation', JOBOBJECT_BASIC_LIMIT_INFORMATION),
350 ('IoInfo', IO_COUNTERS),
351 ('ProcessMemoryLimit', SIZE_T),
352 ('JobMemoryLimit', SIZE_T),
353 ('PeakProcessMemoryUsed', SIZE_T),
354 ('PeakJobMemoryUsed', SIZE_T))
355
356 prototype = WINFUNCTYPE(BOOL, Handle, c_int, c_void_p, DWORD)
357 flags = (1, 'job'), (1, 'infoclass'), (1, 'info'), (1, 'infosize')
358 try:
359 _setjobinfo = prototype(('SetInformationJobObject',windll.kernel32), flags)
360 except AttributeError:
361 # Available on 2000 and up, NT line only
362 SetInformationJobObject = lambda job, infoclass, info: None
363 else:
364 def errcheck(result, func, args):
365 if not result: raise WinError()
366 _setjobinfo.errcheck = errcheck
367 def SetInformationJobObject(job, infoclass, info):
368 return _setjobinfo(job, infoclass, byref(info), sizeof(info))
369
370 (
371 JobObjectBasicAccountingInformation,
372 JobObjectBasicLimitInformation,
373 JobObjectBasicProcessIdList,
374 JobObjectBasicUIRestrictions,
375 JobObjectSecurityLimitInformation,
376 JobObjectEndOfJobTimeInformation,
377 JobObjectAssociateCompletionPortInformation,
378 JobObjectBasicAndIoAccountingInformation,
379 JobObjectExtendedLimitInformation,
380 JobObjectJobSetInformation,
381 MaxJobObjectInfoClass
382 ) = range(1, 12)
383
384
385 prototype = WINFUNCTYPE(DWORD, DWORD, POINTER(HANDLE), BOOL, DWORD)
386 flags = (1, 'count'), (1, 'handles'), (1, 'wait_all'), (1, 'milliseconds')
387 _wait_multiple = prototype(('WaitForMultipleObjects', windll.kernel32), flags)
388 def errcheck(result, func, args):
389 if result == WAIT_FAILED: raise WinError()
390 return args
391 _wait_multiple.errcheck = errcheck
392 def WaitForMultipleObjects(handles, wait_all, timeout):
393 n = len(handles)
394 handles = (Handle.from_param(handle) for handle in handles)
395 timeout = ceil(timeout * 1000)
396 return _wait_multiple(n, (HANDLE * n)(*handles), wait_all, timeout)
397
398 # WAIT_OBJECT_0 defined at the top of the file
399 WAIT_ABANDONED_0 = 0x00000080
400 WAIT_TIMEOUT = 0x00000102
401 WAIT_FAILED = 0xFFFFFFFF
402
403
404 try:
405 _wait_single = WaitForSingleObject
406 except NameError:
407 prototype = WINFUNCTYPE(DWORD, Handle, DWORD)
408 flags = (1, 'handle'), (1, 'milliseconds')
409 _wait_single = prototype(('WaitForSingleObject', windll.kernel32), flags)
410 def errcheck(result, func, args):
411 if result == WAIT_FAILED: raise WinError()
412 return args
413 _wait_single.errcheck = errcheck
414 def WaitForSingleObject(handle, timeout):
415 return _wait_single(handle, ceil(timeout * 1000))
416
417
418 try:
419 GetStdHandle
420 except NameError:
421 prototype = WINFUNCTYPE(HANDLE, DWORD)
422 flags = (1, 'which'),
423 GetStdHandle = prototype(('GetStdHandle', windll.kernel32), flags)
424 def errcheck(result, func, args):
425 if result == INVALID_HANDLE_VALUE: raise WinError()
426 return args if result else None
427 GetStdHandle.errcheck = errcheck
428
429
430 try:
431 TerminateProcess
432 except NameError:
433 prototype = WINFUNCTYPE(BOOL, Handle, UINT)
434 flags = (1, 'process'), (1, 'exitcode')
435 TerminateProcess = prototype(('TerminateProcess', windll.kernel32), flags)
436 def errcheck(result, func, args):
437 if not result: raise WinError()
438 TerminateProcess.errcheck = errcheck
439
440
441 # Do not show error messages due to errors in the program being tested
442 try:
443 errmode = windll.kernel32.GetErrorMode()
444 except AttributeError:
445 # GetErrorMode is available on Vista/2008 and up
446 errmode = windll.kernel32.SetErrorMode(0)
447 windll.kernel32.SetErrorMode(errmode | 0x8003)
448
449 stdin = GetStdHandle(STD_INPUT_HANDLE)
450 try:
451 SetConsoleMode(stdin, ENABLE_PROCESSED_INPUT)
452 except WindowsError:
453 console_input = False
454 else:
455 console_input = True
456 FlushConsoleInputBuffer(stdin)
457
458 def call(*args, **kwargs):
459 case = kwargs.pop('case')
460 job = CreateJobObject(None)
461 flags = 0
462 if case.maxcputime:
463 flags |= JOB_OBJECT_LIMIT_PROCESS_TIME
464 if case.maxmemory:
465 flags |= JOB_OBJECT_LIMIT_PROCESS_MEMORY
466 limits = JOBOBJECT_EXTENDED_LIMIT_INFORMATION(
467 JOBOBJECT_BASIC_LIMIT_INFORMATION(
468 PerProcessUserTimeLimit=ceil((case.maxcputime or 0)*10000000),
469 LimitFlags=flags,
470 ),
471 ProcessMemoryLimit=ceil((case.maxmemory or 0)*1048576),
472 )
473 SetInformationJobObject(job, JobObjectExtendedLimitInformation, limits)
474 try:
475 case.process = Popen(*args, **kwargs)
476 except OSError:
477 raise CannotStartTestee(sys.exc_info()[1])
478 case.time_started = clock()
479 AssignProcessToJobObject(job, case.process._handle)
480 if not console_input:
481 if case.maxwalltime:
482 if (WaitForSingleObject(case.process._handle, case.maxwalltime) !=
483 WAIT_OBJECT_0):
484 raise WallTimeLimitExceeded
485 else:
486 case.process.wait()
487 else:
488 handles = case.process._handle, stdin
489 if case.maxwalltime:
490 time_end = case.time_started + case.maxwalltime
491 while case.process.poll() is None:
492 remaining = time_end - clock()
493 if remaining > 0:
494 if (WaitForMultipleObjects(handles, False, remaining) ==
495 WAIT_OBJECT_0 + 1):
496 ir = ReadConsoleInput(stdin)
497 if (ir and
498 ir.EventType == KEY_EVENT and
499 ir.Event.KeyEvent.bKeyDown and
500 ir.Event.KeyEvent.wVirtualKeyCode == 27):
501 raise CanceledByUser
502 else:
503 raise WallTimeLimitExceeded
504 else:
505 while case.process.poll() is None:
506 if (WaitForMultipleObjects(handles, False, INFINITE) ==
507 WAIT_OBJECT_0 + 1):
508 ir = ReadConsoleInput(stdin)
509 if (ir and
510 ir.EventType == KEY_EVENT and
511 ir.Event.KeyEvent.bKeyDown and
512 ir.Event.KeyEvent.wVirtualKeyCode == 27):
513 raise CanceledByUser
514 case.time_stopped = clock()
515 if GetProcessTimes:
516 try:
517 times = GetProcessTimes(case.process._handle)
518 except WindowsError:
519 pass
520 else:
521 if case.maxcputime or not case.maxwalltime:
522 cputime = times.kernel + times.user
523 case.time_stopped = cputime
524 case.time_started = 0
525 case.time_limit_string = case.cpu_time_limit_string
526 if case.maxcputime and cputime > case.maxcputime:
527 raise CPUTimeLimitExceeded
528 else:
529 case.time_stopped = times.exit
530 case.time_started = times.creation
531 walltime = times.exit - times.creation
532 if case.maxwalltime and walltime > case.maxwalltime:
533 raise WallTimeLimitExceeded
534 if case.maxcputime and case.process.returncode == 1816:
535 raise CPUTimeLimitExceeded
536 if case.maxmemory and GetProcessMemoryInfo:
537 try:
538 counters = GetProcessMemoryInfo(case.process._handle)
539 except WindowsError:
540 pass
541 else:
542 if counters.PeakPagefileUsage > case.maxmemory * 1048576:
543 raise MemoryLimitExceeded
544
545
546 def kill(process):
547 # Give up after three attempts
548 for i in range(3):
549 try:
550 try:
551 process.terminate()
552 except AttributeError:
553 TerminateProcess(process._handle, 1)
554 except WindowsError:
555 time.sleep(0)
556 else:
557 break