Coverage for bzfs_main/util/utils.py: 100%
782 statements
« prev ^ index » next coverage.py v7.15.1, created at 2026-07-21 12:39 +0000
1# Copyright 2024 Wolfgang Hoschek AT mac DOT com
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14#
15"""Collection of helper functions used across bzfs; includes environment variable parsing, process management and lightweight
16concurrency primitives, etc.
18Everything in this module relies only on the Python standard library so other modules remain dependency free. Each utility
19favors simple, predictable behavior on all supported platforms.
20"""
22from __future__ import (
23 annotations,
24)
25import argparse
26import base64
27import bisect
28import contextlib
29import dataclasses
30import errno
31import hashlib
32import itertools
33import logging
34import operator
35import os
36import platform
37import pwd
38import random
39import re
40import signal
41import stat
42import subprocess
43import sys
44import threading
45import time
46import types
47from collections import (
48 defaultdict,
49 deque,
50)
51from collections.abc import (
52 ItemsView,
53 Iterable,
54 Iterator,
55 Sequence,
56)
57from concurrent.futures import (
58 Executor,
59 Future,
60)
61from concurrent.futures.thread import (
62 ThreadPoolExecutor,
63)
64from dataclasses import (
65 dataclass,
66)
67from datetime import (
68 datetime,
69 timedelta,
70 timezone,
71 tzinfo,
72)
73from subprocess import (
74 DEVNULL,
75 PIPE,
76)
77from typing import (
78 IO,
79 Any,
80 Callable,
81 Final,
82 Generic,
83 Literal,
84 NoReturn,
85 Protocol,
86 SupportsIndex,
87 TextIO,
88 TypeVar,
89 cast,
90 final,
91)
93# constants:
94PROG_NAME: Final[str] = "bzfs"
95ENV_VAR_PREFIX: Final[str] = PROG_NAME + "_"
96DIE_STATUS: Final[int] = 3
97DESCENDANTS_RE_SUFFIX: Final[str] = r"(?:/.*)?" # also match descendants of a matching dataset
98LOG_STDERR: Final[int] = (logging.INFO + logging.WARNING) // 2 # custom log level is halfway in between
99LOG_STDOUT: Final[int] = (LOG_STDERR + logging.INFO) // 2 # custom log level is halfway in between
100LOG_DEBUG: Final[int] = logging.DEBUG
101LOG_TRACE: Final[int] = logging.DEBUG // 2 # custom log level is halfway in between
102YEAR_WITH_FOUR_DIGITS_REGEX: Final[re.Pattern] = re.compile(r"[1-9][0-9][0-9][0-9]") # empty shall not match nonempty target
103UNIX_TIME_INFINITY_SECS: Final[int] = 2**64 # billions of years and to be extra safe, larger than the largest ZFS GUID
104DONT_SKIP_DATASET: Final[str] = ""
105SHELL_CHARS: Final[str] = '"' + "'`~!@#$%^&*()+={}[]|;<>?,\\" # intentionally not included: -_.:/
106SHELL_CHARS_AND_SLASH: Final[str] = SHELL_CHARS + "/"
107FILE_PERMISSIONS: Final[int] = stat.S_IRUSR | stat.S_IWUSR # rw------- (user read + write)
108DIR_PERMISSIONS: Final[int] = stat.S_IRWXU # rwx------ (user read + write + execute)
109UMASK: Final[int] = (~DIR_PERMISSIONS) & 0o777 # so intermediate dirs created by os.makedirs() have stricter permissions
110UNIX_DOMAIN_SOCKET_PATH_MAX_LENGTH: Final[int] = 107 if platform.system() == "Linux" else 103 # see Google for 'sun_path'
112RegexList = list[tuple[re.Pattern[str], bool]] # Type alias
115def getenv_any(key: str, default: str | None = None, env_var_prefix: str = ENV_VAR_PREFIX) -> str | None:
116 """All shell environment variable names used for configuration start with this prefix."""
117 return os.getenv(env_var_prefix + key, default)
120def getenv_int(key: str, default: int, env_var_prefix: str = ENV_VAR_PREFIX) -> int:
121 """Returns environment variable ``key`` as int with ``default`` fallback."""
122 return int(cast(str, getenv_any(key, default=str(default), env_var_prefix=env_var_prefix)))
125def getenv_bool(key: str, default: bool = False, env_var_prefix: str = ENV_VAR_PREFIX) -> bool:
126 """Returns environment variable ``key`` as bool with ``default`` fallback."""
127 return cast(str, getenv_any(key, default=str(default), env_var_prefix=env_var_prefix)).lower().strip() == "true"
130def cut(field: int, separator: str = "\t", *, lines: list[str]) -> list[str]:
131 """Retains only column number 'field' in a list of TSV/CSV lines; Analog to Unix 'cut' CLI command."""
132 assert lines is not None
133 assert isinstance(lines, list)
134 assert len(separator) == 1
135 if field == 1:
136 return [line[: line.index(separator)] for line in lines]
137 elif field == 2:
138 return [line[line.index(separator) + 1 :] for line in lines]
139 else:
140 raise ValueError(f"Invalid field value: {field}")
143def drain(iterable: Iterable[Any]) -> None:
144 """Consumes all items in the iterable, effectively draining it."""
145 for _ in iterable:
146 del _ # help gc (iterable can block)
149_K_ = TypeVar("_K_")
150_V_ = TypeVar("_V_")
151_R_ = TypeVar("_R_")
154def shuffle_dict(dictionary: dict[_K_, _V_], /, rand: random.Random = random.SystemRandom()) -> dict[_K_, _V_]: # noqa: B008
155 """Returns a new dict with items shuffled randomly."""
156 items: list[tuple[_K_, _V_]] = list(dictionary.items())
157 rand.shuffle(items)
158 return dict(items)
161def sorted_dict(
162 dictionary: dict[_K_, _V_], /, *, key: Callable[[tuple[_K_, _V_]], Any] | None = None, reverse: bool = False
163) -> dict[_K_, _V_]:
164 """Returns a new dict with items sorted, primarily by key and secondarily by value (unless a custom key is supplied)."""
165 return dict(sorted(dictionary.items(), key=key, reverse=reverse))
168def tail(file: str, *, n: int, errors: str | None = None) -> Sequence[str]:
169 """Return the last ``n`` lines of ``file`` without following symlinks."""
170 if not os.path.isfile(file):
171 return []
172 with open_nofollow(file, "r", encoding="utf-8", errors=errors, check_owner=False, check_perm=False) as fd:
173 return deque(fd, maxlen=n)
176_NAMED_CAPTURING_GROUP: Final[re.Pattern[str]] = re.compile(r"^" + re.escape("(?P<") + r"[^\W\d]\w*" + re.escape(">"))
177_NUMERIC_BACKREFERENCE_REGEX: Final[re.Pattern[str]] = re.compile(r"\\\d+") # example: \1
180def replace_capturing_groups_with_non_capturing_groups(regex: str) -> str:
181 """Replaces regex capturing groups with non-capturing groups for better matching performance (unless it's tricky).
183 Unnamed capturing groups example: '(.*/)?tmp(foo|bar)(?!public)\\(' --> '(?:.*/)?tmp(?:foo|bar)(?!public)\\('
184 Aka replaces parenthesis '(' followed by a char other than question mark '?', but not preceded by a backslash
185 with the replacement string '(?:'
187 Named capturing group example: '(?P<name>abc)' --> '(?:abc)'
188 Aka replaces '(?P<' followed by a valid name followed by '>', but not preceded by a backslash
189 with the replacement string '(?:'
191 Also see https://docs.python.org/3/howto/regex.html#non-capturing-and-named-groups
192 """
193 i = regex.find("[")
194 if i >= 0 and regex.find("(", i) >= 0:
195 # Conservative fallback to minimize code complexity: skip the rewrite entirely in the case where the regex might
196 # contain a regex character class that contains parenthesis.
197 # Rewriting a regex is a performance optimization; correctness comes first.
198 return regex
200 if "(?P=" in regex or "(?(" in regex or _NUMERIC_BACKREFERENCE_REGEX.search(regex):
201 # Conservative fallback to minimize code complexity: skip the rewrite entirely if the regex might contain a
202 # (named or conditional or numeric) backreference.
203 # Rewriting a regex is a performance optimization; correctness comes first.
204 return regex
206 i = len(regex) - 2
207 while i >= 0:
208 i = regex.rfind("(", 0, i + 1)
209 if i >= 0 and (i == 0 or regex[i - 1] != "\\"):
210 if regex[i + 1] != "?":
211 regex = f"{regex[0:i]}(?:{regex[i + 1:]}" # unnamed capturing group
212 else: # potentially a valid named capturing group
213 regex = regex[0:i] + _NAMED_CAPTURING_GROUP.sub(repl="(?:", string=regex[i:], count=1)
214 i -= 1
215 return regex
218def get_home_directory() -> str:
219 """Reliably detects home dir without using HOME env var."""
220 # thread-safe version of: os.environ.pop('HOME', None); os.path.expanduser('~')
221 return pwd.getpwuid(os.getuid()).pw_dir
224def human_readable_bytes(num_bytes: float, *, separator: str = " ", precision: int | None = None) -> str:
225 """Formats 'num_bytes' as a human-readable size; for example "567 MiB"."""
226 sign = "-" if num_bytes < 0 else ""
227 s = abs(num_bytes)
228 units = ("B", "KiB", "MiB", "GiB", "TiB", "PiB", "EiB", "ZiB", "YiB", "RiB", "QiB")
229 n = len(units) - 1
230 i = 0
231 while s >= 1024 and i < n:
232 s /= 1024
233 i += 1
234 formatted_num = human_readable_float(s) if precision is None else f"{s:.{precision}f}"
235 return f"{sign}{formatted_num}{separator}{units[i]}"
238def human_readable_duration(duration: float, *, unit: str = "ns", separator: str = "", precision: int | None = None) -> str:
239 """Formats a duration in human units, automatically scaling as needed; for example "567ms"."""
240 sign = "-" if duration < 0 else ""
241 t = abs(duration)
242 units = ("ns", "μs", "ms", "s", "m", "h", "d")
243 i = units.index(unit)
244 if t < 1 and t != 0:
245 nanos = (1, 1_000, 1_000_000, 1_000_000_000, 60 * 1_000_000_000, 60 * 60 * 1_000_000_000, 3600 * 24 * 1_000_000_000)
246 t *= nanos[i]
247 i = 0
248 while t >= 1000 and i < 3:
249 t /= 1000
250 i += 1
251 if i >= 3:
252 while t >= 60 and i < 5:
253 t /= 60
254 i += 1
255 if i >= 5:
256 while t >= 24 and i < len(units) - 1:
257 t /= 24
258 i += 1
259 formatted_num = human_readable_float(t) if precision is None else f"{t:.{precision}f}"
260 return f"{sign}{formatted_num}{separator}{units[i]}"
263def human_readable_float(number: float) -> str:
264 """Formats ``number`` with a variable precision depending on magnitude.
266 This design mirrors the way humans round values when scanning logs.
268 If the number has one digit before the decimal point (0 <= abs(number) < 10):
269 Round and use two decimals after the decimal point (e.g., 3.14559 --> "3.15").
271 If the number has two digits before the decimal point (10 <= abs(number) < 100):
272 Round and use one decimal after the decimal point (e.g., 12.36 --> "12.4").
274 If the number has three or more digits before the decimal point (abs(number) >= 100):
275 Round and use zero decimals after the decimal point (e.g., 123.556 --> "124").
277 Ensures no unnecessary trailing zeroes are retained: Example: 1.500 --> "1.5", 1.00 --> "1"
278 """
279 abs_number = abs(number)
280 precision = 2 if abs_number < 10 else 1 if abs_number < 100 else 0
281 if precision == 0:
282 return str(round(number))
283 result = f"{number:.{precision}f}"
284 assert "." in result
285 result = result.rstrip("0").rstrip(".") # Remove trailing zeros and trailing decimal point if empty
286 if result == "-0":
287 result = "0"
288 return result
291def percent(number: int, total: int, *, print_total: bool = False) -> str:
292 """Returns percentage string of ``number`` relative to ``total``."""
293 tot: str = f"/{total}" if print_total else ""
294 return f"{number}{tot}={'inf' if total == 0 else human_readable_float(100 * number / total)}%"
297def open_nofollow(
298 path: str,
299 mode: str = "r",
300 buffering: int = -1,
301 encoding: str | None = None,
302 errors: str | None = None,
303 newline: str | None = None,
304 *,
305 perm: int = FILE_PERMISSIONS,
306 check_owner: bool = True,
307 check_perm: bool = True,
308 **kwargs: Any,
309) -> IO[Any]:
310 """Behaves exactly like built-in open(), except that it refuses to follow symlinks, i.e. raises OSError with
311 errno.ELOOP/EMLINK if basename of path is a symlink.
313 Also, can specify custom permissions on O_CREAT, and verify secure ownership and mode bits.
315 If check_owner=True, write-capable opens require ownership by the effective UID; read-only opens also allow ownership by
316 uid 0 (root). This allows safe reads of root-owned system files while preventing writes to files not owned by the caller.
318 If check_perm=True, the file must also not be writable by group or others, and must not be executable by anyone.
319 """
320 if not mode:
321 raise ValueError("Must have exactly one of create/read/write/append mode and at most one plus")
322 flags = {
323 "r": os.O_RDONLY,
324 "w": os.O_WRONLY | os.O_CREAT | os.O_TRUNC,
325 "a": os.O_WRONLY | os.O_CREAT | os.O_APPEND,
326 "x": os.O_WRONLY | os.O_CREAT | os.O_EXCL,
327 }.get(mode[0])
328 if flags is None:
329 raise ValueError(f"invalid mode {mode!r}")
330 if "+" in mode: # enable read-write access for r+, w+, a+, x+
331 flags = (flags & ~os.O_WRONLY) | os.O_RDWR # clear os.O_WRONLY and set os.O_RDWR while preserving all other flags
332 flags |= os.O_NOFOLLOW | os.O_CLOEXEC
333 fd: int = os.open(path, flags=flags, mode=perm)
334 try:
335 stats: os.stat_result | None = None
336 if check_owner:
337 stats = os.fstat(fd)
338 st_uid: int = stats.st_uid
339 if st_uid != os.geteuid(): # verify ownership is current effective UID
340 if (flags & (os.O_WRONLY | os.O_RDWR)) != 0: # require that writer owns the file
341 raise PermissionError(errno.EPERM, f"{path!r} is owned by uid {st_uid}, not {os.geteuid()}", path)
342 elif st_uid != 0: # it's ok for root to own a file that we'll merely read
343 raise PermissionError(errno.EPERM, f"{path!r} is owned by uid {st_uid}, not {os.geteuid()} or 0", path)
344 if check_perm:
345 stats = os.fstat(fd) if stats is None else stats
346 if (stats.st_mode & (stat.S_IWGRP | stat.S_IWOTH | stat.S_IXUSR | stat.S_IXGRP | stat.S_IXOTH)) != 0:
347 st_mode: int = stat.S_IMODE(stats.st_mode)
348 raise PermissionError(
349 errno.EPERM,
350 f"{path!r} has permissions {st_mode:03o} aka {stat.filemode(st_mode)[1:]} "
351 "but must not be writable by group or others, or executable by anyone",
352 path,
353 )
354 return os.fdopen(fd, mode, buffering=buffering, encoding=encoding, errors=errors, newline=newline, **kwargs)
355 except Exception:
356 try:
357 os.close(fd)
358 except OSError:
359 pass
360 raise
363def close_quietly(fd: int) -> None:
364 """Closes the given file descriptor while silently swallowing any OSError that might arise as part of this."""
365 if fd >= 0:
366 try:
367 os.close(fd)
368 except OSError:
369 pass
372_P = TypeVar("_P")
375def find_match(
376 seq: Sequence[_P],
377 predicate: Callable[[_P], bool],
378 start: SupportsIndex | None = None,
379 end: SupportsIndex | None = None,
380 *,
381 reverse: bool = False,
382 raises: bool | object | Callable[[], object] = False, # raises: bool | object | Callable = False, # python >= 3.10
383) -> int:
384 """Returns the integer index within ``seq`` of the first item (or last item if reverse=True) that matches the given
385 predicate condition.
387 If no matching item is found returns -1 or ValueError, depending on the ``raises`` parameter, which is a bool indicating
388 whether to raise an error, or an object containing the error message, but can also be a Callable/lambda in order to
389 support efficient deferred generation of error messages.
391 Analog to ``str.find()``, including slicing semantics with parameters start and end, i.e. respects Python slicing
392 semantics for start/end (including clamping). For example, seq can be a list, tuple or str.
394 Example usage:
395 lst = ["a", "b", "-c", "d"]
396 i = find_match(lst, lambda arg: arg.startswith("-"), start=1, end=3, reverse=True)
397 if i >= 0:
398 print(lst[i])
399 i = find_match(lst, lambda arg: arg.startswith("-"), raises=f"Tag {tag} not found in {file}")
400 i = find_match(lst, lambda arg: arg.startswith("-"), raises=lambda: f"Tag {tag} not found in {file}")
401 """
402 if start is None and end is None:
403 for i in range(len(seq) - 1, -1, -1) if reverse else range(len(seq)):
404 if predicate(seq[i]):
405 return i
406 else:
407 slice_start, slice_end, _ = slice(start, end).indices(len(seq))
408 for i in range(slice_end - 1, slice_start - 1, -1) if reverse else range(slice_start, slice_end):
409 if predicate(seq[i]):
410 return i
411 if raises is False or raises is None:
412 return -1
413 if raises is True:
414 raise ValueError("No matching item found in sequence")
415 if callable(raises):
416 raises = raises()
417 raise ValueError(raises)
420def is_descendant(dataset: str, of_root_dataset: str) -> bool:
421 """Returns True if ZFS ``dataset`` lies under ``of_root_dataset`` in the dataset hierarchy, or is the same."""
422 return dataset == of_root_dataset or dataset.startswith(of_root_dataset + "/")
425def has_duplicates(sorted_list: list[Any]) -> bool:
426 """Returns True if any adjacent items within the given sorted sequence are equal."""
427 return any(map(operator.eq, sorted_list, itertools.islice(sorted_list, 1, None)))
430def dry(msg: str, is_dry_run: bool) -> str:
431 """Prefix ``msg`` with 'Dry' when in dry-run mode."""
432 return "Dry " + msg if is_dry_run else msg
435def relativize_dataset(dataset: str, root_dataset: str) -> str:
436 """Converts an absolute dataset path to one relative to ``root_dataset``.
438 Example: root_dataset=tank/foo, dataset=tank/foo/bar/baz --> relative_path=/bar/baz.
439 """
440 return dataset[len(root_dataset) :]
443def dataset_paths(dataset: str) -> Iterator[str]:
444 """Enumerates all paths of a valid ZFS dataset name; Example: "a/b/c" --> yields "a", "a/b", "a/b/c"."""
445 i: int = 0
446 while i >= 0:
447 i = dataset.find("/", i)
448 if i < 0:
449 yield dataset
450 else:
451 yield dataset[:i]
452 i += 1
455def replace_prefix(s: str, old_prefix: str, new_prefix: str) -> str:
456 """In a string s, replaces a leading old_prefix string with new_prefix; assumes the leading string is present."""
457 assert s.startswith(old_prefix)
458 return new_prefix + s[len(old_prefix) :]
461def replace_in_lines(lines: list[str], old: str, new: str, count: int = -1) -> None:
462 """Replaces ``old`` with ``new`` in-place for every string in ``lines``."""
463 for i in range(len(lines)):
464 lines[i] = lines[i].replace(old, new, count)
467_TAPPEND = TypeVar("_TAPPEND")
470def append_if_absent(lst: list[_TAPPEND], *items: _TAPPEND) -> list[_TAPPEND]:
471 """Appends items to list if they are not already present."""
472 for item in items:
473 if item not in lst:
474 lst.append(item)
475 return lst
478def xappend(lst: list[_TAPPEND], *items: _TAPPEND | Iterable[_TAPPEND]) -> list[_TAPPEND]:
479 """Appends each of the items to the given list if the item is "truthy", for example not None and not an empty string; If
480 an item is an iterable does so recursively, flattening the output."""
481 for item in items:
482 if isinstance(item, str) or not isinstance(item, Iterable):
483 if item:
484 lst.append(item)
485 else:
486 xappend(lst, *item)
487 return lst
490def is_included(name: str, include_regexes: RegexList, exclude_regexes: RegexList) -> bool:
491 """Returns True if the name matches at least one of the include regexes but none of the exclude regexes; else False.
493 A regex that starts with a `!` is a negation - the regex matches if the regex without the `!` prefix does not match.
494 """
495 for regex, is_negation in exclude_regexes:
496 is_match = regex.fullmatch(name) if regex.pattern != ".*" else True
497 if is_negation:
498 is_match = not is_match
499 if is_match:
500 return False
502 for regex, is_negation in include_regexes:
503 is_match = regex.fullmatch(name) if regex.pattern != ".*" else True
504 if is_negation:
505 is_match = not is_match
506 if is_match:
507 return True
509 return False
512def compile_regexes(regexes: list[str], *, suffix: str = "") -> RegexList:
513 """Compiles regex strings and keeps track of negations."""
514 assert isinstance(regexes, list)
515 compiled_regexes: RegexList = []
516 for regex in regexes:
517 if suffix: # disallow non-trailing end-of-str symbol in dataset regexes to ensure descendants will also match
518 if regex.endswith("\\$"):
519 pass # trailing literal $ is ok
520 elif regex.endswith("$"):
521 regex = regex[0:-1] # ok because all users of compile_regexes() call re.fullmatch()
522 elif "$" in regex:
523 raise re.error("Must not use non-trailing '$' character", regex)
524 if is_negation := regex.startswith("!"):
525 regex = regex[1:]
526 regex = replace_capturing_groups_with_non_capturing_groups(regex)
527 if regex != ".*" or not (suffix.startswith("(") and suffix.endswith(")?")):
528 regex = f"{regex}{suffix}"
529 compiled_regexes.append((re.compile(regex), is_negation))
530 return compiled_regexes
533def list_formatter(iterable: Iterable[Any], separator: str = " ", lstrip: bool = False) -> Any:
534 """Lazy formatter joining items with ``separator`` used to avoid overhead in disabled log levels."""
536 @final
537 class CustomListFormatter:
538 """Formatter object that joins items when converted to ``str``."""
540 def __str__(self) -> str:
541 s = separator.join(map(str, iterable))
542 return s.lstrip() if lstrip else s
544 return CustomListFormatter()
547def pretty_print_formatter(obj_to_format: Any) -> Any:
548 """Lazy pprint formatter used to avoid overhead in disabled log levels."""
550 @final
551 class PrettyPrintFormatter:
552 """Formatter that pretty-prints the object on conversion to ``str``."""
554 def __str__(self) -> str:
555 import pprint # lazy import for startup perf
557 return pprint.pformat(vars(obj_to_format))
559 return PrettyPrintFormatter()
562def stderr_to_str(stderr: Any) -> str:
563 """Workaround for https://github.com/python/cpython/issues/87597."""
564 return str(stderr) if not isinstance(stderr, bytes) else stderr.decode("utf-8", errors="replace")
567def xprint(log: logging.Logger, value: Any, *, run: bool = True, end: str = "\n", file: TextIO | None = None) -> None:
568 """Optionally logs ``value`` at stdout/stderr level."""
569 if run and value:
570 value = value if end else str(value).rstrip()
571 level = LOG_STDOUT if file is sys.stdout else LOG_STDERR
572 log.log(level, "%s", value)
575def sha256_hex(text: str) -> str:
576 """Returns the sha256 hex string for the given text."""
577 return hashlib.sha256(text.encode()).hexdigest()
580def sha256_urlsafe_base64(text: str, *, padding: bool = True) -> str:
581 """Returns the URL-safe base64-encoded sha256 value for the given text."""
582 digest: bytes = hashlib.sha256(text.encode()).digest()
583 s: str = base64.urlsafe_b64encode(digest).decode()
584 return s if padding else s.rstrip("=")
587def sha256_128_urlsafe_base64(text: str) -> str:
588 """Returns the left half portion of the unpadded URL-safe base64-encoded sha256 value for the given text."""
589 s: str = sha256_urlsafe_base64(text, padding=False)
590 return s[: len(s) // 2]
593def sha256_85_urlsafe_base64(text: str) -> str:
594 """Returns the left one third portion of the unpadded URL-safe base64-encoded sha256 value for the given text."""
595 s: str = sha256_urlsafe_base64(text, padding=False)
596 return s[: len(s) // 3]
599def urlsafe_base64(
600 value: int, max_value: int = 2**64 - 1, *, padding: bool = True, byteorder: Literal["little", "big"] = "big"
601) -> str:
602 """Returns the URL-safe base64 string encoding of the int value, assuming it is contained in the range [0..max_value]."""
603 assert 0 <= value <= max_value
604 max_bytes: int = (max_value.bit_length() + 7) // 8
605 value_bytes: bytes = value.to_bytes(max_bytes, byteorder)
606 s: str = base64.urlsafe_b64encode(value_bytes).decode()
607 return s if padding else s.rstrip("=")
610def die(msg: str, exit_code: int = DIE_STATUS, parser: argparse.ArgumentParser | None = None) -> NoReturn:
611 """Exits the program with ``exit_code`` after logging ``msg``."""
612 if parser is None:
613 ex = SystemExit(msg)
614 ex.code = exit_code
615 raise ex
616 else:
617 parser.error(msg)
620def subprocess_run(*args: Any, **kwargs: Any) -> subprocess.CompletedProcess:
621 """Drop-in replacement for subprocess.run() that mimics its behavior except it enhances cleanup on TimeoutExpired, and
622 provides optional child PID tracking, and optional logging of execution status via ``log`` and ``loglevel`` params."""
623 input_value = kwargs.pop("input", None)
624 timeout = kwargs.pop("timeout", None)
625 check = kwargs.pop("check", False)
626 subprocesses: Subprocesses | None = kwargs.pop("subprocesses", None)
627 if input_value is not None:
628 if kwargs.get("stdin") is not None:
629 raise ValueError("input and stdin are mutually exclusive")
630 kwargs["stdin"] = subprocess.PIPE
632 log: logging.Logger | None = kwargs.pop("log", None)
633 loglevel: int | None = kwargs.pop("loglevel", None)
634 start_time_nanos: int = time.monotonic_ns()
635 is_timeout: bool = False
636 is_cancel: bool = False
637 exitcode: int | None = None
639 def log_status() -> None:
640 if log is not None:
641 _loglevel: int = loglevel if loglevel is not None else getenv_int("subprocess_run_loglevel", LOG_TRACE)
642 if log.isEnabledFor(_loglevel):
643 elapsed_time: str = human_readable_float((time.monotonic_ns() - start_time_nanos) / 1_000_000) + "ms"
644 status = "cancel" if is_cancel else "timeout" if is_timeout else "success" if exitcode == 0 else "failure"
645 cmd = kwargs["args"] if "args" in kwargs else (args[0] if args else None)
646 cmd_str: str = " ".join(str(arg) for arg in iter(cmd)) if isinstance(cmd, (list, tuple)) else str(cmd)
647 log.log(_loglevel, f"Executed [{status}] [{elapsed_time}]: %s", cmd_str)
649 with xfinally(log_status):
650 ctx: contextlib.AbstractContextManager[subprocess.Popen]
651 if subprocesses is None:
652 ctx = subprocess.Popen(*args, **kwargs)
653 else:
654 ctx = subprocesses.popen_and_track(*args, **kwargs)
655 with ctx as proc:
656 try:
657 sp = subprocesses
658 if sp is not None and sp._is_terminated(): # noqa: SLF001 pylint: disable=protected-access
659 is_cancel = True
660 timeout = 0.0
661 stdout, stderr = proc.communicate(input_value, timeout=timeout)
662 except BaseException as e:
663 try:
664 if isinstance(e, subprocess.TimeoutExpired):
665 is_timeout = True
666 terminate_process_subtree(root_pids=[proc.pid]) # send SIGTERM to child proc and descendants
667 finally:
668 proc.kill()
669 raise
670 else:
671 exitcode = proc.poll()
672 assert exitcode is not None
673 if check and exitcode:
674 raise subprocess.CalledProcessError(exitcode, proc.args, output=stdout, stderr=stderr)
675 return subprocess.CompletedProcess(proc.args, exitcode, stdout, stderr)
678def terminate_process_subtree(
679 *, except_current_process: bool = True, root_pids: list[int] | None = None, sig: signal.Signals = signal.SIGTERM
680) -> None:
681 """For each root PID: Sends the given signal to the root PID and all its descendant processes."""
682 current_pid: int = os.getpid()
683 root_pids = [current_pid] if root_pids is None else root_pids
684 all_pids: list[list[int]] = _get_descendant_processes(root_pids)
685 assert len(all_pids) == len(root_pids)
686 for i, pids in enumerate(all_pids):
687 root_pid = root_pids[i]
688 if root_pid == current_pid:
689 pids += [] if except_current_process else [current_pid]
690 else:
691 pids.insert(0, root_pid)
692 for pid in pids:
693 with contextlib.suppress(OSError):
694 os.kill(pid, sig)
697def _get_descendant_processes(root_pids: list[int]) -> list[list[int]]:
698 """For each root PID, returns the list of all descendant process IDs for the given root PID, on POSIX systems."""
699 if len(root_pids) == 0:
700 return []
701 cmd: list[str] = ["ps", "-Ao", "pid,ppid"]
702 try:
703 lines: list[str] = subprocess.run(cmd, stdin=DEVNULL, stdout=PIPE, text=True, check=True).stdout.splitlines()
704 except PermissionError:
705 # degrade gracefully in sandbox environments that deny executing `ps` entirely
706 return [[] for _ in root_pids]
707 procs: dict[int, list[int]] = defaultdict(list)
708 for line in lines[1:]: # all lines except the header line
709 splits: list[str] = line.split()
710 assert len(splits) == 2
711 pid = int(splits[0])
712 ppid = int(splits[1])
713 procs[ppid].append(pid)
715 def recursive_append(ppid: int, descendants: list[int]) -> None:
716 """Recursively collect descendant PIDs starting from ``ppid``."""
717 for child_pid in procs[ppid]:
718 descendants.append(child_pid)
719 recursive_append(child_pid, descendants)
721 all_descendants: list[list[int]] = []
722 for root_pid in root_pids:
723 descendants: list[int] = []
724 recursive_append(root_pid, descendants)
725 all_descendants.append(descendants)
726 return all_descendants
729@contextlib.contextmanager
730def termination_signal_handler(
731 termination_events: list[threading.Event],
732 *,
733 termination_handler: Callable[[], None] = lambda: terminate_process_subtree(),
734) -> Iterator[None]:
735 """Context manager that installs SIGINT/SIGTERM handlers that set all ``termination_events`` and, by default, terminate
736 all descendant processes."""
737 termination_events = list(termination_events) # shallow copy
739 def _handler(_sig: int, _frame: object) -> None:
740 for event in termination_events:
741 event.set()
742 termination_handler()
744 previous_int_handler = signal.signal(signal.SIGINT, _handler) # install new signal handler
745 previous_term_handler = signal.signal(signal.SIGTERM, _handler) # install new signal handler
746 try:
747 yield # run body of context manager
748 finally:
749 signal.signal(signal.SIGINT, previous_int_handler) # restore original signal handler
750 signal.signal(signal.SIGTERM, previous_term_handler) # restore original signal handler
753def return_false() -> bool:
754 """Always returns ``False``; picklable."""
755 return False
758def sleep_nanos(delay_nanos: int) -> None:
759 """Same as time.sleep() but expects a relative sleep duration in nanoseconds as input value; picklable."""
760 time.sleep(delay_nanos / 1_000_000_000)
763#############################################################################
764@dataclass(frozen=True)
765@final
766class TaskTiming:
767 """Customizable callbacks for reading the current monotonic time, sleeping and optional async termination; immutable."""
769 monotonic_ns: Callable[[], int] = time.monotonic_ns
771 is_terminated: Callable[[], bool] = return_false
772 """Returns whether a predicate has become true; can be used to indicate system shutdown or similar cancellation
773 conditions; default is to always return ``False``."""
775 sleep: Callable[[int], None] = sleep_nanos
776 """Sleeps N nanoseconds; thread-safe."""
778 def copy(self, **override_kwargs: Any) -> TaskTiming:
779 """Creates a new object copying an existing one with the specified fields overridden for customization; thread-
780 safe."""
781 return dataclasses.replace(self, **override_kwargs)
783 @staticmethod
784 def make_from(termination_event: threading.Event | None) -> TaskTiming:
785 """Convenience factory that creates an object that performs async termination when ``termination_event`` is set."""
786 if termination_event is None:
787 return TaskTiming()
789 def _sleep(delay_nanos: int) -> None:
790 termination_event.wait(delay_nanos / 1_000_000_000) # allow early wakeup on async termination
792 return TaskTiming(is_terminated=termination_event.is_set, sleep=_sleep)
795#############################################################################
796@final
797class Subprocesses:
798 """Provides per-job tracking of child PIDs so a job can safely terminate only the subprocesses it spawned itself; used
799 when multiple jobs run concurrently within the same Python process.
801 Optionally binds to an ``_is_terminated`` predicate to enforce async cancellation by forcing immediate timeouts for newly
802 spawned subprocesses once cancellation is requested.
803 """
805 def __init__(self, is_terminated: Callable[[], bool] = return_false) -> None:
806 self._is_terminated: Final[Callable[[], bool]] = is_terminated
807 self._lock: Final[threading.Lock] = threading.Lock()
808 self._child_pids: Final[dict[int, None]] = {} # a set that preserves insertion order
810 @contextlib.contextmanager
811 def popen_and_track(self, *popen_args: Any, **popen_kwargs: Any) -> Iterator[subprocess.Popen]:
812 """Context manager that calls subprocess.Popen() and tracks the child PID for per-job termination.
814 Holds a lock across Popen+PID registration to prevent a race when terminate_process_subtrees() is invoked (e.g. from
815 SIGINT/SIGTERM handlers), ensuring newly spawned child processes cannot escape termination. The child PID is
816 unregistered on context exit.
817 """
818 with self._lock:
819 proc: subprocess.Popen = subprocess.Popen(*popen_args, **popen_kwargs)
820 self._child_pids[proc.pid] = None
821 try:
822 yield proc
823 finally:
824 with self._lock:
825 self._child_pids.pop(proc.pid, None)
827 def subprocess_run(self, *args: Any, **kwargs: Any) -> subprocess.CompletedProcess:
828 """Wrapper around utils.subprocess_run() that auto-registers/unregisters child PIDs for per-job termination."""
829 return subprocess_run(*args, **kwargs, subprocesses=self)
831 def terminate_process_subtrees(self, sig: signal.Signals = signal.SIGTERM) -> None:
832 """Sends the given signal to all tracked child PIDs and their descendants, ignoring errors for dead PIDs."""
833 with self._lock:
834 pids: list[int] = list(self._child_pids)
835 self._child_pids.clear()
836 terminate_process_subtree(root_pids=pids, sig=sig)
839#############################################################################
840def pid_exists(pid: int) -> bool | None:
841 """Returns True if a process with PID exists, False if not, or None on error."""
842 if pid <= 0:
843 return False
844 try: # with signal=0, no signal is actually sent, but error checking is still performed
845 os.kill(pid, 0) # ... which can be used to check for process existence on POSIX systems
846 except OSError as err:
847 if err.errno == errno.ESRCH: # No such process
848 return False
849 if err.errno == errno.EPERM: # Operation not permitted
850 return True
851 return None
852 return True
855def nprefix(s: str) -> str:
856 """Returns a canonical snapshot prefix with trailing underscore."""
857 return sys.intern(s + "_")
860def ninfix(s: str) -> str:
861 """Returns a canonical infix with trailing underscore when not empty."""
862 return sys.intern(s + "_") if s else ""
865def nsuffix(s: str) -> str:
866 """Returns a canonical suffix with leading underscore when not empty."""
867 return sys.intern("_" + s) if s else ""
870def format_dict(dictionary: dict[Any, Any]) -> str:
871 """Returns a formatted dictionary using repr for consistent output."""
872 return f'"{dictionary}"'
875def format_obj(obj: object) -> str:
876 """Returns a formatted str using repr for consistent output."""
877 return f'"{obj}"'
880def validate_dataset_name(dataset: str, input_text: str) -> None:
881 """'zfs create' CLI does not accept dataset names that are empty or start or end in a slash, etc."""
882 # Also see https://github.com/openzfs/zfs/issues/439#issuecomment-2784424
883 # and https://github.com/openzfs/zfs/issues/8798
884 # and (by now no longer accurate): https://docs.oracle.com/cd/E26505_01/html/E37384/gbcpt.html
885 invalid_chars: str = SHELL_CHARS
886 if (
887 dataset in ("", ".", "..")
888 or dataset.startswith(("/", "./", "../"))
889 or dataset.endswith(("/", "/.", "/.."))
890 or any(substring in dataset for substring in ("//", "/./", "/../"))
891 or any(char in invalid_chars or (char.isspace() and char != " ") for char in dataset)
892 or not dataset[0].isalpha()
893 ):
894 die(f"Invalid ZFS dataset name: '{dataset}' for: '{input_text}'")
897def validate_property_name(propname: str, input_text: str) -> str:
898 """Checks that the ZFS property name contains no spaces or shell chars, etc."""
899 invalid_chars: str = SHELL_CHARS
900 if (not propname) or propname.startswith("-") or any(char.isspace() or char in invalid_chars for char in propname):
901 die(f"Invalid ZFS property name: '{propname}' for: '{input_text}'")
902 return propname
905def validate_is_not_a_symlink(msg: str, path: str, parser: argparse.ArgumentParser | None = None) -> None:
906 """Checks that the given path is not a symbolic link."""
907 if os.path.islink(path):
908 die(f"{msg}must not be a symlink: {path}", parser=parser)
911def validate_file_permissions(path: str, mode: int) -> None:
912 """Verify permissions and that ownership is current effective UID."""
913 stats: os.stat_result = os.stat(path, follow_symlinks=False)
914 st_uid: int = stats.st_uid
915 if st_uid != os.geteuid(): # verify ownership is current effective UID
916 die(f"{path!r} is owned by uid {st_uid}, not {os.geteuid()}")
917 st_mode = stat.S_IMODE(stats.st_mode)
918 if st_mode != mode:
919 die(
920 f"{path!r} has permissions {st_mode:03o} aka {stat.filemode(st_mode)[1:]}, "
921 f"not {mode:03o} aka {stat.filemode(mode)[1:]})"
922 )
925def parse_duration_to_milliseconds(duration: str, *, regex_suffix: str = "", context: str = "") -> int:
926 """Parses human duration strings like '5minutes' or '2 hours' to milliseconds."""
927 unit_milliseconds: dict[str, int] = {
928 "milliseconds": 1,
929 "millis": 1,
930 "seconds": 1000,
931 "secs": 1000,
932 "minutes": 60 * 1000,
933 "mins": 60 * 1000,
934 "hours": 60 * 60 * 1000,
935 "days": 86400 * 1000,
936 "weeks": 7 * 86400 * 1000,
937 "months": round(30.5 * 86400 * 1000),
938 "years": 365 * 86400 * 1000,
939 }
940 match = re.fullmatch(
941 r"(\d+)\s*(milliseconds|millis|seconds|secs|minutes|mins|hours|days|weeks|months|years)" + regex_suffix,
942 duration,
943 )
944 if not match:
945 if context:
946 die(f"Invalid duration format: {duration} within {context}")
947 else:
948 raise ValueError(f"Invalid duration format: {duration}")
949 assert match
950 quantity: int = int(match.group(1))
951 unit: str = match.group(2)
952 return quantity * unit_milliseconds[unit]
955def unixtime_fromisoformat(datetime_str: str) -> int:
956 """Converts ISO 8601 datetime string into UTC Unix time in integer seconds."""
957 return int(datetime.fromisoformat(datetime_str).timestamp())
960def isotime_from_unixtime(unixtime_in_seconds: int) -> str:
961 """Converts UTC Unix time seconds into ISO 8601 datetime string."""
962 tz: tzinfo = timezone.utc
963 dt: datetime = datetime.fromtimestamp(unixtime_in_seconds, tz=tz)
964 return dt.isoformat(sep="_", timespec="seconds")
967def current_datetime(
968 tz_spec: str | None = None,
969 now_fn: Callable[[tzinfo | None], datetime] | None = None,
970) -> datetime:
971 """Returns current time in ``tz_spec`` timezone or local timezone."""
972 if now_fn is None:
973 now_fn = datetime.now
974 return now_fn(get_timezone(tz_spec))
977def get_timezone(tz_spec: str | None = None) -> tzinfo | None:
978 """Returns timezone from spec or local timezone if unspecified."""
979 tz: tzinfo | None
980 if tz_spec is None:
981 tz = None
982 elif tz_spec == "UTC":
983 tz = timezone.utc
984 elif match := re.fullmatch(r"([+-])(\d\d):?(\d\d)", tz_spec):
985 sign, hours, minutes = match.groups()
986 offset: int = int(hours) * 60 + int(minutes)
987 offset = -offset if sign == "-" else offset
988 tz = timezone(timedelta(minutes=offset))
989 elif "/" in tz_spec:
990 from zoneinfo import ZoneInfo # lazy import for startup perf
992 tz = ZoneInfo(tz_spec)
993 else:
994 raise ValueError(f"Invalid timezone specification: {tz_spec}")
995 return tz
998###############################################################################
999@final
1000class SnapshotPeriods: # thread-safe
1001 """Parses snapshot suffix strings and converts between durations."""
1003 def __init__(self) -> None:
1004 """Initialize lookup tables of suffixes and corresponding millis."""
1005 self.suffix_milliseconds: Final[dict[str, int]] = {
1006 "yearly": 365 * 86400 * 1000,
1007 "monthly": round(30.5 * 86400 * 1000),
1008 "weekly": 7 * 86400 * 1000,
1009 "daily": 86400 * 1000,
1010 "hourly": 60 * 60 * 1000,
1011 "minutely": 60 * 1000,
1012 "secondly": 1000,
1013 "millisecondly": 1,
1014 }
1015 self.period_labels: Final[dict[str, str]] = {
1016 "yearly": "years",
1017 "monthly": "months",
1018 "weekly": "weeks",
1019 "daily": "days",
1020 "hourly": "hours",
1021 "minutely": "minutes",
1022 "secondly": "seconds",
1023 "millisecondly": "milliseconds",
1024 }
1025 self._suffix_regex0: Final[re.Pattern] = re.compile(rf"([1-9][0-9]*)?({'|'.join(self.suffix_milliseconds.keys())})")
1026 self._suffix_regex1: Final[re.Pattern] = re.compile("_" + self._suffix_regex0.pattern)
1028 def suffix_to_duration0(self, suffix: str) -> tuple[int, str]:
1029 """Parse suffix like '10minutely' to (10, 'minutely')."""
1030 return self._suffix_to_duration(suffix, self._suffix_regex0)
1032 def suffix_to_duration1(self, suffix: str) -> tuple[int, str]:
1033 """Like :meth:`suffix_to_duration0` but expects an underscore prefix."""
1034 return self._suffix_to_duration(suffix, self._suffix_regex1)
1036 @staticmethod
1037 def _suffix_to_duration(suffix: str, regex: re.Pattern) -> tuple[int, str]:
1038 """Example: Converts '2 hourly' to (2, 'hourly') and 'hourly' to (1, 'hourly')."""
1039 if match := regex.fullmatch(suffix):
1040 duration_amount: int = int(match.group(1)) if match.group(1) else 1
1041 assert duration_amount > 0
1042 duration_unit: str = match.group(2)
1043 return duration_amount, duration_unit
1044 else:
1045 return 0, ""
1047 def label_milliseconds(self, snapshot: str) -> int:
1048 """Returns duration encoded in ``snapshot`` suffix, in milliseconds."""
1049 i = snapshot.rfind("_")
1050 snapshot = "" if i < 0 else snapshot[i + 1 :]
1051 duration_amount, duration_unit = self._suffix_to_duration(snapshot, self._suffix_regex0)
1052 return duration_amount * self.suffix_milliseconds.get(duration_unit, 0)
1055#############################################################################
1056@final
1057class JobStats:
1058 """Simple thread-safe counters summarizing job progress."""
1060 def __init__(self, jobs_all: int) -> None:
1061 assert jobs_all >= 0
1062 self.lock: Final[threading.Lock] = threading.Lock()
1063 self.jobs_all: int = jobs_all
1064 self.jobs_started: int = 0
1065 self.jobs_completed: int = 0
1066 self.jobs_failed: int = 0
1067 self.jobs_running: int = 0
1068 self.sum_elapsed_nanos: int = 0
1069 self.started_job_names: Final[set[str]] = set()
1071 def submit_job(self, job_name: str) -> str:
1072 """Counts a job submission."""
1073 with self.lock:
1074 self.jobs_started += 1
1075 self.jobs_running += 1
1076 self.started_job_names.add(job_name)
1077 return str(self)
1079 def complete_job(self, failed: bool, elapsed_nanos: int) -> str:
1080 """Counts a job completion."""
1081 assert elapsed_nanos >= 0
1082 with self.lock:
1083 self.jobs_running -= 1
1084 self.jobs_completed += 1
1085 self.jobs_failed += 1 if failed else 0
1086 self.sum_elapsed_nanos += elapsed_nanos
1087 msg = str(self)
1088 assert self.sum_elapsed_nanos >= 0, msg
1089 assert self.jobs_running >= 0, msg
1090 assert self.jobs_failed >= 0, msg
1091 assert self.jobs_failed <= self.jobs_completed, msg
1092 assert self.jobs_completed <= self.jobs_started, msg
1093 assert self.jobs_started <= self.jobs_all, msg
1094 return msg
1096 def __repr__(self) -> str:
1097 def pct(number: int) -> str:
1098 """Returns percentage string relative to total jobs."""
1099 return percent(number, total=self.jobs_all, print_total=True)
1101 al, started, completed, failed = self.jobs_all, self.jobs_started, self.jobs_completed, self.jobs_failed
1102 running = self.jobs_running
1103 t = "avg_completion_time:" + human_readable_duration(self.sum_elapsed_nanos / max(1, completed))
1104 return f"all:{al}, started:{pct(started)}, completed:{pct(completed)}, failed:{pct(failed)}, running:{running}, {t}"
1107#############################################################################
1108class Comparable(Protocol):
1109 """Partial ordering protocol."""
1111 def __lt__(self, other: Any) -> bool: ...
1114TComparable = TypeVar("TComparable", bound=Comparable) # Generic type variable for elements stored in UpdatablePriorityQueue
1117@final
1118class UpdatablePriorityQueue(Generic[TComparable]):
1119 """A priority queue that can handle updates to the priority of any element that is already contained in the queue, and
1120 does so very efficiently if there are a small number of elements in the queue (no more than thousands), as is the case
1121 for us.
1123 Could be implemented using a SortedList via https://github.com/grantjenks/python-sortedcontainers or using an indexed
1124 priority queue via
1125 https://github.com/nvictus/pqdict.
1126 But, to avoid an external dependency, is actually implemented
1127 using a simple yet effective binary search-based sorted list that can handle updates to the priority of elements that
1128 are already contained in the queue, via removal of the element, followed by update of the element, followed by
1129 (re)insertion. Duplicate elements (if any) are maintained in their order of insertion relative to other duplicates.
1130 """
1132 def __init__(self, reverse: bool = False) -> None:
1133 """Creates an empty queue; sort order flips when ``reverse`` is True."""
1134 self._lst: Final[list[TComparable]] = []
1135 self._reverse: Final[bool] = reverse
1137 def clear(self) -> None:
1138 """Removes all elements from the queue."""
1139 self._lst.clear()
1141 def push(self, element: TComparable) -> None:
1142 """Inserts ``element`` while maintaining sorted order."""
1143 bisect.insort(self._lst, element)
1145 def pop(self) -> TComparable:
1146 """Removes and returns the smallest (or largest if reverse == True) element from the queue."""
1147 return self._lst.pop() if self._reverse else self._lst.pop(0)
1149 def peek(self) -> TComparable:
1150 """Returns the smallest (or largest if reverse == True) element without removing it."""
1151 return self._lst[-1] if self._reverse else self._lst[0]
1153 def remove(self, element: TComparable) -> bool:
1154 """Removes the first occurrence (in insertion order aka FIFO) of ``element`` and returns True if it was present."""
1155 lst = self._lst
1156 i = bisect.bisect_left(lst, element)
1157 is_contained = i < len(lst) and lst[i] == element
1158 if is_contained:
1159 del lst[i] # is an optimized memmove()
1160 return is_contained
1162 def __len__(self) -> int:
1163 """Returns the number of queued elements."""
1164 return len(self._lst)
1166 def __contains__(self, element: TComparable) -> bool:
1167 """Returns ``True`` if ``element`` is present."""
1168 lst = self._lst
1169 i = bisect.bisect_left(lst, element)
1170 return i < len(lst) and lst[i] == element
1172 def __iter__(self) -> Iterator[TComparable]:
1173 """Iterates over queued elements in priority order."""
1174 return reversed(self._lst) if self._reverse else iter(self._lst)
1176 def __repr__(self) -> str:
1177 """Representation showing queue contents in current order."""
1178 return repr(list(reversed(self._lst))) if self._reverse else repr(self._lst)
1181###############################################################################
1182@final
1183class SortedInterner(Generic[TComparable]):
1184 """Same as sys.intern() except that it isn't global and that it assumes the input list is sorted (for binary search)."""
1186 def __init__(self, sorted_list: list[TComparable]) -> None:
1187 self._lst: Final[list[TComparable]] = sorted_list
1189 def interned(self, element: TComparable) -> TComparable:
1190 """Returns the interned (aka deduped) item if an equal item is contained, else returns the non-interned item."""
1191 lst = self._lst
1192 i = binary_search(lst, element)
1193 return lst[i] if i >= 0 else element
1195 def __contains__(self, element: TComparable) -> bool:
1196 """Returns ``True`` if ``element`` is present."""
1197 return binary_search(self._lst, element) >= 0
1200def binary_search(sorted_list: list[TComparable], item: TComparable) -> int:
1201 """Java-style binary search; Returns index >= 0 if an equal item is found in list, else '-insertion_point-1'; If it
1202 returns index >= 0, the index will be the left-most index in case multiple such equal items are contained."""
1203 i = bisect.bisect_left(sorted_list, item)
1204 return i if i < len(sorted_list) and sorted_list[i] == item else -i - 1
1207###############################################################################
1208_S = TypeVar("_S")
1211@final
1212class HashedInterner(Generic[_S]):
1213 """Same as sys.intern() except that it isn't global and can also be used for types other than str."""
1215 def __init__(self, items: Iterable[_S] = frozenset()) -> None:
1216 self._items: Final[dict[_S, _S]] = {v: v for v in items}
1218 def intern(self, item: _S) -> _S:
1219 """Interns the given item."""
1220 return self._items.setdefault(item, item)
1222 def interned(self, item: _S) -> _S:
1223 """Returns the interned (aka deduped) item if an equal item is contained, else returns the non-interned item."""
1224 return self._items.get(item, item)
1226 def __contains__(self, item: _S) -> bool:
1227 return item in self._items
1230#############################################################################
1231@final
1232class SynchronizedBool:
1233 """Thread-safe wrapper around a regular bool."""
1235 def __init__(self, val: bool) -> None:
1236 assert isinstance(val, bool)
1237 self._lock: Final[threading.Lock] = threading.Lock()
1238 self._value: bool = val
1240 @property
1241 def value(self) -> bool:
1242 """Returns the current boolean value."""
1243 with self._lock:
1244 return self._value
1246 @value.setter
1247 def value(self, new_value: bool) -> None:
1248 """Atomically assign ``new_value``."""
1249 with self._lock:
1250 self._value = new_value
1252 def get_and_set(self, new_value: bool) -> bool:
1253 """Swaps in ``new_value`` and return the previous value."""
1254 with self._lock:
1255 old_value = self._value
1256 self._value = new_value
1257 return old_value
1259 def compare_and_set(self, expected_value: bool, new_value: bool) -> bool:
1260 """Sets to ``new_value`` only if current value equals ``expected_value``."""
1261 with self._lock:
1262 eq: bool = self._value == expected_value
1263 if eq:
1264 self._value = new_value
1265 return eq
1267 def __bool__(self) -> bool:
1268 return self.value
1270 def __repr__(self) -> str:
1271 return repr(self.value)
1273 def __str__(self) -> str:
1274 return str(self.value)
1277#############################################################################
1278_K = TypeVar("_K")
1279_V = TypeVar("_V")
1282@final
1283class SynchronizedDict(Generic[_K, _V]):
1284 """Thread-safe wrapper around a regular dict."""
1286 def __init__(self, val: dict[_K, _V]) -> None:
1287 assert isinstance(val, dict)
1288 self._lock: Final[threading.Lock] = threading.Lock()
1289 self._dict: Final[dict[_K, _V]] = val
1291 def __getitem__(self, key: _K) -> _V:
1292 with self._lock:
1293 return self._dict[key]
1295 def __setitem__(self, key: _K, value: _V) -> None:
1296 with self._lock:
1297 self._dict[key] = value
1299 def __delitem__(self, key: _K) -> None:
1300 with self._lock:
1301 self._dict.pop(key)
1303 def __contains__(self, key: _K) -> bool:
1304 with self._lock:
1305 return key in self._dict
1307 def __len__(self) -> int:
1308 with self._lock:
1309 return len(self._dict)
1311 def __repr__(self) -> str:
1312 with self._lock:
1313 return repr(self._dict)
1315 def __str__(self) -> str:
1316 with self._lock:
1317 return str(self._dict)
1319 def get(self, key: _K, default: _V | None = None) -> _V | None:
1320 """Returns ``self[key]`` or ``default`` if missing."""
1321 with self._lock:
1322 return self._dict.get(key, default)
1324 def pop(self, key: _K, default: _V | None = None) -> _V | None:
1325 """Removes ``key`` and returns its value."""
1326 with self._lock:
1327 return self._dict.pop(key, default)
1329 def clear(self) -> None:
1330 """Removes all items atomically."""
1331 with self._lock:
1332 self._dict.clear()
1334 def items(self) -> ItemsView[_K, _V]:
1335 """Returns a snapshot of dictionary items."""
1336 with self._lock:
1337 return self._dict.copy().items()
1340#############################################################################
1341@final
1342class InterruptibleSleep:
1343 """Provides a sleep(timeout) function that can be interrupted by another thread; The underlying lock is configurable."""
1345 def __init__(self, lock: threading.Lock | None = None) -> None:
1346 self._is_stopping: bool = False
1347 self._lock: Final[threading.Lock] = lock if lock is not None else threading.Lock()
1348 self._condition: Final[threading.Condition] = threading.Condition(self._lock)
1350 def sleep(self, duration_nanos: int) -> bool:
1351 """Delays the current thread by the given number of nanoseconds; Returns True if the sleep got interrupted;
1352 Equivalent to threading.Event.wait()."""
1353 end_time_nanos: int = time.monotonic_ns() + duration_nanos
1354 with self._lock:
1355 while not self._is_stopping:
1356 diff_nanos: int = end_time_nanos - time.monotonic_ns()
1357 if diff_nanos <= 0:
1358 return False
1359 self._condition.wait(timeout=diff_nanos / 1_000_000_000) # release, then block until notified or timeout
1360 return True
1362 def interrupt(self) -> None:
1363 """Wakes sleeping threads and makes any future sleep()s a no-op; Equivalent to threading.Event.set()."""
1364 with self._lock:
1365 if not self._is_stopping:
1366 self._is_stopping = True
1367 self._condition.notify_all()
1369 def reset(self) -> None:
1370 """Makes any future sleep()s no longer a no-op; Equivalent to threading.Event.clear()."""
1371 with self._lock:
1372 self._is_stopping = False
1375#############################################################################
1376@final
1377class SynchronousExecutor(Executor):
1378 """Executor that runs tasks inline in the calling thread, sequentially."""
1380 def __init__(self) -> None:
1381 self._shutdown: bool = False
1383 def submit(self, fn: Callable[..., _R_], /, *args: Any, **kwargs: Any) -> Future[_R_]:
1384 """Executes `fn(*args, **kwargs)` immediately and returns its Future."""
1385 future: Future[_R_] = Future()
1386 if self._shutdown:
1387 raise RuntimeError("cannot schedule new futures after shutdown")
1388 try:
1389 result: _R_ = fn(*args, **kwargs)
1390 except BaseException as exc:
1391 future.set_exception(exc)
1392 else:
1393 future.set_result(result)
1394 return future
1396 def shutdown(self, wait: bool = True, *, cancel_futures: bool = False) -> None:
1397 """Prevents new submissions; no worker resources to join/cleanup."""
1398 self._shutdown = True
1400 @classmethod
1401 def executor_for(cls, max_workers: int) -> Executor:
1402 """Factory returning a SynchronousExecutor if 0 <= max_workers <= 1; else a ThreadPoolExecutor."""
1403 return cls() if 0 <= max_workers <= 1 else ThreadPoolExecutor(max_workers=max_workers)
1406#############################################################################
1407@final
1408class _XFinally(contextlib.AbstractContextManager):
1409 """Context manager ensuring cleanup code executes after ``with`` blocks."""
1411 def __init__(self, cleanup: Callable[[], None]) -> None:
1412 """Records the callable to run upon exit."""
1413 self._cleanup: Final = cleanup # Zero-argument callable executed after the `with` block exits.
1415 def __exit__(
1416 self, exc_type: type[BaseException] | None, exc: BaseException | None, tb: types.TracebackType | None
1417 ) -> Literal[False]:
1418 """Runs cleanup and propagate any exceptions appropriately."""
1419 try:
1420 self._cleanup()
1421 except BaseException as cleanup_exc:
1422 if exc is None:
1423 raise # No main error --> propagate cleanup error normally
1424 # Both failed
1425 # if sys.version_info >= (3, 11):
1426 # raise ExceptionGroup("main error and cleanup error", [exc, cleanup_exc]) from None
1427 # <= 3.10: attach so it shows up in traceback but doesn't mask
1428 exc.__context__ = cleanup_exc
1429 return False # reraise original exception
1430 return False # propagate main exception if any
1433def xfinally(cleanup: Callable[[], None]) -> _XFinally:
1434 """Usage: with xfinally(lambda: cleanup()): ...
1435 Returns a context manager that guarantees that cleanup() runs on exit and guarantees any error in cleanup() will never
1436 mask an exception raised earlier inside the body of the `with` block, while still surfacing both problems when possible.
1438 Problem it solves
1439 -----------------
1440 A naive ``try ... finally`` may lose the original exception:
1442 try:
1443 work()
1444 finally:
1445 cleanup() # <-- if this raises an exception, it replaces the real error!
1447 `_XFinally` preserves exception priority:
1449 * Body raises, cleanup succeeds --> original body exception is re-raised.
1450 * Body raises, cleanup also raises --> re-raises body exception; cleanup exception is linked via ``__context__``.
1451 * Body succeeds, cleanup raises --> cleanup exception propagates normally.
1453 Example:
1454 -------
1455 >>> with xfinally(lambda: release_resources()): # doctest: +SKIP
1456 ... run_tasks()
1458 The single *with* line replaces verbose ``try/except/finally`` boilerplate while preserving full error information.
1459 """
1460 return _XFinally(cleanup)