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