dycw-utilities 0.129.10__py3-none-any.whl → 0.175.17__py3-none-any.whl

utilities/concurrent.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 
 from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor
 from functools import partial
-from typing import TYPE_CHECKING, Any, TypeVar, assert_never
+from typing import TYPE_CHECKING, Any, assert_never
 
 from utilities.iterables import apply_to_tuple
 from utilities.os import get_cpu_use
@@ -15,11 +15,8 @@ if TYPE_CHECKING:
     from utilities.os import IntOrAll
 
 
-_T = TypeVar("_T")
-
-
-def concurrent_map(
-    func: Callable[..., _T],
+def concurrent_map[T](
+    func: Callable[..., T],
     /,
     *iterables: Iterable[Any],
     parallelism: Parallelism = "processes",
@@ -31,7 +28,7 @@ def concurrent_map(
     thread_name_prefix: str = "",
     timeout: float | None = None,
     chunksize: int = 1,
-) -> list[_T]:
+) -> list[T]:
     """Concurrent map."""
     return concurrent_starmap(
         func,
@@ -51,8 +48,8 @@ def concurrent_map(
 ##
 
 
-def concurrent_starmap(
-    func: Callable[..., _T],
+def concurrent_starmap[T](
+    func: Callable[..., T],
     iterable: Iterable[tuple[Any, ...]],
     /,
     *,
@@ -65,7 +62,7 @@ def concurrent_starmap(
     thread_name_prefix: str = "",
     timeout: float | None = None,
     chunksize: int = 1,
-) -> list[_T]:
+) -> list[T]:
     """Concurrent map."""
     max_workers_use = get_cpu_use(n=max_workers)
     apply = partial(apply_to_tuple, func)
@@ -87,7 +84,7 @@ def concurrent_starmap(
                 initargs=initargs,
             ) as pool:
                 result = pool.map(apply, iterable, timeout=timeout, chunksize=chunksize)
-        case _ as never:
+        case never:
             assert_never(never)
     return list(result)
 

utilities/contextlib.py

@@ -1,28 +1,223 @@
 from __future__ import annotations
 
 import re
-from contextlib import contextmanager
-from typing import TYPE_CHECKING
+from asyncio import create_task, get_event_loop
+from contextlib import (
+    _AsyncGeneratorContextManager,
+    _GeneratorContextManager,
+    asynccontextmanager,
+    contextmanager,
+)
+from functools import partial, wraps
+from signal import SIGABRT, SIGFPE, SIGILL, SIGINT, SIGSEGV, SIGTERM, getsignal, signal
+from typing import TYPE_CHECKING, Any, assert_never, cast, overload
 
 if TYPE_CHECKING:
-    from collections.abc import Iterator
-    from types import TracebackType
+    from collections.abc import AsyncIterator, Callable, Iterator
+    from signal import _HANDLER, _SIGNUM
+    from types import FrameType
 
 
-class NoOpContextManager:
-    """Context-manager for no-op."""
+@overload
+def enhanced_context_manager[**P, T_co](
+    func: Callable[P, Iterator[T_co]],
+    /,
+    *,
+    sigabrt: bool = True,
+    sigfpe: bool = True,
+    sigill: bool = True,
+    sigint: bool = True,
+    sigsegv: bool = True,
+    sigterm: bool = True,
+) -> Callable[P, _GeneratorContextManager[T_co]]: ...
+@overload
+def enhanced_context_manager[**P, T_co](
+    func: None = None,
+    /,
+    *,
+    sigabrt: bool = True,
+    sigfpe: bool = True,
+    sigill: bool = True,
+    sigint: bool = True,
+    sigsegv: bool = True,
+    sigterm: bool = True,
+) -> Callable[
+    [Callable[P, Iterator[T_co]]], Callable[P, _GeneratorContextManager[T_co]]
+]: ...
+def enhanced_context_manager[**P, T_co](
+    func: Callable[P, Iterator[T_co]] | None = None,
+    /,
+    *,
+    sigabrt: bool = True,
+    sigfpe: bool = True,
+    sigill: bool = True,
+    sigint: bool = True,
+    sigsegv: bool = True,
+    sigterm: bool = True,
+) -> (
+    Callable[P, _GeneratorContextManager[T_co]]
+    | Callable[
+        [Callable[P, Iterator[T_co]]], Callable[P, _GeneratorContextManager[T_co]]
+    ]
+):
+    if func is None:
+        result = partial(
+            enhanced_context_manager,
+            sigabrt=sigabrt,
+            sigfpe=sigfpe,
+            sigill=sigill,
+            sigint=sigint,
+            sigsegv=sigsegv,
+            sigterm=sigterm,
+        )
+        return cast(
+            "Callable[[Callable[P, Iterator[T_co]]], Callable[P, _GeneratorContextManager[T_co]]]",
+            result,
+        )
+    make_gcm = contextmanager(func)
 
-    def __enter__(self) -> None:
-        return None
+    @contextmanager
+    @wraps(func)
+    def wrapped(*args: P.args, **kwargs: P.kwargs) -> Iterator[T_co]:
+        gcm = make_gcm(*args, **kwargs)
+        sigabrt0 = _swap_handler(SIGABRT, gcm) if sigabrt else None
+        sigfpe0 = _swap_handler(SIGFPE, gcm) if sigfpe else None
+        sigill0 = _swap_handler(SIGILL, gcm) if sigill else None
+        sigint0 = _swap_handler(SIGINT, gcm) if sigint else None
+        sigsegv0 = _swap_handler(SIGSEGV, gcm) if sigsegv else None
+        sigterm0 = _swap_handler(SIGTERM, gcm) if sigterm else None
+        try:
+            with gcm as value:
+                yield value
+        finally:
+            _ = signal(SIGABRT, sigabrt0) if sigabrt else None
+            _ = signal(SIGFPE, sigfpe0) if sigfpe else None
+            _ = signal(SIGILL, sigill0) if sigill else None
+            _ = signal(SIGINT, sigint0) if sigint else None
+            _ = signal(SIGSEGV, sigsegv0) if sigsegv else None
+            _ = signal(SIGTERM, sigterm0) if sigterm else None
 
-    def __exit__(
-        self,
-        exc_type: type[BaseException] | None,
-        exc_val: BaseException | None,
-        traceback: TracebackType | None,
-    ) -> bool:
-        _ = (exc_type, exc_val, traceback)
-        return False
+    return wrapped
+
+
+@overload
+def enhanced_async_context_manager[**P, T_co](
+    func: Callable[P, AsyncIterator[T_co]],
+    /,
+    *,
+    sigabrt: bool = True,
+    sigfpe: bool = True,
+    sigill: bool = True,
+    sigint: bool = True,
+    sigsegv: bool = True,
+    sigterm: bool = True,
+) -> Callable[P, _AsyncGeneratorContextManager[T_co]]: ...
+@overload
+def enhanced_async_context_manager[**P, T_co](
+    func: None = None,
+    /,
+    *,
+    sigabrt: bool = True,
+    sigfpe: bool = True,
+    sigill: bool = True,
+    sigint: bool = True,
+    sigsegv: bool = True,
+    sigterm: bool = True,
+) -> Callable[
+    [Callable[P, AsyncIterator[T_co]]], Callable[P, _AsyncGeneratorContextManager[T_co]]
+]: ...
+def enhanced_async_context_manager[**P, T_co](
+    func: Callable[P, AsyncIterator[T_co]] | None = None,
+    /,
+    *,
+    sigabrt: bool = True,
+    sigfpe: bool = True,
+    sigill: bool = True,
+    sigint: bool = True,
+    sigsegv: bool = True,
+    sigterm: bool = True,
+) -> (
+    Callable[P, _AsyncGeneratorContextManager[T_co]]
+    | Callable[
+        [Callable[P, AsyncIterator[T_co]]],
+        Callable[P, _AsyncGeneratorContextManager[T_co]],
+    ]
+):
+    if func is None:
+        result = partial(
+            enhanced_async_context_manager,
+            sigabrt=sigabrt,
+            sigfpe=sigfpe,
+            sigill=sigill,
+            sigint=sigint,
+            sigsegv=sigsegv,
+            sigterm=sigterm,
+        )
+        return cast(
+            "Callable[[Callable[P, AsyncIterator[T_co]]], Callable[P, _AsyncGeneratorContextManager[T_co]]]",
+            result,
+        )
+    make_agcm = asynccontextmanager(func)
+
+    @asynccontextmanager
+    @wraps(func)
+    async def wrapped(*args: P.args, **kwargs: P.kwargs) -> AsyncIterator[T_co]:
+        agcm = make_agcm(*args, **kwargs)
+        sigabrt0 = _swap_handler(SIGABRT, agcm) if sigabrt else None
+        sigfpe0 = _swap_handler(SIGFPE, agcm) if sigfpe else None
+        sigill0 = _swap_handler(SIGILL, agcm) if sigill else None
+        sigint0 = _swap_handler(SIGINT, agcm) if sigint else None
+        sigsegv0 = _swap_handler(SIGSEGV, agcm) if sigsegv else None
+        sigterm0 = _swap_handler(SIGTERM, agcm) if sigterm else None
+        try:
+            async with agcm as value:
+                yield value
+        finally:
+            _ = signal(SIGABRT, sigabrt0) if sigabrt else None
+            _ = signal(SIGFPE, sigfpe0) if sigfpe else None
+            _ = signal(SIGILL, sigill0) if sigill else None
+            _ = signal(SIGINT, sigint0) if sigint else None
+            _ = signal(SIGSEGV, sigsegv0) if sigsegv else None
+            _ = signal(SIGTERM, sigterm0) if sigterm else None
+
+    return wrapped
+
+
+def _swap_handler(
+    signum: _SIGNUM,
+    obj: _GeneratorContextManager[Any, None, None]
+    | _AsyncGeneratorContextManager[Any, None],
+    /,
+) -> _HANDLER:
+    orig_handler = getsignal(signum)
+    new_handler = _make_handler(signum, obj)
+    _ = signal(signum, new_handler)
+    return orig_handler
+
+
+def _make_handler(
+    signum: _SIGNUM,
+    obj: _GeneratorContextManager[Any, None, None]
+    | _AsyncGeneratorContextManager[Any, None],
+    /,
+) -> Callable[[int, FrameType | None], None]:
+    orig_handler = getsignal(signum)
+
+    def new_handler(signum: int, frame: FrameType | None) -> None:
+        match obj:  # pragma: no cover
+            case _GeneratorContextManager() as gcm:
+                _ = gcm.__exit__(None, None, None)
+            case _AsyncGeneratorContextManager() as agcm:
+                loop = get_event_loop()
+                _ = loop.call_soon_threadsafe(
+                    create_task, agcm.__aexit__(None, None, None)
+                )
+            case never:
+                assert_never(never)
+        if callable(orig_handler):  # pragma: no cover
+            orig_handler(signum, frame)
+
+    return new_handler
 
 
 ##
@@ -41,4 +236,8 @@ def suppress_super_object_attribute_error() -> Iterator[None]:
             raise
 
 
-__all__ = ["NoOpContextManager", "suppress_super_object_attribute_error"]
+__all__ = [
+    "enhanced_async_context_manager",
+    "enhanced_context_manager",
+    "suppress_super_object_attribute_error",
+]

utilities/contextvars.py

@@ -1,6 +1,12 @@
 from __future__ import annotations
 
+from contextlib import contextmanager
 from contextvars import ContextVar
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from collections.abc import Iterator
+
 
 ##
 
@@ -19,4 +25,17 @@ def set_global_breakpoint() -> None:
     _ = _GLOBAL_BREAKPOINT.set(True)
 
 
-__all__ = ["global_breakpoint", "set_global_breakpoint"]
+##
+
+
+@contextmanager
+def yield_set_context(var: ContextVar[bool], /) -> Iterator[None]:
+    """Yield a context var as being set."""
+    token = var.set(True)
+    try:
+        yield
+    finally:
+        _ = var.reset(token)
+
+
+__all__ = ["global_breakpoint", "set_global_breakpoint", "yield_set_context"]

utilities/cryptography.py

@@ -11,18 +11,18 @@ _ENV_VAR = "FERNET_KEY"
 
 def encrypt(text: str, /, *, env_var: str = _ENV_VAR) -> bytes:
     """Encrypt a string."""
-    return get_fernet(env_var=env_var).encrypt(text.encode())
+    return get_fernet(env_var).encrypt(text.encode())
 
 
 def decrypt(text: bytes, /, *, env_var: str = _ENV_VAR) -> str:
     """Encrypt a string."""
-    return get_fernet(env_var=env_var).decrypt(text).decode()
+    return get_fernet(env_var).decrypt(text).decode()
 
 
 ##
 
 
-def get_fernet(*, env_var: str = _ENV_VAR) -> Fernet:
+def get_fernet(env_var: str = _ENV_VAR, /) -> Fernet:
     """Get the Fernet key."""
     if (key := getenv(env_var)) is None:
         raise GetFernetError(env_var=env_var)