360 lines
11 KiB
Python
360 lines
11 KiB
Python
|
"""
|
||
|
|
A small collection of useful functional tools for working with iterables.
|
||
|
|
"""
|
||
|
|
import errno
|
||
|
|
import locale
|
||
|
|
import os
|
||
|
|
import stat
|
||
|
|
import subprocess
|
||
|
|
import time
|
||
|
|
import warnings
|
||
|
|
from functools import partial
|
||
|
|
from itertools import count, islice
|
||
|
|
from typing import Any, Iterable
|
||
|
|
|
||
|
|
DIRECTORY_CLEANUP_TIMEOUT = 1.0
|
||
|
|
|
||
|
|
|
||
|
|
def _is_iterable(elem: Any) -> bool:
|
||
|
|
if getattr(elem, "__iter__", False) or isinstance(elem, Iterable):
|
||
|
|
return True
|
||
|
|
return False
|
||
|
|
|
||
|
|
|
||
|
|
def take(n: int, iterable: Iterable) -> Iterable:
|
||
|
|
"""Take n elements from the supplied iterable without consuming it.
|
||
|
|
|
||
|
|
:param int n: Number of unique groups
|
||
|
|
:param iter iterable: An iterable to split up
|
||
|
|
"""
|
||
|
|
return list(islice(iterable, n))
|
||
|
|
|
||
|
|
|
||
|
|
def chunked(n: int, iterable: Iterable) -> Iterable:
|
||
|
|
"""Split an iterable into lists of length *n*.
|
||
|
|
|
||
|
|
:param int n: Number of unique groups
|
||
|
|
:param iter iterable: An iterable to split up
|
||
|
|
|
||
|
|
"""
|
||
|
|
return iter(partial(take, n, iter(iterable)), [])
|
||
|
|
|
||
|
|
|
||
|
|
def unnest(elem: Iterable) -> Any:
|
||
|
|
"""Flatten an arbitrarily nested iterable.
|
||
|
|
|
||
|
|
:param elem: An iterable to flatten
|
||
|
|
:type elem: :class:`~collections.Iterable`
|
||
|
|
>>> nested_iterable = (
|
||
|
|
1234, (3456, 4398345, (234234)), (
|
||
|
|
2396, (
|
||
|
|
928379, 29384, (
|
||
|
|
293759, 2347, (
|
||
|
|
2098, 7987, 27599
|
||
|
|
)
|
||
|
|
)
|
||
|
|
)
|
||
|
|
)
|
||
|
|
)
|
||
|
|
>>> list(unnest(nested_iterable))
|
||
|
|
[1234, 3456, 4398345, 234234, 2396, 928379, 29384, 293759,
|
||
|
|
2347, 2098, 7987, 27599]
|
||
|
|
"""
|
||
|
|
|
||
|
|
if isinstance(elem, Iterable) and not isinstance(elem, str):
|
||
|
|
for el in elem:
|
||
|
|
if isinstance(el, Iterable) and not isinstance(el, str):
|
||
|
|
yield from unnest(el)
|
||
|
|
else:
|
||
|
|
yield el
|
||
|
|
else:
|
||
|
|
yield elem
|
||
|
|
|
||
|
|
|
||
|
|
def dedup(iterable: Iterable) -> Iterable:
|
||
|
|
"""Deduplicate an iterable object like iter(set(iterable)) but order-
|
||
|
|
preserved."""
|
||
|
|
|
||
|
|
return iter(dict.fromkeys(iterable))
|
||
|
|
|
||
|
|
|
||
|
|
def is_readonly_path(fn: os.PathLike) -> bool:
|
||
|
|
"""check if a provided path exists and is readonly.
|
||
|
|
|
||
|
|
permissions check is `bool(path.stat & stat.s_iread)` or `not
|
||
|
|
os.access(path, os.w_ok)`
|
||
|
|
"""
|
||
|
|
if os.path.exists(fn):
|
||
|
|
file_stat = os.stat(fn).st_mode
|
||
|
|
return not bool(file_stat & stat.s_iwrite) or not os.access(fn, os.w_ok)
|
||
|
|
return False
|
||
|
|
|
||
|
|
|
||
|
|
def _wait_for_files(path): # pragma: no cover
|
||
|
|
"""Retry with backoff up to 1 second to delete files from a directory.
|
||
|
|
|
||
|
|
:param str path: The path to crawl to delete files from
|
||
|
|
:return: A list of remaining paths or None
|
||
|
|
:rtype: Optional[List[str]]
|
||
|
|
"""
|
||
|
|
timeout = 0.001 # noqa:S101
|
||
|
|
remaining = []
|
||
|
|
while timeout < DIRECTORY_CLEANUP_TIMEOUT:
|
||
|
|
remaining = []
|
||
|
|
if os.path.isdir(path):
|
||
|
|
L = os.listdir(path)
|
||
|
|
for target in L:
|
||
|
|
_remaining = _wait_for_files(target)
|
||
|
|
if _remaining:
|
||
|
|
remaining.extend(_remaining)
|
||
|
|
continue
|
||
|
|
try:
|
||
|
|
os.unlink(path)
|
||
|
|
except FileNotFoundError as e:
|
||
|
|
if e.errno == errno.ENOENT:
|
||
|
|
return
|
||
|
|
except (OSError, PermissionError): # noqa:B014
|
||
|
|
time.sleep(timeout)
|
||
|
|
timeout *= 2
|
||
|
|
remaining.append(path)
|
||
|
|
else:
|
||
|
|
return
|
||
|
|
return remaining
|
||
|
|
|
||
|
|
|
||
|
|
def _walk_for_powershell(directory):
|
||
|
|
for _, dirs, files in os.walk(directory):
|
||
|
|
powershell = next(
|
||
|
|
iter(fn for fn in files if fn.lower() == "powershell.exe"), None
|
||
|
|
)
|
||
|
|
if powershell is not None:
|
||
|
|
return os.path.join(directory, powershell)
|
||
|
|
for subdir in dirs:
|
||
|
|
powershell = _walk_for_powershell(os.path.join(directory, subdir))
|
||
|
|
if powershell:
|
||
|
|
return powershell
|
||
|
|
return None
|
||
|
|
|
||
|
|
|
||
|
|
def _get_powershell_path():
|
||
|
|
paths = [
|
||
|
|
os.path.expandvars(r"%windir%\{0}\WindowsPowerShell").format(subdir)
|
||
|
|
for subdir in ("SysWOW64", "system32")
|
||
|
|
]
|
||
|
|
powershell_path = next(iter(_walk_for_powershell(pth) for pth in paths), None)
|
||
|
|
if not powershell_path:
|
||
|
|
powershell_path = subprocess.run(["where", "powershell"], check=False)
|
||
|
|
if powershell_path.stdout:
|
||
|
|
return powershell_path.stdout.strip()
|
||
|
|
|
||
|
|
|
||
|
|
def _get_sid_with_powershell():
|
||
|
|
powershell_path = _get_powershell_path()
|
||
|
|
if not powershell_path:
|
||
|
|
return None
|
||
|
|
args = [
|
||
|
|
powershell_path,
|
||
|
|
"-ExecutionPolicy",
|
||
|
|
"Bypass",
|
||
|
|
"-Command",
|
||
|
|
"Invoke-Expression '[System.Security.Principal.WindowsIdentity]::GetCurrent().user | Write-Host'",
|
||
|
|
]
|
||
|
|
sid = subprocess.run(args, capture_output=True, check=False)
|
||
|
|
return sid.stdout.strip()
|
||
|
|
|
||
|
|
|
||
|
|
def get_value_from_tuple(value, value_type):
|
||
|
|
try:
|
||
|
|
import winreg
|
||
|
|
except ImportError:
|
||
|
|
import _winreg as winreg
|
||
|
|
if value_type in (winreg.REG_SZ, winreg.REG_EXPAND_SZ):
|
||
|
|
if "\0" in value:
|
||
|
|
return value[: value.index("\0")]
|
||
|
|
return value
|
||
|
|
return None
|
||
|
|
|
||
|
|
|
||
|
|
def query_registry_value(root, key_name, value):
|
||
|
|
try:
|
||
|
|
import winreg
|
||
|
|
except ImportError:
|
||
|
|
import _winreg as winreg
|
||
|
|
try:
|
||
|
|
with winreg.OpenKeyEx(root, key_name, 0, winreg.KEY_READ) as key:
|
||
|
|
return get_value_from_tuple(*winreg.QueryValueEx(key, value))
|
||
|
|
except OSError:
|
||
|
|
return None
|
||
|
|
|
||
|
|
|
||
|
|
def _get_sid_from_registry():
|
||
|
|
try:
|
||
|
|
import winreg
|
||
|
|
except ImportError:
|
||
|
|
import _winreg as winreg
|
||
|
|
var_names = ("%USERPROFILE%", "%HOME%")
|
||
|
|
current_user_home = next(iter(os.path.expandvars(v) for v in var_names if v), None)
|
||
|
|
root, subkey = (
|
||
|
|
winreg.HKEY_LOCAL_MACHINE,
|
||
|
|
r"Software\Microsoft\Windows NT\CurrentVersion\ProfileList",
|
||
|
|
)
|
||
|
|
subkey_names = []
|
||
|
|
value = None
|
||
|
|
matching_key = None
|
||
|
|
try:
|
||
|
|
with winreg.OpenKeyEx(root, subkey, 0, winreg.KEY_READ) as key:
|
||
|
|
for i in count():
|
||
|
|
key_name = winreg.EnumKey(key, i)
|
||
|
|
subkey_names.append(key_name)
|
||
|
|
value = query_registry_value(
|
||
|
|
root, rf"{subkey}\{key_name}", "ProfileImagePath"
|
||
|
|
)
|
||
|
|
if value and value.lower() == current_user_home.lower():
|
||
|
|
matching_key = key_name
|
||
|
|
break
|
||
|
|
except OSError:
|
||
|
|
pass
|
||
|
|
if matching_key is not None:
|
||
|
|
return matching_key
|
||
|
|
|
||
|
|
|
||
|
|
def _get_current_user():
|
||
|
|
fns = (_get_sid_from_registry, _get_sid_with_powershell)
|
||
|
|
for fn in fns:
|
||
|
|
result = fn()
|
||
|
|
if result:
|
||
|
|
return result
|
||
|
|
return None
|
||
|
|
|
||
|
|
|
||
|
|
def _find_icacls_exe():
|
||
|
|
if os.name == "nt":
|
||
|
|
paths = [
|
||
|
|
os.path.expandvars(r"%windir%\{0}").format(subdir)
|
||
|
|
for subdir in ("system32", "SysWOW64")
|
||
|
|
]
|
||
|
|
for path in paths:
|
||
|
|
icacls_path = next(
|
||
|
|
iter(fn for fn in os.listdir(path) if fn.lower() == "icacls.exe"), None
|
||
|
|
)
|
||
|
|
if icacls_path is not None:
|
||
|
|
icacls_path = os.path.join(path, icacls_path)
|
||
|
|
return icacls_path
|
||
|
|
return None
|
||
|
|
|
||
|
|
|
||
|
|
def set_write_bit(fn: str) -> None:
|
||
|
|
"""Set read-write permissions for the current user on the target path. Fail
|
||
|
|
silently if the path doesn't exist.
|
||
|
|
|
||
|
|
:param str fn: The target filename or path
|
||
|
|
:return: None
|
||
|
|
"""
|
||
|
|
if not os.path.exists(fn):
|
||
|
|
return
|
||
|
|
file_stat = os.stat(fn).st_mode
|
||
|
|
os.chmod(fn, file_stat | stat.S_IRWXU | stat.S_IRWXG | stat.S_IRWXO)
|
||
|
|
if os.name == "nt":
|
||
|
|
user_sid = _get_current_user()
|
||
|
|
icacls_exe = _find_icacls_exe() or "icacls"
|
||
|
|
|
||
|
|
if user_sid:
|
||
|
|
c = subprocess.run(
|
||
|
|
[
|
||
|
|
icacls_exe,
|
||
|
|
f"''{fn}''",
|
||
|
|
"/grant",
|
||
|
|
f"{user_sid}:WD",
|
||
|
|
"/T",
|
||
|
|
"/C",
|
||
|
|
"/Q",
|
||
|
|
],
|
||
|
|
capture_output=True,
|
||
|
|
# 2020-06-12 Yukihiko Shinoda
|
||
|
|
# There are 3 way to get system default encoding in Stack Overflow.
|
||
|
|
# see: https://stackoverflow.com/questions/37506535/how-to-get-the-system-default-encoding-in-python-2-x
|
||
|
|
# I investigated these way by using Shift-JIS Windows.
|
||
|
|
# >>> import locale
|
||
|
|
# >>> locale.getpreferredencoding()
|
||
|
|
# "cp932" (Shift-JIS)
|
||
|
|
# >>> import sys
|
||
|
|
# >>> sys.getdefaultencoding()
|
||
|
|
# "utf-8"
|
||
|
|
# >>> sys.stdout.encoding
|
||
|
|
# "UTF8"
|
||
|
|
encoding=locale.getpreferredencoding(),
|
||
|
|
check=False,
|
||
|
|
)
|
||
|
|
if not c.err and c.returncode == 0:
|
||
|
|
return
|
||
|
|
|
||
|
|
if not os.path.isdir(fn):
|
||
|
|
for path in [fn, os.path.dirname(fn)]:
|
||
|
|
try:
|
||
|
|
os.chflags(path, 0)
|
||
|
|
except AttributeError: # noqa: PERF203
|
||
|
|
pass
|
||
|
|
return None
|
||
|
|
for root, dirs, files in os.walk(fn, topdown=False):
|
||
|
|
for dir_ in [os.path.join(root, d) for d in dirs]:
|
||
|
|
set_write_bit(dir_)
|
||
|
|
for file_ in [os.path.join(root, f) for f in files]:
|
||
|
|
set_write_bit(file_)
|
||
|
|
|
||
|
|
|
||
|
|
def handle_remove_readonly(func, path, exc):
|
||
|
|
"""Error handler for shutil.rmtree.
|
||
|
|
|
||
|
|
Windows source repo folders are read-only by default, so this error handler
|
||
|
|
attempts to set them as writeable and then proceed with deletion.
|
||
|
|
|
||
|
|
:param function func: The caller function
|
||
|
|
:param str path: The target path for removal
|
||
|
|
:param Exception exc: The raised exception
|
||
|
|
|
||
|
|
This function will call check :func:`is_readonly_path` before attempting to call
|
||
|
|
:func:`set_write_bit` on the target path and try again.
|
||
|
|
"""
|
||
|
|
|
||
|
|
PERM_ERRORS = (errno.EACCES, errno.EPERM, errno.ENOENT)
|
||
|
|
default_warning_message = "Unable to remove file due to permissions restriction: {!r}"
|
||
|
|
# split the initial exception out into its type, exception, and traceback
|
||
|
|
exc_type, exc_exception, exc_tb = exc
|
||
|
|
if is_readonly_path(path):
|
||
|
|
# Apply write permission and call original function
|
||
|
|
set_write_bit(path)
|
||
|
|
try:
|
||
|
|
func(path)
|
||
|
|
except (OSError, FileNotFoundError, PermissionError) as e: # pragma: no cover
|
||
|
|
if e.errno in PERM_ERRORS:
|
||
|
|
if e.errno == errno.ENOENT:
|
||
|
|
return
|
||
|
|
remaining = None
|
||
|
|
if os.path.isdir(path):
|
||
|
|
remaining = _wait_for_files(path)
|
||
|
|
if remaining:
|
||
|
|
warnings.warn(
|
||
|
|
default_warning_message.format(path),
|
||
|
|
ResourceWarning,
|
||
|
|
stacklevel=2,
|
||
|
|
)
|
||
|
|
else:
|
||
|
|
func(path, ignore_errors=True)
|
||
|
|
return
|
||
|
|
|
||
|
|
if exc_exception.errno in PERM_ERRORS:
|
||
|
|
set_write_bit(path)
|
||
|
|
remaining = _wait_for_files(path)
|
||
|
|
try:
|
||
|
|
func(path)
|
||
|
|
except (OSError, FileNotFoundError, PermissionError) as e: # noqa:B014
|
||
|
|
if e.errno in PERM_ERRORS and e.errno != errno.ENOENT: # File still exists
|
||
|
|
warnings.warn(
|
||
|
|
default_warning_message.format(path),
|
||
|
|
ResourceWarning,
|
||
|
|
stacklevel=2,
|
||
|
|
)
|
||
|
|
return
|
||
|
|
else:
|
||
|
|
raise exc_exception
|