Coverage for bzfs_main/util/utils.py: 100%
821 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-19 18:44 +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 fcntl
32import functools
33import hashlib
34import itertools
35import logging
36import operator
37import os
38import platform
39import pwd
40import random
41import re
42import signal
43import stat
44import subprocess
45import sys
46import threading
47import time
48import types
49from collections import (
50 defaultdict,
51 deque,
52)
53from collections.abc import (
54 Generator,
55 ItemsView,
56 Iterable,
57 Iterator,
58 Mapping,
59 Sequence,
60)
61from concurrent.futures import (
62 Executor,
63 Future,
64)
65from concurrent.futures.thread import (
66 ThreadPoolExecutor,
67)
68from dataclasses import (
69 dataclass,
70)
71from datetime import (
72 datetime,
73 timedelta,
74 timezone,
75 tzinfo,
76)
77from pathlib import (
78 Path,
79)
80from subprocess import (
81 DEVNULL,
82 PIPE,
83)
84from typing import (
85 IO,
86 Any,
87 Callable,
88 Final,
89 Generic,
90 Literal,
91 NoReturn,
92 Protocol,
93 SupportsIndex,
94 TextIO,
95 TypeVar,
96 cast,
97 final,
98)
100# constants:
101PROG_NAME: Final[str] = "bzfs"
102ENV_VAR_PREFIX: Final[str] = PROG_NAME + "_"
103DIE_STATUS: Final[int] = 3
104DESCENDANTS_RE_SUFFIX: Final[str] = r"(?:/.*)?" # also match descendants of a matching dataset
105LOG_STDERR: Final[int] = (logging.INFO + logging.WARNING) // 2 # custom log level is halfway in between
106LOG_STDOUT: Final[int] = (LOG_STDERR + logging.INFO) // 2 # custom log level is halfway in between
107LOG_DEBUG: Final[int] = logging.DEBUG
108LOG_TRACE: Final[int] = logging.DEBUG // 2 # custom log level is halfway in between
109YEAR_WITH_FOUR_DIGITS_REGEX: Final[re.Pattern] = re.compile(r"[1-9][0-9][0-9][0-9]") # empty shall not match nonempty target
110UNIX_TIME_INFINITY_SECS: Final[int] = 2**64 # billions of years and to be extra safe, larger than the largest ZFS GUID
111DONT_SKIP_DATASET: Final[str] = ""
112ZFS_DATASET_NAME_REGEX: Final[re.Pattern] = re.compile(r"[a-zA-Z][a-zA-Z0-9 _.:/-]*") # see zfs_namecheck.c
113ZFS_PROPERTY_NAME_REGEX: Final[re.Pattern] = re.compile(r"[a-z0-9_.:-]+") # see zfs_prop.c
114SHELL_CHARS: Final[str] = '"' + "'`~!@#$%^&*()+={}[]|;<>?,\\" # intentionally not included: -_.:/
115SHELL_CHARS_AND_SLASH: Final[str] = SHELL_CHARS + "/"
116FILE_PERMISSIONS: Final[int] = stat.S_IRUSR | stat.S_IWUSR # rw------- (user read + write)
117DIR_PERMISSIONS: Final[int] = stat.S_IRWXU # rwx------ (user read + write + execute)
118UMASK: Final[int] = (~DIR_PERMISSIONS) & 0o777 # so intermediate dirs created by os.makedirs() have stricter permissions
119UNIX_DOMAIN_SOCKET_PATH_MAX_LENGTH: Final[int] = 107 if platform.system() == "Linux" else 103 # see Google for 'sun_path'
121RegexList = list[tuple[re.Pattern[str], bool]] # Type alias
124def getenv_any(key: str, default: str | None = None, env_var_prefix: str = ENV_VAR_PREFIX) -> str | None:
125 """All shell environment variable names used for configuration start with this prefix."""
126 return os.getenv(env_var_prefix + key, default)
129def getenv_int(key: str, default: int, env_var_prefix: str = ENV_VAR_PREFIX) -> int:
130 """Returns environment variable ``key`` as int with ``default`` fallback."""
131 return int(cast(str, getenv_any(key, default=str(default), env_var_prefix=env_var_prefix)))
134def getenv_bool(key: str, default: bool = False, env_var_prefix: str = ENV_VAR_PREFIX) -> bool:
135 """Returns environment variable ``key`` as bool with ``default`` fallback."""
136 return cast(str, getenv_any(key, default=str(default), env_var_prefix=env_var_prefix)).lower().strip() == "true"
139def cut(field: int, separator: str = "\t", *, lines: list[str]) -> list[str]:
140 """Retains only column number 'field' in a list of TSV/CSV lines; Analog to Unix 'cut' CLI command."""
141 assert lines is not None
142 assert isinstance(lines, list)
143 assert len(separator) == 1
144 if field == 1:
145 return [line[: line.index(separator)] for line in lines]
146 elif field == 2:
147 return [line[line.index(separator) + 1 :] for line in lines]
148 else:
149 raise ValueError(f"Invalid field value: {field}")
152def drain(iterable: Iterable[object]) -> None:
153 """Consumes all items in the iterable, effectively draining it."""
154 for _ in iterable:
155 del _ # help gc (iterable can block)
158_K_ = TypeVar("_K_")
159_V_ = TypeVar("_V_")
160_R_ = TypeVar("_R_")
163def shuffle_dict(dictionary: dict[_K_, _V_], /, rand: random.Random = random.SystemRandom()) -> dict[_K_, _V_]: # noqa: B008
164 """Returns a new dict with items shuffled randomly."""
165 items: list[tuple[_K_, _V_]] = list(dictionary.items())
166 rand.shuffle(items)
167 return dict(items)
170def sorted_dict(
171 dictionary: dict[_K_, _V_], /, *, key: Callable[[tuple[_K_, _V_]], Comparable] | None = None, reverse: bool = False
172) -> dict[_K_, _V_]:
173 """Returns a new dict with items sorted, primarily by key and secondarily by value (unless a custom key is supplied)."""
174 return dict(sorted(dictionary.items(), key=key, reverse=reverse))
177def sort_datasets_key(dataset: str) -> str:
178 """Ensures that the '/' ZFS dataset component separator sorts before any other character, such as '-', '.', ' '."""
179 return dataset.replace("/", "\x00") # replace '/' with the ASCII NUL character, which sorts before any other character
182def sort_datasets(datasets: Iterable[str]) -> list[str]:
183 """Sorts the given datasets such that the '/' ZFS dataset component separator sorts before any other character."""
184 # perf: sorted() is fast because Powersort is close to O(N) for nearly sorted input, which is our case
185 return sorted(datasets, key=sort_datasets_key)
188def tail(file: str, *, n: int, errors: str | None = None) -> Sequence[str]:
189 """Return the last ``n`` lines of ``file`` without following symlinks."""
190 if not os.path.isfile(file):
191 return []
192 with open_nofollow(file, "r", encoding="utf-8", errors=errors, check_owner=False, check_perm=False) as fd:
193 return deque(fd, maxlen=n)
196_NAMED_CAPTURING_GROUP: Final[re.Pattern[str]] = re.compile(r"^" + re.escape("(?P<") + r"[^\W\d]\w*" + re.escape(">"))
197_NUMERIC_BACKREFERENCE_REGEX: Final[re.Pattern[str]] = re.compile(r"\\\d+") # example: \1
200def replace_capturing_groups_with_non_capturing_groups(regex: str) -> str:
201 """Replaces regex capturing groups with non-capturing groups for better matching performance (unless it's tricky).
203 Unnamed capturing groups example: '(.*/)?tmp(foo|bar)(?!public)\\(' --> '(?:.*/)?tmp(?:foo|bar)(?!public)\\('
204 Aka replaces parenthesis '(' followed by a char other than question mark '?', but not preceded by a backslash
205 with the replacement string '(?:'
207 Named capturing group example: '(?P<name>abc)' --> '(?:abc)'
208 Aka replaces '(?P<' followed by a valid name followed by '>', but not preceded by a backslash
209 with the replacement string '(?:'
211 Also see https://docs.python.org/3/howto/regex.html#non-capturing-and-named-groups
212 """
213 i = regex.find("[")
214 if i >= 0 and regex.find("(", i) >= 0:
215 # Conservative fallback to minimize code complexity: skip the rewrite entirely in the case where the regex might
216 # contain a regex character class that contains parenthesis.
217 # Rewriting a regex is a performance optimization; correctness comes first.
218 return regex
220 if "(?P=" in regex or "(?(" in regex or _NUMERIC_BACKREFERENCE_REGEX.search(regex):
221 # Conservative fallback to minimize code complexity: skip the rewrite entirely if the regex might contain a
222 # (named or conditional or numeric) backreference.
223 # Rewriting a regex is a performance optimization; correctness comes first.
224 return regex
226 i = len(regex) - 2
227 while i >= 0:
228 i = regex.rfind("(", 0, i + 1)
229 if i >= 0 and (i == 0 or regex[i - 1] != "\\"):
230 if regex[i + 1] != "?":
231 regex = f"{regex[0:i]}(?:{regex[i + 1:]}" # unnamed capturing group
232 else: # potentially a valid named capturing group
233 regex = regex[0:i] + _NAMED_CAPTURING_GROUP.sub(repl="(?:", string=regex[i:], count=1)
234 i -= 1
235 return regex
238def get_home_directory() -> str:
239 """Reliably detects home dir without using HOME env var."""
240 # thread-safe version of: os.environ.pop('HOME', None); os.path.expanduser('~')
241 return pwd.getpwuid(os.getuid()).pw_dir
244def human_readable_bytes(num_bytes: float, *, separator: str = " ", precision: int | None = None) -> str:
245 """Formats 'num_bytes' as a human-readable size; for example "567 MiB"."""
246 sign = "-" if num_bytes < 0 else ""
247 s = abs(num_bytes)
248 units = ("B", "KiB", "MiB", "GiB", "TiB", "PiB", "EiB", "ZiB", "YiB", "RiB", "QiB")
249 n = len(units) - 1
250 i = 0
251 while s >= 1024 and i < n:
252 s /= 1024
253 i += 1
254 formatted_num = human_readable_float(s) if precision is None else f"{s:.{precision}f}"
255 return f"{sign}{formatted_num}{separator}{units[i]}"
258def human_readable_duration(duration: float, *, unit: str = "ns", separator: str = "", precision: int | None = None) -> str:
259 """Formats a duration in human units, automatically scaling as needed; for example "567ms"."""
260 sign = "-" if duration < 0 else ""
261 t = abs(duration)
262 units = ("ns", "μs", "ms", "s", "m", "h", "d")
263 i = units.index(unit)
264 if t < 1 and t != 0:
265 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)
266 t *= nanos[i]
267 i = 0
268 while t >= 1000 and i < 3:
269 t /= 1000
270 i += 1
271 if i >= 3:
272 while t >= 60 and i < 5:
273 t /= 60
274 i += 1
275 if i >= 5:
276 while t >= 24 and i < len(units) - 1:
277 t /= 24
278 i += 1
279 formatted_num = human_readable_float(t) if precision is None else f"{t:.{precision}f}"
280 return f"{sign}{formatted_num}{separator}{units[i]}"
283def human_readable_float(number: float) -> str:
284 """Formats ``number`` with a variable precision depending on magnitude.
286 This design mirrors the way humans round values when scanning logs.
288 If the number has one digit before the decimal point (0 <= abs(number) < 10):
289 Round and use two decimals after the decimal point (e.g., 3.14559 --> "3.15").
291 If the number has two digits before the decimal point (10 <= abs(number) < 100):
292 Round and use one decimal after the decimal point (e.g., 12.36 --> "12.4").
294 If the number has three or more digits before the decimal point (abs(number) >= 100):
295 Round and use zero decimals after the decimal point (e.g., 123.556 --> "124").
297 Ensures no unnecessary trailing zeroes are retained: Example: 1.500 --> "1.5", 1.00 --> "1"
298 """
299 abs_number = abs(number)
300 precision = 2 if abs_number < 10 else 1 if abs_number < 100 else 0
301 if precision == 0:
302 return str(round(number))
303 result = f"{number:.{precision}f}"
304 assert "." in result
305 result = result.rstrip("0").rstrip(".") # Remove trailing zeros and trailing decimal point if empty
306 if result == "-0":
307 result = "0"
308 return result
311def percent(number: int, total: int, *, print_total: bool = False) -> str:
312 """Returns percentage string of ``number`` relative to ``total``."""
313 tot: str = f"/{total}" if print_total else ""
314 pct: str = "inf" if total == 0 else human_readable_float(100 * number / total)
315 return f"{number}{tot}={pct}%"
318def open_nofollow(
319 path: str,
320 mode: str = "r",
321 *,
322 buffering: int = -1,
323 encoding: str | None = None,
324 errors: str | None = None,
325 newline: str | None = None,
326 perm: int = FILE_PERMISSIONS,
327 check_owner: bool = True,
328 check_perm: bool = True,
329 **kwargs: Any,
330) -> IO[Any]:
331 """Behaves exactly like built-in open(), except that it refuses to follow symlinks, i.e. raises OSError with
332 errno.ELOOP/EMLINK if basename of path is a symlink.
334 Also, can specify custom permissions on O_CREAT, and verify secure ownership and mode bits.
336 If check_owner=True, write-capable opens require ownership by the effective UID; read-only opens also allow ownership by
337 uid 0 (root). This allows safe reads of root-owned system files while preventing writes to files not owned by the caller.
339 If check_perm=True, the file must also not be writable by group or others, and must not be executable by anyone.
340 """
341 if not mode:
342 raise ValueError("Must have exactly one of create/read/write/append mode and at most one plus")
343 flags = {
344 "r": os.O_RDONLY,
345 "w": os.O_WRONLY | os.O_CREAT | os.O_TRUNC,
346 "a": os.O_WRONLY | os.O_CREAT | os.O_APPEND,
347 "x": os.O_WRONLY | os.O_CREAT | os.O_EXCL,
348 }.get(mode[0])
349 if flags is None:
350 raise ValueError(f"invalid mode {mode!r}")
351 if "+" in mode: # enable read-write access for r+, w+, a+, x+
352 flags = (flags & ~os.O_WRONLY) | os.O_RDWR # clear os.O_WRONLY and set os.O_RDWR while preserving all other flags
353 flags |= os.O_NOFOLLOW | os.O_CLOEXEC
354 fd: int = os.open(path, flags=flags, mode=perm)
355 try:
356 stats: os.stat_result | None = None
357 if check_owner:
358 stats = os.fstat(fd)
359 st_uid: int = stats.st_uid
360 if st_uid != os.geteuid(): # verify ownership is current effective UID
361 if (flags & (os.O_WRONLY | os.O_RDWR)) != 0: # require that writer owns the file
362 raise PermissionError(errno.EPERM, f"{path!r} is owned by uid {st_uid}, not {os.geteuid()}", path)
363 elif st_uid != 0: # it's ok for root to own a file that we'll merely read
364 raise PermissionError(errno.EPERM, f"{path!r} is owned by uid {st_uid}, not {os.geteuid()} or 0", path)
365 if check_perm:
366 stats = os.fstat(fd) if stats is None else stats
367 if (stats.st_mode & (stat.S_IWGRP | stat.S_IWOTH | stat.S_IXUSR | stat.S_IXGRP | stat.S_IXOTH)) != 0:
368 st_mode: int = stat.S_IMODE(stats.st_mode)
369 raise PermissionError(
370 errno.EPERM,
371 f"{path!r} has permissions {st_mode:03o} aka {stat.filemode(st_mode)[1:]} "
372 "but must not be writable by group or others, or executable by anyone",
373 path,
374 )
375 return os.fdopen(fd, mode, buffering=buffering, encoding=encoding, errors=errors, newline=newline, **kwargs)
376 except Exception:
377 try:
378 os.close(fd)
379 except OSError:
380 pass
381 raise
384def close_quietly(fd: int) -> None:
385 """Closes the given file descriptor while silently swallowing any OSError that might arise as part of this."""
386 if fd >= 0:
387 try:
388 os.close(fd)
389 except OSError:
390 pass
393@contextlib.contextmanager
394def nonblocking_file_lock(path: str) -> Generator[bool]:
395 """Context manager that yields whether an exclusive non-blocking flock was acquired, deleting its file before releasing
396 a successful acquisition in order to avoid accumulating stale lock files over time."""
398 def _is_current_file(fd: int, path: str) -> bool:
399 """Checks identity while holding the flock; a missing or replaced pathname cannot result in lock acquisition."""
400 try:
401 return os.path.samestat(os.fstat(fd), os.stat(path, follow_symlinks=False))
402 except FileNotFoundError:
403 return False
405 while True:
406 fd: int = os.open(path, os.O_WRONLY | os.O_CREAT | os.O_NOFOLLOW | os.O_CLOEXEC, FILE_PERMISSIONS)
407 with xfinally(functools.partial(os.close, fd)):
408 try:
409 # Acquire an exclusive lock; will raise a BlockingIOError if lock is already held by this process or
410 # another process. The (advisory) lock is auto-released when the process terminates or the fd is closed.
411 fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB) # LOCK_NB ... non-blocking
412 except BlockingIOError:
413 yield False
414 return
416 if not _is_current_file(fd, path):
417 continue # retry on race
419 # xfinally: unlink the lock_file while still holding the flock on its fd - it's a correct and safe
420 # standard POSIX pattern:
421 # - Performing unlink() before close(fd) avoids a race where a subsequent bzfs process could recreate and
422 # lock a fresh inode for the same path between our close() and a later unlink(). In that case, a late
423 # unlink would delete the newer process's lock_file path.
424 with xfinally(lambda: Path(path).unlink(missing_ok=True)): # don't accumulate stale files
425 yield True
426 return
429_P = TypeVar("_P")
432def find_match(
433 seq: Sequence[_P],
434 predicate: Callable[[_P], bool],
435 start: SupportsIndex | None = None,
436 end: SupportsIndex | None = None,
437 *,
438 reverse: bool = False,
439 raises: bool | object | Callable[[], object] = False, # raises: bool | object | Callable = False, # python >= 3.10
440) -> int:
441 """Returns the integer index within ``seq`` of the first item (or last item if reverse=True) that matches the given
442 predicate condition.
444 If no matching item is found returns -1 or ValueError, depending on the ``raises`` parameter, which is a bool indicating
445 whether to raise an error, or an object containing the error message, but can also be a Callable/lambda in order to
446 support efficient deferred generation of error messages.
448 Analog to ``str.find()``, including slicing semantics with parameters start and end, i.e. respects Python slicing
449 semantics for start/end (including clamping). For example, seq can be a list, tuple or str.
451 Example usage:
452 lst = ["a", "b", "-c", "d"]
453 i = find_match(lst, lambda arg: arg.startswith("-"), start=1, end=3, reverse=True)
454 if i >= 0:
455 print(lst[i])
456 i = find_match(lst, lambda arg: arg.startswith("-"), raises=f"Tag {tag} not found in {file}")
457 i = find_match(lst, lambda arg: arg.startswith("-"), raises=lambda: f"Tag {tag} not found in {file}")
458 """
459 if start is None and end is None:
460 for i in range(len(seq) - 1, -1, -1) if reverse else range(len(seq)):
461 if predicate(seq[i]):
462 return i
463 else:
464 slice_start, slice_end, _ = slice(start, end).indices(len(seq))
465 for i in range(slice_end - 1, slice_start - 1, -1) if reverse else range(slice_start, slice_end):
466 if predicate(seq[i]):
467 return i
468 if raises is False or raises is None:
469 return -1
470 if raises is True:
471 raise ValueError("No matching item found in sequence")
472 if callable(raises):
473 raises = raises()
474 raise ValueError(raises)
477def is_descendant(dataset: str, of_root_dataset: str) -> bool:
478 """Returns True if ZFS ``dataset`` lies under ``of_root_dataset`` in the dataset hierarchy, or is the same."""
479 return dataset == of_root_dataset or dataset.startswith(of_root_dataset + "/")
482def has_prefix(sorted_strings: list[str], prefix: str, *, is_test_mode: bool = False) -> bool:
483 """Returns whether any string in lexicographically sorted_strings starts with the given prefix."""
484 assert (not is_test_mode) or sorted_strings == sorted(sorted_strings), "List is not sorted"
485 i = bisect.bisect_left(sorted_strings, prefix)
486 return i < len(sorted_strings) and sorted_strings[i].startswith(prefix)
489def has_duplicates(sorted_list: Sequence[object]) -> bool:
490 """Returns True if any adjacent items within the given sorted sequence are equal."""
491 return any(map(operator.eq, sorted_list, itertools.islice(sorted_list, 1, None)))
494def dry(msg: str, is_dry_run: bool) -> str:
495 """Prefix ``msg`` with 'Dry' when in dry-run mode."""
496 return "Dry " + msg if is_dry_run else msg
499def relativize_dataset(dataset: str, root_dataset: str) -> str:
500 """Converts an absolute dataset path to one relative to ``root_dataset``.
502 Example: root_dataset=tank/foo, dataset=tank/foo/bar/baz --> relative_path=/bar/baz.
503 """
504 return dataset[len(root_dataset) :]
507def dataset_paths(dataset: str) -> Iterator[str]:
508 """Enumerates all paths of a valid ZFS dataset name; Example: "a/b/c" --> yields "a", "a/b", "a/b/c"."""
509 i: int = 0
510 while i >= 0:
511 i = dataset.find("/", i)
512 if i < 0:
513 yield dataset
514 else:
515 yield dataset[:i]
516 i += 1
519def replace_prefix(s: str, old_prefix: str, new_prefix: str) -> str:
520 """In a string s, replaces a leading old_prefix string with new_prefix; assumes the leading string is present."""
521 assert s.startswith(old_prefix)
522 return new_prefix + s[len(old_prefix) :]
525def replace_in_lines(lines: list[str], old: str, new: str, count: int = -1) -> None:
526 """Replaces ``old`` with ``new`` in-place for every string in ``lines``."""
527 for i in range(len(lines)):
528 lines[i] = lines[i].replace(old, new, count)
531_TAPPEND = TypeVar("_TAPPEND")
534def append_if_absent(lst: list[_TAPPEND], *items: _TAPPEND) -> list[_TAPPEND]:
535 """Appends items to list if they are not already present."""
536 for item in items:
537 if item not in lst:
538 lst.append(item)
539 return lst
542def xappend(lst: list[_TAPPEND], *items: _TAPPEND | Iterable[_TAPPEND]) -> list[_TAPPEND]:
543 """Appends each of the items to the given list if the item is "truthy", for example not None and not an empty string; If
544 an item is an iterable does so recursively, flattening the output."""
545 for item in items:
546 if isinstance(item, str) or not isinstance(item, Iterable):
547 if item:
548 lst.append(item)
549 else:
550 xappend(lst, *item)
551 return lst
554def is_included(name: str, include_regexes: RegexList, exclude_regexes: RegexList) -> bool:
555 """Returns True if the name matches at least one of the include regexes but none of the exclude regexes; else False.
557 A regex that starts with a `!` is a negation - the regex matches if the regex without the `!` prefix does not match.
558 """
559 for regex, is_negation in exclude_regexes:
560 is_match = regex.fullmatch(name) is not None
561 if is_negation:
562 is_match = not is_match
563 if is_match:
564 return False
566 for regex, is_negation in include_regexes:
567 is_match = regex.fullmatch(name) is not None
568 if is_negation:
569 is_match = not is_match
570 if is_match:
571 return True
573 return False
576def compile_regexes(regexes: list[str], *, suffix: str = "") -> RegexList:
577 """Compiles regex strings and keeps track of negations."""
578 assert isinstance(regexes, list)
579 compiled_regexes: RegexList = []
580 for regex in regexes:
581 if suffix: # disallow non-trailing end-of-str symbol in dataset regexes to ensure descendants will also match
582 original = regex
583 if regex.endswith("\\$"):
584 pass # trailing literal $ is ok
585 elif regex.endswith("$"):
586 regex = regex[0:-1] # ok because all users of compile_regexes() call re.fullmatch()
587 if "$" in regex.removesuffix("\\$"):
588 raise re.error("Must not use non-trailing '$' character", original)
589 if "\\z" in regex or "\\Z" in regex:
590 raise re.error("Must not use '\\z' or '\\Z' character", original)
591 if is_negation := regex.startswith("!"):
592 regex = regex[1:]
593 regex = replace_capturing_groups_with_non_capturing_groups(regex)
594 if regex != ".*" or not (suffix.startswith("(") and suffix.endswith(")?")):
595 if suffix:
596 # wrap in (non-capturing) group to retain semantics of regex alternations
597 # while preserving global inline flags as-is, e.g. for case-insensitive matching
598 prefix: re.Match[str] | None = re.match(r"(\(\?[aimsux]+\))+", regex) # e.g. (?i)
599 end: int = prefix.end() if prefix else 0
600 regex = f"{regex[:end]}(?:{regex[end:]}){suffix}"
601 compiled_regexes.append((re.compile(regex), is_negation))
602 return compiled_regexes
605def list_formatter(iterable: Iterable[object], separator: str = " ", lstrip: bool = False) -> object:
606 """Lazy formatter joining items with ``separator`` used to avoid overhead in disabled log levels."""
608 @final
609 class CustomListFormatter:
610 """Formatter object that joins items when converted to ``str``."""
612 def __str__(self) -> str:
613 s = separator.join(map(str, iterable))
614 return s.lstrip() if lstrip else s
616 return CustomListFormatter()
619def pretty_print_formatter(obj_to_format: object) -> object:
620 """Lazy pprint formatter used to avoid overhead in disabled log levels."""
622 @final
623 class PrettyPrintFormatter:
624 """Formatter that pretty-prints the object on conversion to ``str``."""
626 def __str__(self) -> str:
627 import pprint # lazy import for startup perf
629 return pprint.pformat(vars(obj_to_format))
631 return PrettyPrintFormatter()
634def stderr_to_str(stderr: object) -> str:
635 """Workaround for https://github.com/python/cpython/issues/87597."""
636 return str(stderr) if not isinstance(stderr, bytes) else stderr.decode("utf-8", errors="replace")
639def xprint(log: logging.Logger, value: object, *, run: bool = True, end: str = "\n", file: TextIO | None = None) -> None:
640 """Optionally logs ``value`` at stdout/stderr level."""
641 if run and value:
642 value = value if end else str(value).rstrip()
643 level = LOG_STDOUT if file is sys.stdout else LOG_STDERR
644 log.log(level, "%s", value)
647def sha256_hex(text: str) -> str:
648 """Returns the sha256 hex string for the given text."""
649 return hashlib.sha256(text.encode()).hexdigest()
652def sha256_urlsafe_base64(text: str, *, padding: bool = True) -> str:
653 """Returns the URL-safe base64-encoded sha256 value for the given text."""
654 digest: bytes = hashlib.sha256(text.encode()).digest()
655 s: str = base64.urlsafe_b64encode(digest).decode()
656 return s if padding else s.rstrip("=")
659def sha256_128_urlsafe_base64(text: str) -> str:
660 """Returns the left half portion of the unpadded URL-safe base64-encoded sha256 value for the given text."""
661 s: str = sha256_urlsafe_base64(text, padding=False)
662 return s[: len(s) // 2]
665def sha256_85_urlsafe_base64(text: str) -> str:
666 """Returns the left one third portion of the unpadded URL-safe base64-encoded sha256 value for the given text."""
667 s: str = sha256_urlsafe_base64(text, padding=False)
668 return s[: len(s) // 3]
671def urlsafe_base64(
672 value: int, max_value: int = 2**64 - 1, *, padding: bool = True, byteorder: Literal["little", "big"] = "big"
673) -> str:
674 """Returns the URL-safe base64 string encoding of the int value, assuming it is contained in the range [0..max_value]."""
675 assert 0 <= value <= max_value
676 max_bytes: int = (max_value.bit_length() + 7) // 8
677 value_bytes: bytes = value.to_bytes(max_bytes, byteorder)
678 s: str = base64.urlsafe_b64encode(value_bytes).decode()
679 return s if padding else s.rstrip("=")
682def die(msg: str, exit_code: int = DIE_STATUS, parser: argparse.ArgumentParser | None = None) -> NoReturn:
683 """Exits the program with ``exit_code`` after logging ``msg``."""
684 if parser is None:
685 ex = SystemExit(msg)
686 ex.code = exit_code
687 raise ex
688 else:
689 parser.error(msg)
692def subprocess_run(*args: Any, **kwargs: Any) -> subprocess.CompletedProcess:
693 """Drop-in replacement for subprocess.run() that mimics its behavior except it enhances cleanup on TimeoutExpired, and
694 provides optional child PID tracking, and optional logging of execution status via ``log`` and ``loglevel`` params."""
695 input_value = kwargs.pop("input", None)
696 timeout = kwargs.pop("timeout", None)
697 check = kwargs.pop("check", False)
698 subprocesses: Subprocesses | None = kwargs.pop("subprocesses", None)
699 if input_value is not None:
700 if kwargs.get("stdin") is not None:
701 raise ValueError("input and stdin are mutually exclusive")
702 kwargs["stdin"] = subprocess.PIPE
704 log: logging.Logger | None = kwargs.pop("log", None)
705 loglevel: int | None = kwargs.pop("loglevel", None)
706 start_time_nanos: int = time.monotonic_ns()
707 is_timeout: bool = False
708 is_cancel: bool = False
709 exitcode: int | None = None
711 def log_status() -> None:
712 if log is not None:
713 _loglevel: int = loglevel if loglevel is not None else getenv_int("subprocess_run_loglevel", LOG_TRACE)
714 if log.isEnabledFor(_loglevel):
715 elapsed_time: str = human_readable_float((time.monotonic_ns() - start_time_nanos) / 1_000_000) + "ms"
716 status = "cancel" if is_cancel else "timeout" if is_timeout else "success" if exitcode == 0 else "failure"
717 cmd = kwargs["args"] if "args" in kwargs else (args[0] if args else None)
718 cmd_str: str = " ".join(str(arg) for arg in iter(cmd)) if isinstance(cmd, (list, tuple)) else str(cmd)
719 log.log(_loglevel, f"Executed [{status}] [{elapsed_time}]: %s", cmd_str)
721 with xfinally(log_status):
722 ctx: contextlib.AbstractContextManager[subprocess.Popen]
723 if subprocesses is None:
724 ctx = subprocess.Popen(*args, **kwargs)
725 else:
726 ctx = subprocesses.popen_and_track(*args, **kwargs)
727 with ctx as proc:
728 try:
729 sp = subprocesses
730 if sp is not None and sp._is_terminated(): # noqa: SLF001 pylint: disable=protected-access
731 is_cancel = True
732 timeout = 0.0
733 stdout, stderr = proc.communicate(input_value, timeout=timeout)
734 except BaseException as e:
735 try:
736 if isinstance(e, subprocess.TimeoutExpired):
737 is_timeout = True
738 terminate_process_subtree(root_pids=[proc.pid]) # send SIGTERM to child proc and descendants
739 finally:
740 proc.kill()
741 raise
742 else:
743 exitcode = proc.poll()
744 assert exitcode is not None
745 if check and exitcode:
746 raise subprocess.CalledProcessError(exitcode, proc.args, output=stdout, stderr=stderr)
747 return subprocess.CompletedProcess(proc.args, exitcode, stdout, stderr)
750def terminate_process_subtree(
751 *, except_current_process: bool = True, root_pids: list[int] | None = None, sig: signal.Signals = signal.SIGTERM
752) -> None:
753 """For each root PID: Sends the given signal to the root PID and all its descendant processes."""
754 current_pid: int = os.getpid()
755 root_pids = [current_pid] if root_pids is None else root_pids
756 all_pids: list[list[int]] = _get_descendant_processes(root_pids)
757 assert len(all_pids) == len(root_pids)
758 for i, pids in enumerate(all_pids):
759 root_pid = root_pids[i]
760 if root_pid == current_pid:
761 pids += [] if except_current_process else [current_pid]
762 else:
763 pids.insert(0, root_pid)
764 for pid in pids:
765 with contextlib.suppress(OSError):
766 os.kill(pid, sig)
769def _get_descendant_processes(root_pids: list[int]) -> list[list[int]]:
770 """For each root PID, returns the list of all descendant process IDs for the given root PID, on POSIX systems."""
771 if len(root_pids) == 0:
772 return []
773 cmd: list[str] = ["ps", "-Ao", "pid,ppid"]
774 try:
775 lines: list[str] = subprocess.run(cmd, stdin=DEVNULL, stdout=PIPE, text=True, check=True).stdout.splitlines()
776 except PermissionError:
777 # degrade gracefully in sandbox environments that deny executing `ps` entirely
778 return [[] for _ in root_pids]
779 procs: dict[int, list[int]] = defaultdict(list)
780 for line in lines[1:]: # all lines except the header line
781 splits: list[str] = line.split()
782 assert len(splits) == 2
783 pid = int(splits[0])
784 ppid = int(splits[1])
785 procs[ppid].append(pid)
787 def recursive_append(ppid: int, descendants: list[int]) -> None:
788 """Recursively collect descendant PIDs starting from ``ppid``."""
789 for child_pid in procs[ppid]:
790 descendants.append(child_pid)
791 recursive_append(child_pid, descendants)
793 all_descendants: list[list[int]] = []
794 for root_pid in root_pids:
795 descendants: list[int] = []
796 recursive_append(root_pid, descendants)
797 all_descendants.append(descendants)
798 return all_descendants
801@contextlib.contextmanager
802def termination_signal_handler(
803 termination_events: list[threading.Event],
804 *,
805 termination_handler: Callable[[], None] = lambda: terminate_process_subtree(),
806) -> Generator[None]:
807 """Context manager that installs SIGINT/SIGTERM handlers that set all ``termination_events`` and, by default, terminate
808 all descendant processes."""
809 termination_events = list(termination_events) # shallow copy
811 def _handler(_sig: int, _frame: object) -> None:
812 for event in termination_events:
813 event.set()
814 termination_handler()
816 previous_int_handler = signal.signal(signal.SIGINT, _handler) # install new signal handler
817 previous_term_handler = signal.signal(signal.SIGTERM, _handler) # install new signal handler
818 try:
819 yield # run body of context manager
820 finally:
821 signal.signal(signal.SIGINT, previous_int_handler) # restore original signal handler
822 signal.signal(signal.SIGTERM, previous_term_handler) # restore original signal handler
825def return_false() -> bool:
826 """Always returns ``False``; picklable."""
827 return False
830def sleep_nanos(delay_nanos: int) -> None:
831 """Same as time.sleep() but expects a relative sleep duration in nanoseconds as input value; picklable."""
832 time.sleep(delay_nanos / 1_000_000_000)
835#############################################################################
836@dataclass(frozen=True)
837@final
838class TaskTiming:
839 """Customizable callbacks for reading the current monotonic time, sleeping and optional async termination; immutable."""
841 monotonic_ns: Callable[[], int] = time.monotonic_ns
843 is_terminated: Callable[[], bool] = return_false
844 """Returns whether a predicate has become true; can be used to indicate system shutdown or similar cancellation
845 conditions; default is to always return ``False``."""
847 sleep: Callable[[int], None] = sleep_nanos
848 """Sleeps N nanoseconds; thread-safe."""
850 def copy(self, **override_kwargs: Any) -> TaskTiming:
851 """Creates a new object copying an existing one with the specified fields overridden for customization; thread-
852 safe."""
853 return dataclasses.replace(self, **override_kwargs)
855 @staticmethod
856 def make_from(termination_event: threading.Event | None) -> TaskTiming:
857 """Convenience factory that creates an object that performs async termination when ``termination_event`` is set."""
858 if termination_event is None:
859 return TaskTiming()
861 def _sleep(delay_nanos: int) -> None:
862 termination_event.wait(delay_nanos / 1_000_000_000) # allow early wakeup on async termination
864 return TaskTiming(is_terminated=termination_event.is_set, sleep=_sleep)
867#############################################################################
868@final
869class Subprocesses:
870 """Provides per-job tracking of child PIDs so a job can safely terminate only the subprocesses it spawned itself; used
871 when multiple jobs run concurrently within the same Python process.
873 Optionally binds to an ``_is_terminated`` predicate to enforce async cancellation by forcing immediate timeouts for newly
874 spawned subprocesses once cancellation is requested.
875 """
877 def __init__(self, is_terminated: Callable[[], bool] = return_false) -> None:
878 self._is_terminated: Final[Callable[[], bool]] = is_terminated
879 self._lock: Final[threading.Lock] = threading.Lock()
880 self._child_pids: Final[dict[int, None]] = {} # a set that preserves insertion order
882 @contextlib.contextmanager
883 def popen_and_track(self, *popen_args: Any, **popen_kwargs: Any) -> Generator[subprocess.Popen]:
884 """Context manager that calls subprocess.Popen() and tracks the child PID for per-job termination.
886 Holds a lock across Popen+PID registration to prevent a race when terminate_process_subtrees() is invoked (e.g. from
887 SIGINT/SIGTERM handlers), ensuring newly spawned child processes cannot escape termination. The child PID is
888 unregistered on context exit.
889 """
890 with self._lock:
891 proc: subprocess.Popen = subprocess.Popen(*popen_args, **popen_kwargs)
892 self._child_pids[proc.pid] = None
893 try:
894 yield proc
895 finally:
896 with self._lock:
897 self._child_pids.pop(proc.pid, None)
899 def subprocess_run(self, *args: Any, **kwargs: Any) -> subprocess.CompletedProcess:
900 """Wrapper around utils.subprocess_run() that auto-registers/unregisters child PIDs for per-job termination."""
901 return subprocess_run(*args, **kwargs, subprocesses=self)
903 def terminate_process_subtrees(self, sig: signal.Signals = signal.SIGTERM) -> None:
904 """Sends the given signal to all tracked child PIDs and their descendants, ignoring errors for dead PIDs."""
905 with self._lock:
906 pids: list[int] = list(self._child_pids)
907 self._child_pids.clear()
908 terminate_process_subtree(root_pids=pids, sig=sig)
911#############################################################################
912def pid_exists(pid: int) -> bool | None:
913 """Returns True if a process with PID exists, False if not, or None on error."""
914 if pid <= 0:
915 return False
916 try: # with signal=0, no signal is actually sent, but error checking is still performed
917 os.kill(pid, 0) # ... which can be used to check for process existence on POSIX systems
918 except OSError as err:
919 if err.errno == errno.ESRCH: # No such process
920 return False
921 if err.errno == errno.EPERM: # Operation not permitted
922 return True
923 return None
924 return True
927def nprefix(s: str) -> str:
928 """Returns a canonical snapshot prefix with trailing underscore."""
929 return sys.intern(s + "_")
932def ninfix(s: str) -> str:
933 """Returns a canonical infix with trailing underscore when not empty."""
934 return sys.intern(s + "_") if s else ""
937def nsuffix(s: str) -> str:
938 """Returns a canonical suffix with leading underscore when not empty."""
939 return sys.intern("_" + s) if s else ""
942def format_dict(dictionary: Mapping[_K_, _V_]) -> str:
943 """Returns a formatted dictionary using repr for consistent output."""
944 return f'"{dictionary}"'
947def format_obj(obj: object) -> str:
948 """Returns a formatted str using repr for consistent output."""
949 return f'"{obj}"'
952def validate_dataset_name(dataset: str, input_text: str) -> None:
953 """'zfs create' CLI does not accept dataset names that are empty or start or end in a slash, etc."""
954 # Also see zfs_namecheck.c and https://github.com/openzfs/zfs/issues/439#issuecomment-2784424
955 # and https://github.com/openzfs/zfs/issues/8798
956 # and (by now no longer accurate): https://docs.oracle.com/cd/E26505_01/html/E37384/gbcpt.html
957 if (
958 dataset in ("", ".", "..")
959 or dataset.startswith(("/", "./", "../"))
960 or dataset.endswith(("/", "/.", "/.."))
961 or any(substring in dataset for substring in ("//", "/./", "/../"))
962 or not ZFS_DATASET_NAME_REGEX.fullmatch(dataset)
963 ):
964 die(f"Invalid ZFS dataset name: '{dataset}' for: '{input_text}'")
967def validate_property_name(propname: str, input_text: str) -> str:
968 """Checks that the ZFS property name contains no spaces or shell chars, etc."""
969 if propname.startswith("-") or not ZFS_PROPERTY_NAME_REGEX.fullmatch(propname):
970 die(f"Invalid ZFS property name: '{propname}' for: '{input_text}'")
971 return propname
974def validate_is_not_a_symlink(msg: str, path: str, parser: argparse.ArgumentParser | None = None) -> None:
975 """Checks that the given path is not a symbolic link."""
976 if os.path.islink(path):
977 die(f"{msg}must not be a symlink: {path}", parser=parser)
980def validate_file_permissions(path: str, mode: int) -> None:
981 """Verify permissions and that ownership is current effective UID."""
982 stats: os.stat_result = os.stat(path, follow_symlinks=False)
983 st_uid: int = stats.st_uid
984 if st_uid != os.geteuid(): # verify ownership is current effective UID
985 die(f"{path!r} is owned by uid {st_uid}, not {os.geteuid()}")
986 st_mode = stat.S_IMODE(stats.st_mode)
987 if st_mode != mode:
988 die(
989 f"{path!r} has permissions {st_mode:03o} aka {stat.filemode(st_mode)[1:]}, "
990 f"not {mode:03o} aka {stat.filemode(mode)[1:]})"
991 )
994def parse_duration_to_milliseconds(duration: str, *, regex_suffix: str = "", context: str = "") -> int:
995 """Parses human duration strings like '5minutes' or '2 hours' to milliseconds."""
996 unit_milliseconds: dict[str, int] = {
997 "milliseconds": 1,
998 "millis": 1,
999 "seconds": 1000,
1000 "secs": 1000,
1001 "minutes": 60 * 1000,
1002 "mins": 60 * 1000,
1003 "hours": 60 * 60 * 1000,
1004 "days": 86400 * 1000,
1005 "weeks": 7 * 86400 * 1000,
1006 "months": round(30.5 * 86400 * 1000),
1007 "years": 365 * 86400 * 1000,
1008 }
1009 match = re.fullmatch(
1010 r"(\d+)\s*(milliseconds|millis|seconds|secs|minutes|mins|hours|days|weeks|months|years)" + regex_suffix,
1011 duration,
1012 )
1013 if not match:
1014 if context:
1015 die(f"Invalid duration format: {duration} within {context}")
1016 else:
1017 raise ValueError(f"Invalid duration format: {duration}")
1018 assert match
1019 quantity: int = int(match.group(1))
1020 unit: str = match.group(2)
1021 return quantity * unit_milliseconds[unit]
1024def unixtime_fromisoformat(datetime_str: str) -> int:
1025 """Converts ISO 8601 datetime string into UTC Unix time in integer seconds."""
1026 return int(datetime.fromisoformat(datetime_str).timestamp())
1029def isotime_from_unixtime(unixtime_in_seconds: int) -> str:
1030 """Converts UTC Unix time seconds into ISO 8601 datetime string."""
1031 tz: tzinfo = timezone.utc
1032 dt: datetime = datetime.fromtimestamp(unixtime_in_seconds, tz=tz)
1033 return dt.isoformat(sep="_", timespec="seconds")
1036def current_datetime(
1037 tz_spec: str | None = None,
1038 now_fn: Callable[[tzinfo | None], datetime] | None = None,
1039) -> datetime:
1040 """Returns current time in ``tz_spec`` timezone or local timezone."""
1041 if now_fn is None:
1042 now_fn = datetime.now
1043 return now_fn(get_timezone(tz_spec))
1046def get_timezone(tz_spec: str | None = None) -> tzinfo | None:
1047 """Returns timezone from spec or local timezone if unspecified."""
1048 tz: tzinfo | None
1049 if tz_spec is None:
1050 tz = None
1051 elif tz_spec == "UTC":
1052 tz = timezone.utc
1053 elif match := re.fullmatch(r"([+-])(\d\d):?(\d\d)", tz_spec):
1054 sign, hours, minutes = match.groups()
1055 offset: int = int(hours) * 60 + int(minutes)
1056 offset = -offset if sign == "-" else offset
1057 tz = timezone(timedelta(minutes=offset))
1058 elif "/" in tz_spec:
1059 from zoneinfo import ZoneInfo # lazy import for startup perf
1061 tz = ZoneInfo(tz_spec)
1062 else:
1063 raise ValueError(f"Invalid timezone specification: {tz_spec}")
1064 return tz
1067###############################################################################
1068@final
1069class SnapshotPeriods: # thread-safe
1070 """Parses snapshot suffix strings and converts between durations."""
1072 def __init__(self) -> None:
1073 """Initialize lookup tables of suffixes and corresponding millis."""
1074 self.suffix_milliseconds: Final[dict[str, int]] = {
1075 "yearly": 365 * 86400 * 1000,
1076 "monthly": round(30.5 * 86400 * 1000),
1077 "weekly": 7 * 86400 * 1000,
1078 "daily": 86400 * 1000,
1079 "hourly": 60 * 60 * 1000,
1080 "minutely": 60 * 1000,
1081 "secondly": 1000,
1082 "millisecondly": 1,
1083 }
1084 self.period_labels: Final[dict[str, str]] = {
1085 "yearly": "years",
1086 "monthly": "months",
1087 "weekly": "weeks",
1088 "daily": "days",
1089 "hourly": "hours",
1090 "minutely": "minutes",
1091 "secondly": "seconds",
1092 "millisecondly": "milliseconds",
1093 }
1094 self._suffix_regex0: Final[re.Pattern] = re.compile(rf"([1-9][0-9]*)?({'|'.join(self.suffix_milliseconds.keys())})")
1095 self._suffix_regex1: Final[re.Pattern] = re.compile("_" + self._suffix_regex0.pattern)
1097 def suffix_to_duration0(self, suffix: str) -> tuple[int, str]:
1098 """Parse suffix like '10minutely' to (10, 'minutely')."""
1099 return self._suffix_to_duration(suffix, self._suffix_regex0)
1101 def suffix_to_duration1(self, suffix: str) -> tuple[int, str]:
1102 """Like :meth:`suffix_to_duration0` but expects an underscore prefix."""
1103 return self._suffix_to_duration(suffix, self._suffix_regex1)
1105 @staticmethod
1106 def _suffix_to_duration(suffix: str, regex: re.Pattern) -> tuple[int, str]:
1107 """Example: Converts '2 hourly' to (2, 'hourly') and 'hourly' to (1, 'hourly')."""
1108 if match := regex.fullmatch(suffix):
1109 duration_amount: int = int(match.group(1)) if match.group(1) else 1
1110 assert duration_amount > 0
1111 duration_unit: str = match.group(2)
1112 return duration_amount, duration_unit
1113 else:
1114 return 0, ""
1116 def label_milliseconds(self, snapshot: str) -> int:
1117 """Returns duration encoded in ``snapshot`` suffix, in milliseconds."""
1118 i = snapshot.rfind("_")
1119 snapshot = "" if i < 0 else snapshot[i + 1 :]
1120 duration_amount, duration_unit = self._suffix_to_duration(snapshot, self._suffix_regex0)
1121 return duration_amount * self.suffix_milliseconds.get(duration_unit, 0)
1124#############################################################################
1125@final
1126class JobStats:
1127 """Simple thread-safe counters summarizing job progress."""
1129 def __init__(self, jobs_all: int) -> None:
1130 assert jobs_all >= 0
1131 self.lock: Final[threading.Lock] = threading.Lock()
1132 self.jobs_all: int = jobs_all
1133 self.jobs_started: int = 0
1134 self.jobs_completed: int = 0
1135 self.jobs_failed: int = 0
1136 self.jobs_running: int = 0
1137 self.sum_elapsed_nanos: int = 0
1138 self.started_job_names: Final[set[str]] = set()
1140 def submit_job(self, job_name: str) -> str:
1141 """Counts a job submission."""
1142 with self.lock:
1143 self.jobs_started += 1
1144 self.jobs_running += 1
1145 self.started_job_names.add(job_name)
1146 return str(self)
1148 def complete_job(self, failed: bool, elapsed_nanos: int) -> str:
1149 """Counts a job completion."""
1150 assert elapsed_nanos >= 0
1151 with self.lock:
1152 self.jobs_running -= 1
1153 self.jobs_completed += 1
1154 self.jobs_failed += 1 if failed else 0
1155 self.sum_elapsed_nanos += elapsed_nanos
1156 msg = str(self)
1157 assert self.sum_elapsed_nanos >= 0, msg
1158 assert self.jobs_running >= 0, msg
1159 assert self.jobs_failed >= 0, msg
1160 assert self.jobs_failed <= self.jobs_completed, msg
1161 assert self.jobs_completed <= self.jobs_started, msg
1162 assert self.jobs_started <= self.jobs_all, msg
1163 return msg
1165 def __repr__(self) -> str:
1166 def pct(number: int) -> str:
1167 """Returns percentage string relative to total jobs."""
1168 return percent(number, total=self.jobs_all, print_total=True)
1170 al, started, completed, failed = self.jobs_all, self.jobs_started, self.jobs_completed, self.jobs_failed
1171 running = self.jobs_running
1172 t = "avg_completion_time:" + human_readable_duration(self.sum_elapsed_nanos / max(1, completed))
1173 return f"all:{al}, started:{pct(started)}, completed:{pct(completed)}, failed:{pct(failed)}, running:{running}, {t}"
1176#############################################################################
1177_TComparableSelf = TypeVar("_TComparableSelf")
1180class Comparable(Protocol):
1181 """Partial ordering protocol."""
1183 def __lt__(self: _TComparableSelf, other: _TComparableSelf, /) -> bool: ...
1186TComparable = TypeVar("TComparable", bound=Comparable) # Generic type variable for elements stored in UpdatablePriorityQueue
1189@final
1190class UpdatablePriorityQueue(Generic[TComparable]):
1191 """A priority queue that can handle updates to the priority of any element that is already contained in the queue, and
1192 does so very efficiently if there are a small number of elements in the queue (no more than thousands), as is the case
1193 for us.
1195 Could be implemented using a SortedList via https://github.com/grantjenks/python-sortedcontainers or using an indexed
1196 priority queue via
1197 https://github.com/nvictus/pqdict.
1198 But, to avoid an external dependency, is actually implemented
1199 using a simple yet effective binary search-based sorted list that can handle updates to the priority of elements that
1200 are already contained in the queue, via removal of the element, followed by update of the element, followed by
1201 (re)insertion. Duplicate elements (if any) are maintained in their order of insertion relative to other duplicates.
1202 """
1204 def __init__(self, reverse: bool = False) -> None:
1205 """Creates an empty queue; sort order flips when ``reverse`` is True."""
1206 self._lst: Final[list[TComparable]] = []
1207 self._reverse: Final[bool] = reverse
1209 def clear(self) -> None:
1210 """Removes all elements from the queue."""
1211 self._lst.clear()
1213 def push(self, element: TComparable) -> None:
1214 """Inserts ``element`` while maintaining sorted order."""
1215 bisect.insort(self._lst, element)
1217 def pop(self) -> TComparable:
1218 """Removes and returns the smallest (or largest if reverse == True) element from the queue."""
1219 return self._lst.pop() if self._reverse else self._lst.pop(0)
1221 def peek(self) -> TComparable:
1222 """Returns the smallest (or largest if reverse == True) element without removing it."""
1223 return self._lst[-1] if self._reverse else self._lst[0]
1225 def remove(self, element: TComparable) -> bool:
1226 """Removes the first occurrence (in insertion order aka FIFO) of ``element`` and returns True if it was present."""
1227 lst = self._lst
1228 i = bisect.bisect_left(lst, element)
1229 is_contained = i < len(lst) and lst[i] == element
1230 if is_contained:
1231 del lst[i] # is an optimized memmove()
1232 return is_contained
1234 def __len__(self) -> int:
1235 """Returns the number of queued elements."""
1236 return len(self._lst)
1238 def __contains__(self, element: TComparable) -> bool:
1239 """Returns ``True`` if ``element`` is present."""
1240 lst = self._lst
1241 i = bisect.bisect_left(lst, element)
1242 return i < len(lst) and lst[i] == element
1244 def __iter__(self) -> Iterator[TComparable]:
1245 """Iterates over queued elements in priority order."""
1246 return reversed(self._lst) if self._reverse else iter(self._lst)
1248 def __repr__(self) -> str:
1249 """Representation showing queue contents in current order."""
1250 return repr(list(reversed(self._lst))) if self._reverse else repr(self._lst)
1253###############################################################################
1254@final
1255class SortedInterner(Generic[TComparable]):
1256 """Same as sys.intern() except that it isn't global and that it assumes the input list is sorted (for binary search)."""
1258 def __init__(self, sorted_list: list[TComparable]) -> None:
1259 self._lst: Final[list[TComparable]] = sorted_list
1261 def interned(self, element: TComparable) -> TComparable:
1262 """Returns the interned (aka deduped) item if an equal item is contained, else returns the non-interned item."""
1263 lst = self._lst
1264 i = binary_search(lst, element)
1265 return lst[i] if i >= 0 else element
1267 def __contains__(self, element: TComparable) -> bool:
1268 """Returns ``True`` if ``element`` is present."""
1269 return binary_search(self._lst, element) >= 0
1272def binary_search(sorted_list: list[TComparable], item: TComparable) -> int:
1273 """Java-style binary search; Returns index >= 0 if an equal item is found in list, else '-insertion_point-1'; If it
1274 returns index >= 0, the index will be the left-most index in case multiple such equal items are contained."""
1275 i = bisect.bisect_left(sorted_list, item)
1276 return i if i < len(sorted_list) and sorted_list[i] == item else -i - 1
1279###############################################################################
1280_S = TypeVar("_S")
1283@final
1284class HashedInterner(Generic[_S]):
1285 """Same as sys.intern() except that it isn't global and can also be used for types other than str."""
1287 def __init__(self, items: Iterable[_S] = frozenset()) -> None:
1288 self._items: Final[dict[_S, _S]] = {v: v for v in items}
1290 def intern(self, item: _S) -> _S:
1291 """Interns the given item."""
1292 return self._items.setdefault(item, item)
1294 def interned(self, item: _S) -> _S:
1295 """Returns the interned (aka deduped) item if an equal item is contained, else returns the non-interned item."""
1296 return self._items.get(item, item)
1298 def __contains__(self, item: _S) -> bool:
1299 return item in self._items
1302#############################################################################
1303@final
1304class SynchronizedBool:
1305 """Thread-safe wrapper around a regular bool."""
1307 def __init__(self, val: bool) -> None:
1308 assert isinstance(val, bool)
1309 self._lock: Final[threading.Lock] = threading.Lock()
1310 self._value: bool = val
1312 @property
1313 def value(self) -> bool:
1314 """Returns the current boolean value."""
1315 with self._lock:
1316 return self._value
1318 @value.setter
1319 def value(self, new_value: bool) -> None:
1320 """Atomically assign ``new_value``."""
1321 with self._lock:
1322 self._value = new_value
1324 def get_and_set(self, new_value: bool) -> bool:
1325 """Swaps in ``new_value`` and return the previous value."""
1326 with self._lock:
1327 old_value = self._value
1328 self._value = new_value
1329 return old_value
1331 def compare_and_set(self, expected_value: bool, new_value: bool) -> bool:
1332 """Sets to ``new_value`` only if current value equals ``expected_value``."""
1333 with self._lock:
1334 eq: bool = self._value == expected_value
1335 if eq:
1336 self._value = new_value
1337 return eq
1339 def __bool__(self) -> bool:
1340 return self.value
1342 def __repr__(self) -> str:
1343 return repr(self.value)
1345 def __str__(self) -> str:
1346 return str(self.value)
1349#############################################################################
1350_K = TypeVar("_K")
1351_V = TypeVar("_V")
1354@final
1355class SynchronizedDict(Generic[_K, _V]):
1356 """Thread-safe wrapper around a regular dict."""
1358 def __init__(self, val: dict[_K, _V]) -> None:
1359 assert isinstance(val, dict)
1360 self._lock: Final[threading.Lock] = threading.Lock()
1361 self._dict: Final[dict[_K, _V]] = val
1363 def __getitem__(self, key: _K) -> _V:
1364 with self._lock:
1365 return self._dict[key]
1367 def __setitem__(self, key: _K, value: _V) -> None:
1368 with self._lock:
1369 self._dict[key] = value
1371 def __delitem__(self, key: _K) -> None:
1372 with self._lock:
1373 self._dict.pop(key)
1375 def __contains__(self, key: _K) -> bool:
1376 with self._lock:
1377 return key in self._dict
1379 def __len__(self) -> int:
1380 with self._lock:
1381 return len(self._dict)
1383 def __repr__(self) -> str:
1384 with self._lock:
1385 return repr(self._dict)
1387 def __str__(self) -> str:
1388 with self._lock:
1389 return str(self._dict)
1391 def get(self, key: _K, default: _V | None = None) -> _V | None:
1392 """Returns ``self[key]`` or ``default`` if missing."""
1393 with self._lock:
1394 return self._dict.get(key, default)
1396 def pop(self, key: _K, default: _V | None = None) -> _V | None:
1397 """Removes ``key`` and returns its value."""
1398 with self._lock:
1399 return self._dict.pop(key, default)
1401 def clear(self) -> None:
1402 """Removes all items atomically."""
1403 with self._lock:
1404 self._dict.clear()
1406 def items(self) -> ItemsView[_K, _V]:
1407 """Returns a snapshot of dictionary items."""
1408 with self._lock:
1409 return self._dict.copy().items()
1412#############################################################################
1413@final
1414class InterruptibleSleep:
1415 """Provides a sleep(timeout) function that can be interrupted by another thread; The underlying lock is configurable."""
1417 def __init__(self, lock: threading.Lock | None = None) -> None:
1418 self._is_stopping: bool = False
1419 self._lock: Final[threading.Lock] = lock if lock is not None else threading.Lock()
1420 self._condition: Final[threading.Condition] = threading.Condition(self._lock)
1422 def sleep(self, duration_nanos: int) -> bool:
1423 """Delays the current thread by the given number of nanoseconds; Returns True if the sleep got interrupted;
1424 Equivalent to threading.Event.wait()."""
1425 end_time_nanos: int = time.monotonic_ns() + duration_nanos
1426 with self._lock:
1427 while not self._is_stopping:
1428 diff_nanos: int = end_time_nanos - time.monotonic_ns()
1429 if diff_nanos <= 0:
1430 return False
1431 self._condition.wait(timeout=diff_nanos / 1_000_000_000) # release, then block until notified or timeout
1432 return True
1434 def interrupt(self) -> None:
1435 """Wakes sleeping threads and makes any future sleep()s a no-op; Equivalent to threading.Event.set()."""
1436 with self._lock:
1437 if not self._is_stopping:
1438 self._is_stopping = True
1439 self._condition.notify_all()
1441 def reset(self) -> None:
1442 """Makes any future sleep()s no longer a no-op; Equivalent to threading.Event.clear()."""
1443 with self._lock:
1444 self._is_stopping = False
1447#############################################################################
1448@final
1449class SynchronousExecutor(Executor):
1450 """Executor that runs tasks inline in the calling thread, sequentially."""
1452 def __init__(self) -> None:
1453 self._shutdown: bool = False
1455 # pyrefly: ignore[bad-override] # Python 3.9-compatible executor signature
1456 def submit(self, fn: Callable[..., _R_], /, *args: Any, **kwargs: Any) -> Future[_R_]:
1457 """Executes `fn(*args, **kwargs)` immediately and returns its Future."""
1458 future: Future[_R_] = Future()
1459 if self._shutdown:
1460 raise RuntimeError("cannot schedule new futures after shutdown")
1461 try:
1462 result: _R_ = fn(*args, **kwargs)
1463 except BaseException as exc:
1464 future.set_exception(exc)
1465 else:
1466 future.set_result(result)
1467 return future
1469 def shutdown(self, wait: bool = True, *, cancel_futures: bool = False) -> None:
1470 """Prevents new submissions; no worker resources to join/cleanup."""
1471 self._shutdown = True
1473 @classmethod
1474 def executor_for(cls, max_workers: int) -> Executor:
1475 """Factory returning a SynchronousExecutor if 0 <= max_workers <= 1; else a ThreadPoolExecutor."""
1476 return cls() if 0 <= max_workers <= 1 else ThreadPoolExecutor(max_workers=max_workers)
1479#############################################################################
1480@final
1481class _XFinally(contextlib.AbstractContextManager):
1482 """Context manager ensuring cleanup code executes after ``with`` blocks."""
1484 def __init__(self, cleanup: Callable[[], None]) -> None:
1485 """Records the callable to run upon exit."""
1486 self._cleanup: Final = cleanup # Zero-argument callable executed after the `with` block exits.
1488 def __exit__(
1489 self, exc_type: type[BaseException] | None, exc: BaseException | None, tb: types.TracebackType | None
1490 ) -> Literal[False]:
1491 """Runs cleanup and propagate any exceptions appropriately."""
1492 try:
1493 self._cleanup()
1494 except BaseException as cleanup_exc:
1495 if exc is None:
1496 raise # No main error --> propagate cleanup error normally
1497 # Both failed
1498 # if sys.version_info >= (3, 11):
1499 # raise ExceptionGroup("main error and cleanup error", [exc, cleanup_exc]) from None
1500 # <= 3.10: attach so it shows up in traceback but doesn't mask
1501 exc.__context__ = cleanup_exc
1502 return False # reraise original exception
1503 return False # propagate main exception if any
1506def xfinally(cleanup: Callable[[], None]) -> _XFinally:
1507 """Usage: with xfinally(lambda: cleanup()): ...
1508 Returns a context manager that guarantees that cleanup() runs on exit and guarantees any error in cleanup() will never
1509 mask an exception raised earlier inside the body of the `with` block, while still surfacing both problems when possible.
1511 Problem it solves
1512 -----------------
1513 A naive ``try ... finally`` may lose the original exception:
1515 try:
1516 work()
1517 finally:
1518 cleanup() # <-- if this raises an exception, it replaces the real error!
1520 `_XFinally` preserves exception priority:
1522 * Body raises, cleanup succeeds --> original body exception is re-raised.
1523 * Body raises, cleanup also raises --> re-raises body exception; cleanup exception is linked via ``__context__``.
1524 * Body succeeds, cleanup raises --> cleanup exception propagates normally.
1526 Example:
1527 -------
1528 >>> with xfinally(lambda: release_resources()): # doctest: +SKIP
1529 ... run_tasks()
1531 The single *with* line replaces verbose ``try/except/finally`` boilerplate while preserving full error information.
1532 """
1533 return _XFinally(cleanup)