You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
217 lines
6.6 KiB
Python
217 lines
6.6 KiB
Python
"""Windows-specific implementation of process utilities.
|
|
|
|
This file is only meant to be imported by process.py, not by end-users.
|
|
"""
|
|
|
|
import ctypes
|
|
import os
|
|
import subprocess
|
|
import sys
|
|
import time
|
|
from ctypes import POINTER, c_int
|
|
from ctypes.wintypes import HLOCAL, LPCWSTR
|
|
from subprocess import STDOUT
|
|
from threading import Thread
|
|
from types import TracebackType
|
|
from typing import List, Optional
|
|
|
|
from . import py3compat
|
|
from ._process_common import arg_split as py_arg_split
|
|
|
|
from ._process_common import process_handler, read_no_interrupt
|
|
from .encoding import DEFAULT_ENCODING
|
|
|
|
|
|
class AvoidUNCPath:
|
|
"""A context manager to protect command execution from UNC paths.
|
|
|
|
In the Win32 API, commands can't be invoked with the cwd being a UNC path.
|
|
This context manager temporarily changes directory to the 'C:' drive on
|
|
entering, and restores the original working directory on exit.
|
|
|
|
The context manager returns the starting working directory *if* it made a
|
|
change and None otherwise, so that users can apply the necessary adjustment
|
|
to their system calls in the event of a change.
|
|
|
|
Examples
|
|
--------
|
|
::
|
|
cmd = 'dir'
|
|
with AvoidUNCPath() as path:
|
|
if path is not None:
|
|
cmd = '"pushd %s &&"%s' % (path, cmd)
|
|
os.system(cmd)
|
|
"""
|
|
|
|
def __enter__(self) -> Optional[str]:
|
|
self.path = os.getcwd()
|
|
self.is_unc_path = self.path.startswith(r"\\")
|
|
if self.is_unc_path:
|
|
# change to c drive (as cmd.exe cannot handle UNC addresses)
|
|
os.chdir("C:")
|
|
return self.path
|
|
else:
|
|
# We return None to signal that there was no change in the working
|
|
# directory
|
|
return None
|
|
|
|
def __exit__(
|
|
self,
|
|
exc_type: Optional[type[BaseException]],
|
|
exc_value: Optional[BaseException],
|
|
traceback: TracebackType,
|
|
) -> None:
|
|
if self.is_unc_path:
|
|
os.chdir(self.path)
|
|
|
|
|
|
def _system_body(p: subprocess.Popen) -> int:
|
|
"""Callback for _system."""
|
|
enc = DEFAULT_ENCODING
|
|
|
|
# Dec 2024: in both of these functions, I'm not sure why we .splitlines()
|
|
# the bytes and then decode each line individually instead of just decoding
|
|
# the whole thing at once.
|
|
def stdout_read() -> None:
|
|
try:
|
|
assert p.stdout is not None
|
|
for byte_line in read_no_interrupt(p.stdout).splitlines():
|
|
line = byte_line.decode(enc, "replace")
|
|
print(line, file=sys.stdout)
|
|
except Exception as e:
|
|
print(f"Error reading stdout: {e}", file=sys.stderr)
|
|
|
|
def stderr_read() -> None:
|
|
try:
|
|
assert p.stderr is not None
|
|
for byte_line in read_no_interrupt(p.stderr).splitlines():
|
|
line = byte_line.decode(enc, "replace")
|
|
print(line, file=sys.stderr)
|
|
except Exception as e:
|
|
print(f"Error reading stderr: {e}", file=sys.stderr)
|
|
|
|
stdout_thread = Thread(target=stdout_read)
|
|
stderr_thread = Thread(target=stderr_read)
|
|
|
|
stdout_thread.start()
|
|
stderr_thread.start()
|
|
|
|
# Wait to finish for returncode. Unfortunately, Python has a bug where
|
|
# wait() isn't interruptible (https://bugs.python.org/issue28168) so poll in
|
|
# a loop instead of just doing `return p.wait()`
|
|
while True:
|
|
result = p.poll()
|
|
if result is None:
|
|
time.sleep(0.01)
|
|
else:
|
|
break
|
|
|
|
# Join the threads to ensure they complete before returning
|
|
stdout_thread.join()
|
|
stderr_thread.join()
|
|
|
|
return result
|
|
|
|
|
|
def system(cmd: str) -> Optional[int]:
|
|
"""Win32 version of os.system() that works with network shares.
|
|
|
|
Note that this implementation returns None, as meant for use in IPython.
|
|
|
|
Parameters
|
|
----------
|
|
cmd : str or list
|
|
A command to be executed in the system shell.
|
|
|
|
Returns
|
|
-------
|
|
int : child process' exit code.
|
|
"""
|
|
# The controller provides interactivity with both
|
|
# stdin and stdout
|
|
# import _process_win32_controller
|
|
# _process_win32_controller.system(cmd)
|
|
|
|
with AvoidUNCPath() as path:
|
|
if path is not None:
|
|
cmd = '"pushd %s &&"%s' % (path, cmd)
|
|
res = process_handler(cmd, _system_body)
|
|
assert isinstance(res, int | type(None))
|
|
return res
|
|
return None
|
|
|
|
|
|
def getoutput(cmd: str) -> str:
|
|
"""Return standard output of executing cmd in a shell.
|
|
|
|
Accepts the same arguments as os.system().
|
|
|
|
Parameters
|
|
----------
|
|
cmd : str or list
|
|
A command to be executed in the system shell.
|
|
|
|
Returns
|
|
-------
|
|
stdout : str
|
|
"""
|
|
|
|
with AvoidUNCPath() as path:
|
|
if path is not None:
|
|
cmd = '"pushd %s &&"%s' % (path, cmd)
|
|
out = process_handler(cmd, lambda p: p.communicate()[0], STDOUT)
|
|
|
|
if out is None:
|
|
out = b""
|
|
return py3compat.decode(out)
|
|
|
|
|
|
try:
|
|
windll = ctypes.windll # type: ignore [attr-defined]
|
|
CommandLineToArgvW = windll.shell32.CommandLineToArgvW
|
|
CommandLineToArgvW.arg_types = [LPCWSTR, POINTER(c_int)]
|
|
CommandLineToArgvW.restype = POINTER(LPCWSTR)
|
|
LocalFree = windll.kernel32.LocalFree
|
|
LocalFree.res_type = HLOCAL
|
|
LocalFree.arg_types = [HLOCAL]
|
|
|
|
def arg_split(
|
|
commandline: str, posix: bool = False, strict: bool = True
|
|
) -> List[str]:
|
|
"""Split a command line's arguments in a shell-like manner.
|
|
|
|
This is a special version for windows that use a ctypes call to CommandLineToArgvW
|
|
to do the argv splitting. The posix parameter is ignored.
|
|
|
|
If strict=False, process_common.arg_split(...strict=False) is used instead.
|
|
"""
|
|
# CommandLineToArgvW returns path to executable if called with empty string.
|
|
if commandline.strip() == "":
|
|
return []
|
|
if not strict:
|
|
# not really a cl-arg, fallback on _process_common
|
|
return py_arg_split(commandline, posix=posix, strict=strict)
|
|
argvn = c_int()
|
|
result_pointer = CommandLineToArgvW(commandline.lstrip(), ctypes.byref(argvn))
|
|
try:
|
|
result_array_type = LPCWSTR * argvn.value
|
|
result = [
|
|
arg
|
|
for arg in result_array_type.from_address(
|
|
ctypes.addressof(result_pointer.contents)
|
|
)
|
|
if arg is not None
|
|
]
|
|
finally:
|
|
# for side effects
|
|
_ = LocalFree(result_pointer)
|
|
return result
|
|
except AttributeError:
|
|
arg_split = py_arg_split
|
|
|
|
|
|
def check_pid(pid: int) -> bool:
|
|
# OpenProcess returns 0 if no such process (of ours) exists
|
|
# positive int otherwise
|
|
return bool(windll.kernel32.OpenProcess(1, 0, pid))
|