##// END OF EJS Templates
Merge pull request #1424 from mwiebe/win32_shell...
Min RK -
r6185:4953f2da merge
parent child Browse files
Show More
This diff has been collapsed as it changes many lines, (574 lines changed) Show them Hide them
@@ -0,0 +1,574 b''
1 """Windows-specific implementation of process utilities with direct WinAPI.
2
3 This file is meant to be used by process.py
4 """
5
6 #-----------------------------------------------------------------------------
7 # Copyright (C) 2010-2011 The IPython Development Team
8 #
9 # Distributed under the terms of the BSD License. The full license is in
10 # the file COPYING, distributed as part of this software.
11 #-----------------------------------------------------------------------------
12
13 from __future__ import print_function
14
15 # stdlib
16 import os, sys, time, threading
17 import ctypes, msvcrt
18
19 # Win32 API types needed for the API calls
20 from ctypes import POINTER
21 from ctypes.wintypes import HANDLE, HLOCAL, LPVOID, WORD, DWORD, BOOL, \
22 ULONG, LPCWSTR
23 LPDWORD = POINTER(DWORD)
24 LPHANDLE = POINTER(HANDLE)
25 ULONG_PTR = POINTER(ULONG)
26 class SECURITY_ATTRIBUTES(ctypes.Structure):
27 _fields_ = [("nLength", DWORD),
28 ("lpSecurityDescriptor", LPVOID),
29 ("bInheritHandle", BOOL)]
30 LPSECURITY_ATTRIBUTES = POINTER(SECURITY_ATTRIBUTES)
31 class STARTUPINFO(ctypes.Structure):
32 _fields_ = [("cb", DWORD),
33 ("lpReserved", LPCWSTR),
34 ("lpDesktop", LPCWSTR),
35 ("lpTitle", LPCWSTR),
36 ("dwX", DWORD),
37 ("dwY", DWORD),
38 ("dwXSize", DWORD),
39 ("dwYSize", DWORD),
40 ("dwXCountChars", DWORD),
41 ("dwYCountChars", DWORD),
42 ("dwFillAttribute", DWORD),
43 ("dwFlags", DWORD),
44 ("wShowWindow", WORD),
45 ("cbReserved2", WORD),
46 ("lpReserved2", LPVOID),
47 ("hStdInput", HANDLE),
48 ("hStdOutput", HANDLE),
49 ("hStdError", HANDLE)]
50 LPSTARTUPINFO = POINTER(STARTUPINFO)
51 class PROCESS_INFORMATION(ctypes.Structure):
52 _fields_ = [("hProcess", HANDLE),
53 ("hThread", HANDLE),
54 ("dwProcessId", DWORD),
55 ("dwThreadId", DWORD)]
56 LPPROCESS_INFORMATION = POINTER(PROCESS_INFORMATION)
57
58 # Win32 API constants needed
59 ERROR_HANDLE_EOF = 38
60 ERROR_BROKEN_PIPE = 109
61 ERROR_NO_DATA = 232
62 HANDLE_FLAG_INHERIT = 0x0001
63 STARTF_USESTDHANDLES = 0x0100
64 CREATE_SUSPENDED = 0x0004
65 CREATE_NEW_CONSOLE = 0x0010
66 CREATE_NO_WINDOW = 0x08000000
67 STILL_ACTIVE = 259
68 WAIT_TIMEOUT = 0x0102
69 WAIT_FAILED = 0xFFFFFFFF
70 INFINITE = 0xFFFFFFFF
71 DUPLICATE_SAME_ACCESS = 0x00000002
72 ENABLE_ECHO_INPUT = 0x0004
73 ENABLE_LINE_INPUT = 0x0002
74 ENABLE_PROCESSED_INPUT = 0x0001
75
76 # Win32 API functions needed
77 GetLastError = ctypes.windll.kernel32.GetLastError
78 GetLastError.argtypes = []
79 GetLastError.restype = DWORD
80
81 CreateFile = ctypes.windll.kernel32.CreateFileW
82 CreateFile.argtypes = [LPCWSTR, DWORD, DWORD, LPVOID, DWORD, DWORD, HANDLE]
83 CreateFile.restype = HANDLE
84
85 CreatePipe = ctypes.windll.kernel32.CreatePipe
86 CreatePipe.argtypes = [POINTER(HANDLE), POINTER(HANDLE),
87 LPSECURITY_ATTRIBUTES, DWORD]
88 CreatePipe.restype = BOOL
89
90 CreateProcess = ctypes.windll.kernel32.CreateProcessW
91 CreateProcess.argtypes = [LPCWSTR, LPCWSTR, LPSECURITY_ATTRIBUTES,
92 LPSECURITY_ATTRIBUTES, BOOL, DWORD, LPVOID, LPCWSTR, LPSTARTUPINFO,
93 LPPROCESS_INFORMATION]
94 CreateProcess.restype = BOOL
95
96 GetExitCodeProcess = ctypes.windll.kernel32.GetExitCodeProcess
97 GetExitCodeProcess.argtypes = [HANDLE, LPDWORD]
98 GetExitCodeProcess.restype = BOOL
99
100 GetCurrentProcess = ctypes.windll.kernel32.GetCurrentProcess
101 GetCurrentProcess.argtypes = []
102 GetCurrentProcess.restype = HANDLE
103
104 ResumeThread = ctypes.windll.kernel32.ResumeThread
105 ResumeThread.argtypes = [HANDLE]
106 ResumeThread.restype = DWORD
107
108 ReadFile = ctypes.windll.kernel32.ReadFile
109 ReadFile.argtypes = [HANDLE, LPVOID, DWORD, LPDWORD, LPVOID]
110 ReadFile.restype = BOOL
111
112 WriteFile = ctypes.windll.kernel32.WriteFile
113 WriteFile.argtypes = [HANDLE, LPVOID, DWORD, LPDWORD, LPVOID]
114 WriteFile.restype = BOOL
115
116 GetConsoleMode = ctypes.windll.kernel32.GetConsoleMode
117 GetConsoleMode.argtypes = [HANDLE, LPDWORD]
118 GetConsoleMode.restype = BOOL
119
120 SetConsoleMode = ctypes.windll.kernel32.SetConsoleMode
121 SetConsoleMode.argtypes = [HANDLE, DWORD]
122 SetConsoleMode.restype = BOOL
123
124 FlushConsoleInputBuffer = ctypes.windll.kernel32.FlushConsoleInputBuffer
125 FlushConsoleInputBuffer.argtypes = [HANDLE]
126 FlushConsoleInputBuffer.restype = BOOL
127
128 WaitForSingleObject = ctypes.windll.kernel32.WaitForSingleObject
129 WaitForSingleObject.argtypes = [HANDLE, DWORD]
130 WaitForSingleObject.restype = DWORD
131
132 DuplicateHandle = ctypes.windll.kernel32.DuplicateHandle
133 DuplicateHandle.argtypes = [HANDLE, HANDLE, HANDLE, LPHANDLE,
134 DWORD, BOOL, DWORD]
135 DuplicateHandle.restype = BOOL
136
137 SetHandleInformation = ctypes.windll.kernel32.SetHandleInformation
138 SetHandleInformation.argtypes = [HANDLE, DWORD, DWORD]
139 SetHandleInformation.restype = BOOL
140
141 CloseHandle = ctypes.windll.kernel32.CloseHandle
142 CloseHandle.argtypes = [HANDLE]
143 CloseHandle.restype = BOOL
144
145 CommandLineToArgvW = ctypes.windll.shell32.CommandLineToArgvW
146 CommandLineToArgvW.argtypes = [LPCWSTR, POINTER(ctypes.c_int)]
147 CommandLineToArgvW.restype = POINTER(LPCWSTR)
148
149 LocalFree = ctypes.windll.kernel32.LocalFree
150 LocalFree.argtypes = [HLOCAL]
151 LocalFree.restype = HLOCAL
152
153 class AvoidUNCPath(object):
154 """A context manager to protect command execution from UNC paths.
155
156 In the Win32 API, commands can't be invoked with the cwd being a UNC path.
157 This context manager temporarily changes directory to the 'C:' drive on
158 entering, and restores the original working directory on exit.
159
160 The context manager returns the starting working directory *if* it made a
161 change and None otherwise, so that users can apply the necessary adjustment
162 to their system calls in the event of a change.
163
164 Example
165 -------
166 ::
167 cmd = 'dir'
168 with AvoidUNCPath() as path:
169 if path is not None:
170 cmd = '"pushd %s &&"%s' % (path, cmd)
171 os.system(cmd)
172 """
173 def __enter__(self):
174 self.path = os.getcwdu()
175 self.is_unc_path = self.path.startswith(r"\\")
176 if self.is_unc_path:
177 # change to c drive (as cmd.exe cannot handle UNC addresses)
178 os.chdir("C:")
179 return self.path
180 else:
181 # We return None to signal that there was no change in the working
182 # directory
183 return None
184
185 def __exit__(self, exc_type, exc_value, traceback):
186 if self.is_unc_path:
187 os.chdir(self.path)
188
189
190 class Win32ShellCommandController(object):
191 """Runs a shell command in a 'with' context.
192
193 This implementation is Win32-specific.
194
195 Example:
196 # Runs the command interactively with default console stdin/stdout
197 with ShellCommandController('python -i') as scc:
198 scc.run()
199
200 # Runs the command using the provided functions for stdin/stdout
201 def my_stdout_func(s):
202 # print or save the string 's'
203 write_to_stdout(s)
204 def my_stdin_func():
205 # If input is available, return it as a string.
206 if input_available():
207 return get_input()
208 # If no input available, return None after a short delay to
209 # keep from blocking.
210 else:
211 time.sleep(0.01)
212 return None
213
214 with ShellCommandController('python -i') as scc:
215 scc.run(my_stdout_func, my_stdin_func)
216 """
217
218 def __init__(self, cmd, mergeout = True):
219 """Initializes the shell command controller.
220
221 The cmd is the program to execute, and mergeout is
222 whether to blend stdout and stderr into one output
223 in stdout. Merging them together in this fashion more
224 reliably keeps stdout and stderr in the correct order
225 especially for interactive shell usage.
226 """
227 self.cmd = cmd
228 self.mergeout = mergeout
229
230 def __enter__(self):
231 cmd = self.cmd
232 mergeout = self.mergeout
233
234 self.hstdout, self.hstdin, self.hstderr = None, None, None
235 self.piProcInfo = None
236 try:
237 p_hstdout, c_hstdout, p_hstderr, \
238 c_hstderr, p_hstdin, c_hstdin = [None]*6
239
240 # SECURITY_ATTRIBUTES with inherit handle set to True
241 saAttr = SECURITY_ATTRIBUTES()
242 saAttr.nLength = ctypes.sizeof(saAttr)
243 saAttr.bInheritHandle = True
244 saAttr.lpSecurityDescriptor = None
245
246 def create_pipe(uninherit):
247 """Creates a Windows pipe, which consists of two handles.
248
249 The 'uninherit' parameter controls which handle is not
250 inherited by the child process.
251 """
252 handles = HANDLE(), HANDLE()
253 if not CreatePipe(ctypes.byref(handles[0]),
254 ctypes.byref(handles[1]), ctypes.byref(saAttr), 0):
255 raise ctypes.WinError()
256 if not SetHandleInformation(handles[uninherit],
257 HANDLE_FLAG_INHERIT, 0):
258 raise ctypes.WinError()
259 return handles[0].value, handles[1].value
260
261 p_hstdout, c_hstdout = create_pipe(uninherit=0)
262 # 'mergeout' signals that stdout and stderr should be merged.
263 # We do that by using one pipe for both of them.
264 if mergeout:
265 c_hstderr = HANDLE()
266 if not DuplicateHandle(GetCurrentProcess(), c_hstdout,
267 GetCurrentProcess(), ctypes.byref(c_hstderr),
268 0, True, DUPLICATE_SAME_ACCESS):
269 raise ctypes.WinError()
270 else:
271 p_hstderr, c_hstderr = create_pipe(uninherit=0)
272 c_hstdin, p_hstdin = create_pipe(uninherit=1)
273
274 # Create the process object
275 piProcInfo = PROCESS_INFORMATION()
276 siStartInfo = STARTUPINFO()
277 siStartInfo.cb = ctypes.sizeof(siStartInfo)
278 siStartInfo.hStdInput = c_hstdin
279 siStartInfo.hStdOutput = c_hstdout
280 siStartInfo.hStdError = c_hstderr
281 siStartInfo.dwFlags = STARTF_USESTDHANDLES
282 dwCreationFlags = CREATE_SUSPENDED | CREATE_NO_WINDOW # | CREATE_NEW_CONSOLE
283
284 if not CreateProcess(None,
285 u"cmd.exe /c " + cmd,
286 None, None, True, dwCreationFlags,
287 None, None, ctypes.byref(siStartInfo),
288 ctypes.byref(piProcInfo)):
289 raise ctypes.WinError()
290
291 # Close this process's versions of the child handles
292 CloseHandle(c_hstdin)
293 c_hstdin = None
294 CloseHandle(c_hstdout)
295 c_hstdout = None
296 if c_hstderr != None:
297 CloseHandle(c_hstderr)
298 c_hstderr = None
299
300 # Transfer ownership of the parent handles to the object
301 self.hstdin = p_hstdin
302 p_hstdin = None
303 self.hstdout = p_hstdout
304 p_hstdout = None
305 if not mergeout:
306 self.hstderr = p_hstderr
307 p_hstderr = None
308 self.piProcInfo = piProcInfo
309
310 finally:
311 if p_hstdin:
312 CloseHandle(p_hstdin)
313 if c_hstdin:
314 CloseHandle(c_hstdin)
315 if p_hstdout:
316 CloseHandle(p_hstdout)
317 if c_hstdout:
318 CloseHandle(c_hstdout)
319 if p_hstderr:
320 CloseHandle(p_hstderr)
321 if c_hstderr:
322 CloseHandle(c_hstderr)
323
324 return self
325
326 def _stdin_thread(self, handle, hprocess, func, stdout_func):
327 exitCode = DWORD()
328 bytesWritten = DWORD(0)
329 while True:
330 #print("stdin thread loop start")
331 # Get the input string (may be bytes or unicode)
332 data = func()
333
334 # None signals to poll whether the process has exited
335 if data is None:
336 #print("checking for process completion")
337 if not GetExitCodeProcess(hprocess, ctypes.byref(exitCode)):
338 raise ctypes.WinError()
339 if exitCode.value != STILL_ACTIVE:
340 return
341 # TESTING: Does zero-sized writefile help?
342 if not WriteFile(handle, "", 0,
343 ctypes.byref(bytesWritten), None):
344 raise ctypes.WinError()
345 continue
346 #print("\nGot str %s\n" % repr(data), file=sys.stderr)
347
348 # Encode the string to the console encoding
349 if isinstance(data, unicode): #FIXME: Python3
350 data = data.encode('utf_8')
351
352 # What we have now must be a string of bytes
353 if not isinstance(data, str): #FIXME: Python3
354 raise RuntimeError("internal stdin function string error")
355
356 # An empty string signals EOF
357 if len(data) == 0:
358 return
359
360 # In a windows console, sometimes the input is echoed,
361 # but sometimes not. How do we determine when to do this?
362 stdout_func(data)
363 # WriteFile may not accept all the data at once.
364 # Loop until everything is processed
365 while len(data) != 0:
366 #print("Calling writefile")
367 if not WriteFile(handle, data, len(data),
368 ctypes.byref(bytesWritten), None):
369 # This occurs at exit
370 if GetLastError() == ERROR_NO_DATA:
371 return
372 raise ctypes.WinError()
373 #print("Called writefile")
374 data = data[bytesWritten.value:]
375
376 def _stdout_thread(self, handle, func):
377 # Allocate the output buffer
378 data = ctypes.create_string_buffer(4096)
379 while True:
380 bytesRead = DWORD(0)
381 if not ReadFile(handle, data, 4096,
382 ctypes.byref(bytesRead), None):
383 le = GetLastError()
384 if le == ERROR_BROKEN_PIPE:
385 return
386 else:
387 raise ctypes.WinError()
388 # FIXME: Python3
389 s = data.value[0:bytesRead.value]
390 #print("\nv: %s" % repr(s), file=sys.stderr)
391 func(s.decode('utf_8', 'replace'))
392
393 def run(self, stdout_func = None, stdin_func = None, stderr_func = None):
394 """Runs the process, using the provided functions for I/O.
395
396 The function stdin_func should return strings whenever a
397 character or characters become available.
398 The functions stdout_func and stderr_func are called whenever
399 something is printed to stdout or stderr, respectively.
400 These functions are called from different threads (but not
401 concurrently, because of the GIL).
402 """
403 if stdout_func == None and stdin_func == None and stderr_func == None:
404 return self._run_stdio()
405
406 if stderr_func != None and self.mergeout:
407 raise RuntimeError("Shell command was initiated with "
408 "merged stdin/stdout, but a separate stderr_func "
409 "was provided to the run() method")
410
411 # Create a thread for each input/output handle
412 stdin_thread = None
413 threads = []
414 if stdin_func:
415 stdin_thread = threading.Thread(target=self._stdin_thread,
416 args=(self.hstdin, self.piProcInfo.hProcess,
417 stdin_func, stdout_func))
418 threads.append(threading.Thread(target=self._stdout_thread,
419 args=(self.hstdout, stdout_func)))
420 if not self.mergeout:
421 if stderr_func == None:
422 stderr_func = stdout_func
423 threads.append(threading.Thread(target=self._stdout_thread,
424 args=(self.hstderr, stderr_func)))
425 # Start the I/O threads and the process
426 if ResumeThread(self.piProcInfo.hThread) == 0xFFFFFFFF:
427 raise ctypes.WinError()
428 if stdin_thread is not None:
429 stdin_thread.start()
430 for thread in threads:
431 thread.start()
432 # Wait for the process to complete
433 if WaitForSingleObject(self.piProcInfo.hProcess, INFINITE) == \
434 WAIT_FAILED:
435 raise ctypes.WinError()
436 # Wait for the I/O threads to complete
437 for thread in threads:
438 thread.join()
439
440 # Wait for the stdin thread to complete
441 if stdin_thread is not None:
442 stdin_thread.join()
443
444 def _stdin_raw_nonblock(self):
445 """Use the raw Win32 handle of sys.stdin to do non-blocking reads"""
446 # WARNING: This is experimental, and produces inconsistent results.
447 # It's possible for the handle not to be appropriate for use
448 # with WaitForSingleObject, among other things.
449 handle = msvcrt.get_osfhandle(sys.stdin.fileno())
450 result = WaitForSingleObject(handle, 100)
451 if result == WAIT_FAILED:
452 raise ctypes.WinError()
453 elif result == WAIT_TIMEOUT:
454 print(".", end='')
455 return None
456 else:
457 data = ctypes.create_string_buffer(256)
458 bytesRead = DWORD(0)
459 print('?', end='')
460
461 if not ReadFile(handle, data, 256,
462 ctypes.byref(bytesRead), None):
463 raise ctypes.WinError()
464 # This ensures the non-blocking works with an actual console
465 # Not checking the error, so the processing will still work with
466 # other handle types
467 FlushConsoleInputBuffer(handle)
468
469 data = data.value
470 data = data.replace('\r\n', '\n')
471 data = data.replace('\r', '\n')
472 print(repr(data) + " ", end='')
473 return data
474
475 def _stdin_raw_block(self):
476 """Use a blocking stdin read"""
477 # The big problem with the blocking read is that it doesn't
478 # exit when it's supposed to in all contexts. An extra
479 # key-press may be required to trigger the exit.
480 try:
481 data = sys.stdin.read(1)
482 data = data.replace('\r', '\n')
483 return data
484 except WindowsError as we:
485 if we.winerror == ERROR_NO_DATA:
486 # This error occurs when the pipe is closed
487 return None
488 else:
489 # Otherwise let the error propagate
490 raise we
491
492 def _stdout_raw(self, s):
493 """Writes the string to stdout"""
494 print(s, end='', file=sys.stdout)
495 sys.stdout.flush()
496
497 def _stderr_raw(self, s):
498 """Writes the string to stdout"""
499 print(s, end='', file=sys.stderr)
500 sys.stderr.flush()
501
502 def _run_stdio(self):
503 """Runs the process using the system standard I/O.
504
505 IMPORTANT: stdin needs to be asynchronous, so the Python
506 sys.stdin object is not used. Instead,
507 msvcrt.kbhit/getwch are used asynchronously.
508 """
509 # Disable Line and Echo mode
510 #lpMode = DWORD()
511 #handle = msvcrt.get_osfhandle(sys.stdin.fileno())
512 #if GetConsoleMode(handle, ctypes.byref(lpMode)):
513 # set_console_mode = True
514 # if not SetConsoleMode(handle, lpMode.value &
515 # ~(ENABLE_ECHO_INPUT | ENABLE_LINE_INPUT | ENABLE_PROCESSED_INPUT)):
516 # raise ctypes.WinError()
517
518 if self.mergeout:
519 return self.run(stdout_func = self._stdout_raw,
520 stdin_func = self._stdin_raw_block)
521 else:
522 return self.run(stdout_func = self._stdout_raw,
523 stdin_func = self._stdin_raw_block,
524 stderr_func = self._stderr_raw)
525
526 # Restore the previous console mode
527 #if set_console_mode:
528 # if not SetConsoleMode(handle, lpMode.value):
529 # raise ctypes.WinError()
530
531 def __exit__(self, exc_type, exc_value, traceback):
532 if self.hstdin:
533 CloseHandle(self.hstdin)
534 self.hstdin = None
535 if self.hstdout:
536 CloseHandle(self.hstdout)
537 self.hstdout = None
538 if self.hstderr:
539 CloseHandle(self.hstderr)
540 self.hstderr = None
541 if self.piProcInfo != None:
542 CloseHandle(self.piProcInfo.hProcess)
543 CloseHandle(self.piProcInfo.hThread)
544 self.piProcInfo = None
545
546
547 def system(cmd):
548 """Win32 version of os.system() that works with network shares.
549
550 Note that this implementation returns None, as meant for use in IPython.
551
552 Parameters
553 ----------
554 cmd : str
555 A command to be executed in the system shell.
556
557 Returns
558 -------
559 None : we explicitly do NOT return the subprocess status code, as this
560 utility is meant to be used extensively in IPython, where any return value
561 would trigger :func:`sys.displayhook` calls.
562 """
563 with AvoidUNCPath() as path:
564 if path is not None:
565 cmd = '"pushd %s &&"%s' % (path, cmd)
566 with Win32ShellCommandController(cmd) as scc:
567 scc.run()
568
569
570 if __name__ == "__main__":
571 print("Test starting!")
572 #system("cmd")
573 system("python -i")
574 print("Test finished!")
@@ -19,6 +19,7 b' from __future__ import print_function'
19 import os
19 import os
20 import sys
20 import sys
21 import ctypes
21 import ctypes
22 import msvcrt
22
23
23 from ctypes import c_int, POINTER
24 from ctypes import c_int, POINTER
24 from ctypes.wintypes import LPCWSTR, HLOCAL
25 from ctypes.wintypes import LPCWSTR, HLOCAL
@@ -121,10 +122,10 b' def system(cmd):'
121 utility is meant to be used extensively in IPython, where any return value
122 utility is meant to be used extensively in IPython, where any return value
122 would trigger :func:`sys.displayhook` calls.
123 would trigger :func:`sys.displayhook` calls.
123 """
124 """
124 with AvoidUNCPath() as path:
125 # The controller provides interactivity with both
125 if path is not None:
126 # stdin and stdout
126 cmd = '"pushd %s &&"%s' % (path, cmd)
127 import _process_win32_controller
127 return process_handler(cmd, _system_body)
128 _process_win32_controller.system(cmd)
128
129
129
130
130 def getoutput(cmd):
131 def getoutput(cmd):
General Comments 0
You need to be logged in to leave comments. Login now