Mercurial > ~astiob > upreckon > hgweb
diff upreckon/win32.py @ 146:d5b6708c1955
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 |
line wrap: on
line diff
--- /dev/null Thu Jan 01 00:00:00 1970 +0000 +++ b/upreckon/win32.py Sat May 28 14:24:25 2011 +0100 @@ -0,0 +1,557 @@ +# Copyright (c) 2010-2011 Chortos-2 <chortos@inbox.lv> + +from __future__ import division, with_statement + +from .compat import * +from .exceptions import * + +from ctypes import * +from ctypes.wintypes import * +from msvcrt import getch as pause +from time import clock +import os, subprocess, sys, time + +try: + from _winreg import * +except ImportError: + from winreg import * + +# Defaults that may be overwritten by values from _subprocess +INFINITE = -1 +STD_INPUT_HANDLE = -10 +WAIT_OBJECT_0 = 0 + +try: + from _subprocess import * +except ImportError: + pass + +try: + from numbers import Integral +except ImportError: + Integral = int, long + +try: + from collections import namedtuple +except ImportError: + from operator import itemgetter + class ProcessTimes(tuple): + __slots__ = () + def __new__(cls, creation, exit, kernel, user): + return tuple.__new__(cls, (creation, exit, kernel, user)) + __getnewargs__ = lambda self: tuple(self) + creation, exit, kernel, user = map(property, map(itemgetter, range(4))) +else: + ProcessTimes = namedtuple('ProcessTimes', 'creation exit kernel user') + +__all__ = 'call', 'kill', 'pause', 'clock' + + +from functools import wraps +pathext = [''] + os.environ['PATHEXT'].split(';') +@wraps(subprocess.Popen) +def Popen(cmdline, *args, **kwargs): + try: + return subprocess.Popen(cmdline, *args, **kwargs) + except WindowsError: + for ext in pathext: + path = cmdline[0] + ext + newcmdline = type(cmdline)((path,)) + cmdline[1:] + try: + return subprocess.Popen(newcmdline, *args, **kwargs) + except WindowsError: + pass + for branch in HKEY_CURRENT_USER, HKEY_LOCAL_MACHINE: + try: + path = (R'SOFTWARE\Microsoft\Windows\CurrentVersion' + R'\App Paths\%s%s' % (cmdline[0], ext)) + path = QueryValue(branch, path) + break + except WindowsError: + pass + else: + continue + if path[0] == '"' == path[-1]: + path = path[1:-1] + newcmdline = type(cmdline)((path,)) + cmdline[1:] + try: + return subprocess.Popen(newcmdline, *args, **kwargs) + except WindowsError: + pass + # I'd like to transparently re-raise the exception generated + # on the very first try, but syntax differences preclude me from + # doing so in Python 2 and it can't be done at all in Python 3 + raise + + +# Automatically convert _subprocess handle objects into low-level +# HANDLEs and replicate their functionality for our own use +try: + _subprocess_handle = type(GetCurrentProcess()) +except NameError: + _subprocess_handle = Integral +class Handle(object): + @staticmethod + def from_param(handle): + if isinstance(handle, (_subprocess_handle, Integral)): + return HANDLE(int(handle)) + elif isinstance(handle, Handle): + return HANDLE(handle.handle) + elif isinstance(handle, HANDLE): + return handle + else: + raise TypeError('cannot convert %s to a handle' % + type(handle).__name__) + + __slots__ = 'handle' + + def __init__(self, handle): + if isinstance(handle, Integral): + self.handle = handle + elif isinstance(handle, HANDLE): + self.handle = handle.value + elif isinstance(handle, Handle): + self.handle = handle.handle + elif isinstance(handle, _subprocess_handle): + handle = HANDLE(int(handle)) + flags = DWORD() + try: + if windll.kernel32.GetHandleInformation(handle, byref(flags)): + flags = flags.value + else: + flags = 0 + except AttributeError: + # Available on NT 3.51 and up, NT line only + flags = 0 + proc = HANDLE(int(GetCurrentProcess())) + handle = DuplicateHandle(proc, handle, proc, 0, flags & 1, 2) + self.handle = handle.Detach() + else: + raise TypeError("Handle() argument must be a handle, not '%s'" % + type(handle).__name__) + + def __int__(self): + return int(self.handle) + + def Detach(self): + handle = self.handle + self.handle = None + return handle + + # This is also __del__, so only locals are accessed + def Close(self, _CloseHandle=windll.kernel32.CloseHandle, _HANDLE=HANDLE): + if getattr(self, 'handle', None): + _CloseHandle(_HANDLE(self.handle)) + self.handle = None + __del__ = Close + +CHAR = c_char +INVALID_HANDLE_VALUE = HANDLE(-1).value +LPDWORD = POINTER(DWORD) +LPFILETIME = POINTER(FILETIME) +SIZE_T = ULONG_PTR = WPARAM +ULONGLONG = c_ulonglong + +try: + unicode +except NameError: + LPCTSTR = LPCWSTR + UNISUFFIX = 'W' +else: + LPCTSTR = LPCSTR + UNISUFFIX = 'A' + + +prototype = WINFUNCTYPE(BOOL, Handle, + LPFILETIME, LPFILETIME, LPFILETIME, LPFILETIME) +flags = ((1, 'process'), + (2, 'creation'), (2, 'exit'), (2, 'kernel'), (2, 'user')) +try: + GetProcessTimes = prototype(('GetProcessTimes', windll.kernel32), flags) +except AttributeError: + # Available on NT 3.5 and up, NT line only + GetProcessTimes = None +else: + def errcheck(result, func, args): + if not result: raise WinError() + times = ((t.dwHighDateTime << 32 | t.dwLowDateTime) / 10000000 + for t in args[1:]) + return ProcessTimes(*times) + GetProcessTimes.errcheck = errcheck + + +class PROCESS_MEMORY_COUNTERS(Structure): + _fields_ = (('cb', DWORD), + ('PageFaultCount', DWORD), + ('PeakWorkingSetSize', SIZE_T), + ('WorkingSetSize', SIZE_T), + ('QuotaPeakPagedPoolUsage', SIZE_T), + ('QuotaPagedPoolUsage', SIZE_T), + ('QuotaPeakNonPagedPoolUsage', SIZE_T), + ('QuotaNonPagedPoolUsage', SIZE_T), + ('PagefileUsage', SIZE_T), + ('PeakPagefileUsage', SIZE_T)) + +prototype = WINFUNCTYPE(BOOL, Handle, POINTER(PROCESS_MEMORY_COUNTERS), DWORD) +flags = ((1, 'process'), (2, 'counters'), + (5, 'cb', sizeof(PROCESS_MEMORY_COUNTERS))) +try: + GetProcessMemoryInfo = prototype(('GetProcessMemoryInfo', windll.psapi), + flags) +except AttributeError: + # Available on NT 4.0 and up, NT line only + GetProcessMemoryInfo = None +else: + def errcheck(result, func, args): + if not result: raise WinError() + return args + GetProcessMemoryInfo.errcheck = errcheck + + +class _uChar_union(Union): + _fields_ = (('UnicodeChar', WCHAR), + ('AsciiChar', CHAR)) + +class KEY_EVENT_RECORD(Structure): + _fields_ = (('bKeyDown', BOOL), + ('wRepeatCount', WORD), + ('wVirtualKeyCode', WORD), + ('wVirtualScanCode', WORD), + ('uChar', _uChar_union), + ('dwControlKeyState', DWORD)) + +RIGHT_ALT_PRESSED = 0x001 +LEFT_ALT_PRESSED = 0x002 +RIGHT_CTRL_PRESSED = 0x004 +LEFT_CTRL_PRESSED = 0x008 +SHIFT_PRESSED = 0x010 +NUMLOCK_ON = 0x020 +SCROLLLOCK_ON = 0x040 +CAPSLOCK_ON = 0x080 +ENHANCED_KEY = 0x100 + +class _Event_union(Union): + _fields_ = ('KeyEvent', KEY_EVENT_RECORD), + +class INPUT_RECORD(Structure): + _fields_ = (('EventType', WORD), + ('Event', _Event_union)) + +KEY_EVENT = 0x01 +MOUSE_EVENT = 0x02 +WINDOW_BUFFER_SIZE_EVENT = 0x04 +MENU_EVENT = 0x08 +FOCUS_EVENT = 0x10 + +prototype = WINFUNCTYPE(BOOL, Handle, POINTER(INPUT_RECORD), DWORD, LPDWORD) +flags = (1, 'input'), (2, 'buffer'), (5, 'length', 1), (2, 'number_read') +ReadConsoleInput = prototype(('ReadConsoleInputA', windll.kernel32), flags) +def errcheck(result, func, args): + if not result: raise WinError() + return args[1] if args[3] else None +ReadConsoleInput.errcheck = errcheck + + +prototype = WINFUNCTYPE(BOOL, Handle) +flags = (1, 'input'), +FlushConsoleInputBuffer = prototype(('FlushConsoleInputBuffer', + windll.kernel32), flags) +def errcheck(result, func, args): + if not result: raise WinError() +FlushConsoleInputBuffer.errcheck = errcheck + + +prototype = WINFUNCTYPE(BOOL, Handle, DWORD) +flags = (1, 'console'), (1, 'mode') +SetConsoleMode = prototype(('SetConsoleMode', windll.kernel32), flags) +def errcheck(result, func, args): + if not result: raise WinError() +SetConsoleMode.errcheck = errcheck + +ENABLE_PROCESSED_INPUT = 0x001 +ENABLE_LINE_INPUT = 0x002 +ENABLE_ECHO_INPUT = 0x004 +ENABLE_WINDOW_INPUT = 0x008 +ENABLE_MOUSE_INPUT = 0x010 +ENABLE_INSERT_MODE = 0x020 +ENABLE_QUICK_EDIT_MODE = 0x040 +ENABLE_EXTENDED_FLAGS = 0x080 + +ENABLE_PROCESSED_OUTPUT = 1 +ENABLE_WRAP_AT_EOL_OUTPUT = 2 + + +prototype = WINFUNCTYPE(HANDLE, c_void_p, LPCTSTR) +flags = (5, 'attributes'), (1, 'name') +try: + CreateJobObject = prototype(('CreateJobObject'+UNISUFFIX, windll.kernel32), + flags) +except AttributeError: + # Available on 2000 and up, NT line only + CreateJobObject = lambda name: None +else: + def errcheck(result, func, args): + if not result: raise WinError() + return Handle(result) + CreateJobObject.errcheck = errcheck + + +prototype = WINFUNCTYPE(BOOL, Handle, Handle) +flags = (1, 'job'), (1, 'handle') +try: + AssignProcessToJobObject = prototype(('AssignProcessToJobObject', + windll.kernel32), flags) +except AttributeError: + # Available on 2000 and up, NT line only + AssignProcessToJobObject = lambda job, handle: None +else: + def errcheck(result, func, args): + if not result: raise WinError() + AssignProcessToJobObject.errcheck = errcheck + + +class JOBOBJECT_BASIC_LIMIT_INFORMATION(Structure): + _fields_ = (('PerProcessUserTimeLimit', LARGE_INTEGER), + ('PerJobUserTimeLimit', LARGE_INTEGER), + ('LimitFlags', DWORD), + ('MinimumWorkingSetSize', SIZE_T), + ('MaximumWorkingSetSize', SIZE_T), + ('ActiveProcessLimit', DWORD), + ('Affinity', ULONG_PTR), + ('PriorityClass', DWORD), + ('SchedulingClass', DWORD)) + +JOB_OBJECT_LIMIT_WORKINGSET = 0x0001 +JOB_OBJECT_LIMIT_PROCESS_TIME = 0x0002 +JOB_OBJECT_LIMIT_JOB_TIME = 0x0004 +JOB_OBJECT_LIMIT_ACTIVE_PROCESS = 0x0008 +JOB_OBJECT_LIMIT_AFFINITY = 0x0010 +JOB_OBJECT_LIMIT_PRIORITY_CLASS = 0x0020 +JOB_OBJECT_LIMIT_PRESERVE_JOB_TIME = 0x0040 +JOB_OBJECT_LIMIT_SCHEDULING_CLASS = 0x0080 +JOB_OBJECT_LIMIT_PROCESS_MEMORY = 0x0100 +JOB_OBJECT_LIMIT_JOB_MEMORY = 0x0200 +JOB_OBJECT_LIMIT_DIE_ON_UNHANDLED_EXCEPTION = 0x0400 +JOB_OBJECT_LIMIT_BREAKAWAY_OK = 0x0800 +JOB_OBJECT_LIMIT_SILENT_BREAKAWAY_OK = 0x1000 +JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE = 0x2000 +JOB_OBJECT_LIMIT_SUBSET_AFFINITY = 0x4000 + +class IO_COUNTERS(Structure): + _fields_ = (('ReadOperationCount', ULONGLONG), + ('WriteOperationCount', ULONGLONG), + ('OtherOperationCount', ULONGLONG), + ('ReadTransferCount', ULONGLONG), + ('WriteTransferCount', ULONGLONG), + ('OtherTransferCount', ULONGLONG)) + +class JOBOBJECT_EXTENDED_LIMIT_INFORMATION(Structure): + _fields_ = (('BasicLimitInformation', JOBOBJECT_BASIC_LIMIT_INFORMATION), + ('IoInfo', IO_COUNTERS), + ('ProcessMemoryLimit', SIZE_T), + ('JobMemoryLimit', SIZE_T), + ('PeakProcessMemoryUsed', SIZE_T), + ('PeakJobMemoryUsed', SIZE_T)) + +prototype = WINFUNCTYPE(BOOL, Handle, c_int, c_void_p, DWORD) +flags = (1, 'job'), (1, 'infoclass'), (1, 'info'), (1, 'infosize') +try: + _setjobinfo = prototype(('SetInformationJobObject',windll.kernel32), flags) +except AttributeError: + # Available on 2000 and up, NT line only + SetInformationJobObject = lambda job, infoclass, info: None +else: + def errcheck(result, func, args): + if not result: raise WinError() + _setjobinfo.errcheck = errcheck + def SetInformationJobObject(job, infoclass, info): + return _setjobinfo(job, infoclass, byref(info), sizeof(info)) + +( + JobObjectBasicAccountingInformation, + JobObjectBasicLimitInformation, + JobObjectBasicProcessIdList, + JobObjectBasicUIRestrictions, + JobObjectSecurityLimitInformation, + JobObjectEndOfJobTimeInformation, + JobObjectAssociateCompletionPortInformation, + JobObjectBasicAndIoAccountingInformation, + JobObjectExtendedLimitInformation, + JobObjectJobSetInformation, + MaxJobObjectInfoClass +) = range(1, 12) + + +prototype = WINFUNCTYPE(DWORD, DWORD, POINTER(HANDLE), BOOL, DWORD) +flags = (1, 'count'), (1, 'handles'), (1, 'wait_all'), (1, 'milliseconds') +_wait_multiple = prototype(('WaitForMultipleObjects', windll.kernel32), flags) +def errcheck(result, func, args): + if result == WAIT_FAILED: raise WinError() + return args +_wait_multiple.errcheck = errcheck +def WaitForMultipleObjects(handles, wait_all, timeout): + n = len(handles) + handles = (Handle.from_param(handle) for handle in handles) + timeout = ceil(timeout * 1000) + return _wait_multiple(n, (HANDLE * n)(*handles), wait_all, timeout) + +# WAIT_OBJECT_0 defined at the top of the file +WAIT_ABANDONED_0 = 0x00000080 +WAIT_TIMEOUT = 0x00000102 +WAIT_FAILED = 0xFFFFFFFF + + +try: + _wait_single = WaitForSingleObject +except NameError: + prototype = WINFUNCTYPE(DWORD, Handle, DWORD) + flags = (1, 'handle'), (1, 'milliseconds') + _wait_single = prototype(('WaitForSingleObject', windll.kernel32), flags) + def errcheck(result, func, args): + if result == WAIT_FAILED: raise WinError() + return args + _wait_single.errcheck = errcheck +def WaitForSingleObject(handle, timeout): + return _wait_single(handle, ceil(timeout * 1000)) + + +try: + GetStdHandle +except NameError: + prototype = WINFUNCTYPE(HANDLE, DWORD) + flags = (1, 'which'), + GetStdHandle = prototype(('GetStdHandle', windll.kernel32), flags) + def errcheck(result, func, args): + if result == INVALID_HANDLE_VALUE: raise WinError() + return args if result else None + GetStdHandle.errcheck = errcheck + + +try: + TerminateProcess +except NameError: + prototype = WINFUNCTYPE(BOOL, Handle, UINT) + flags = (1, 'process'), (1, 'exitcode') + TerminateProcess = prototype(('TerminateProcess', windll.kernel32), flags) + def errcheck(result, func, args): + if not result: raise WinError() + TerminateProcess.errcheck = errcheck + + +# Do not show error messages due to errors in the program being tested +try: + errmode = windll.kernel32.GetErrorMode() +except AttributeError: + # GetErrorMode is available on Vista/2008 and up + errmode = windll.kernel32.SetErrorMode(0) +windll.kernel32.SetErrorMode(errmode | 0x8003) + +stdin = GetStdHandle(STD_INPUT_HANDLE) +try: + SetConsoleMode(stdin, ENABLE_PROCESSED_INPUT) +except WindowsError: + console_input = False +else: + console_input = True + FlushConsoleInputBuffer(stdin) + +def call(*args, **kwargs): + case = kwargs.pop('case') + job = CreateJobObject(None) + flags = 0 + if case.maxcputime: + flags |= JOB_OBJECT_LIMIT_PROCESS_TIME + if case.maxmemory: + flags |= JOB_OBJECT_LIMIT_PROCESS_MEMORY + limits = JOBOBJECT_EXTENDED_LIMIT_INFORMATION( + JOBOBJECT_BASIC_LIMIT_INFORMATION( + PerProcessUserTimeLimit=ceil((case.maxcputime or 0)*10000000), + LimitFlags=flags, + ), + ProcessMemoryLimit=ceil((case.maxmemory or 0)*1048576), + ) + SetInformationJobObject(job, JobObjectExtendedLimitInformation, limits) + try: + case.process = Popen(*args, **kwargs) + except OSError: + raise CannotStartTestee(sys.exc_info()[1]) + case.time_started = clock() + AssignProcessToJobObject(job, case.process._handle) + if not console_input: + if case.maxwalltime: + if (WaitForSingleObject(case.process._handle, case.maxwalltime) != + WAIT_OBJECT_0): + raise WallTimeLimitExceeded + else: + case.process.wait() + else: + handles = case.process._handle, stdin + if case.maxwalltime: + time_end = case.time_started + case.maxwalltime + while case.process.poll() is None: + remaining = time_end - clock() + if remaining > 0: + if (WaitForMultipleObjects(handles, False, remaining) == + WAIT_OBJECT_0 + 1): + ir = ReadConsoleInput(stdin) + if (ir and + ir.EventType == KEY_EVENT and + ir.Event.KeyEvent.bKeyDown and + ir.Event.KeyEvent.wVirtualKeyCode == 27): + raise CanceledByUser + else: + raise WallTimeLimitExceeded + else: + while case.process.poll() is None: + if (WaitForMultipleObjects(handles, False, INFINITE) == + WAIT_OBJECT_0 + 1): + ir = ReadConsoleInput(stdin) + if (ir and + ir.EventType == KEY_EVENT and + ir.Event.KeyEvent.bKeyDown and + ir.Event.KeyEvent.wVirtualKeyCode == 27): + raise CanceledByUser + case.time_stopped = clock() + if GetProcessTimes: + try: + times = GetProcessTimes(case.process._handle) + except WindowsError: + pass + else: + if case.maxcputime or not case.maxwalltime: + cputime = times.kernel + times.user + case.time_stopped = cputime + case.time_started = 0 + case.time_limit_string = case.cpu_time_limit_string + if case.maxcputime and cputime > case.maxcputime: + raise CPUTimeLimitExceeded + else: + case.time_stopped = times.exit + case.time_started = times.creation + walltime = times.exit - times.creation + if case.maxwalltime and walltime > case.maxwalltime: + raise WallTimeLimitExceeded + if case.maxcputime and case.process.returncode == 1816: + raise CPUTimeLimitExceeded + if case.maxmemory and GetProcessMemoryInfo: + try: + counters = GetProcessMemoryInfo(case.process._handle) + except WindowsError: + pass + else: + if counters.PeakPagefileUsage > case.maxmemory * 1048576: + raise MemoryLimitExceeded + + +def kill(process): + # Give up after three attempts + for i in range(3): + try: + try: + process.terminate() + except AttributeError: + TerminateProcess(process._handle, 1) + except WindowsError: + time.sleep(0) + else: + break