Spaces:
Running
Running
"""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)) | |