dycw-utilities 0.135.0__py3-none-any.whl → 0.178.1__py3-none-any.whl

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)

utilities/dataclasses.py

@@ -6,11 +6,7 @@ from dataclasses import MISSING, dataclass, field, fields, replace
 from typing import TYPE_CHECKING, Any, Literal, assert_never, overload, override
 
 from utilities.errors import ImpossibleCaseError
-from utilities.functions import (
-    get_class_name,
-    is_dataclass_class,
-    is_dataclass_instance,
-)
+from utilities.functions import get_class_name
 from utilities.iterables import (
     OneStrEmptyError,
     OneStrNonUniqueError,
@@ -25,7 +21,7 @@ from utilities.parse import (
     serialize_object,
 )
 from utilities.re import ExtractGroupError, extract_group
-from utilities.sentinel import Sentinel, sentinel
+from utilities.sentinel import Sentinel, is_sentinel, sentinel
 from utilities.text import (
     BRACKETS,
     LIST_SEPARATOR,
@@ -34,8 +30,8 @@ from utilities.text import (
     _SplitKeyValuePairsSplitError,
     split_key_value_pairs,
 )
-from utilities.types import SupportsLT
-from utilities.typing import get_type_hints
+from utilities.types import MaybeType, SupportsLT
+from utilities.typing import get_type_hints, is_dataclass_class, is_dataclass_instance
 
 if TYPE_CHECKING:
     from collections.abc import Callable, Iterable, Iterator
@@ -214,7 +210,7 @@ def is_nullable_lt[T: SupportsLT](x: T | None, y: T | None, /) -> bool | None:
             return True
         case 0:
             return None
-        case _ as never:
+        case never:
             assert_never(never)
 
 
@@ -275,8 +271,7 @@ def mapping_to_dataclass[T: Dataclass](
     default = {
         f.name
         for f in fields_use
-        if (not isinstance(f.default, Sentinel))
-        or (not isinstance(f.default_factory, Sentinel))
+        if (not is_sentinel(f.default)) or (not is_sentinel(f.default_factory))
     }
     have = set(field_names_to_values) | default
     missing = {f.name for f in fields_use} - have
@@ -434,12 +429,10 @@ def replace_non_sentinel[T: Dataclass](
     """Replace attributes on a dataclass, filtering out sentinel values."""
     if in_place:
         for k, v in kwargs.items():
-            if not isinstance(v, Sentinel):
+            if not is_sentinel(v):
                 setattr(obj, k, v)
         return None
-    return replace(
-        obj, **{k: v for k, v in kwargs.items() if not isinstance(v, Sentinel)}
-    )
+    return replace(obj, **{k: v for k, v in kwargs.items() if not is_sentinel(v)})
 
 
 ##
@@ -520,7 +513,7 @@ def parse_dataclass[T: Dataclass](
             )
         case Mapping() as keys_to_serializes:
             ...
-        case _ as never:
+        case never:
             assert_never(never)
     fields = list(
         yield_fields(
@@ -833,7 +826,7 @@ def yield_fields(
     warn_name_errors: bool = False,
 ) -> Iterator[_YieldFieldsClass[Any]]: ...
 def yield_fields(
-    obj: Dataclass | type[Dataclass],
+    obj: MaybeType[Dataclass],
     /,
     *,
     globalns: StrMapping | None = None,
@@ -912,17 +905,11 @@ class _YieldFieldsInstance[T]:
         extra: Mapping[type[U], Callable[[U, U], bool]] | None = None,
     ) -> bool:
         """Check if the field value equals its default."""
-        if isinstance(self.default, Sentinel) and isinstance(
-            self.default_factory, Sentinel
-        ):
+        if is_sentinel(self.default) and is_sentinel(self.default_factory):
             return False
-        if (not isinstance(self.default, Sentinel)) and isinstance(
-            self.default_factory, Sentinel
-        ):
+        if (not is_sentinel(self.default)) and is_sentinel(self.default_factory):
             expected = self.default
-        elif isinstance(self.default, Sentinel) and (
-            not isinstance(self.default_factory, Sentinel)
-        ):
+        elif is_sentinel(self.default) and (not is_sentinel(self.default_factory)):
             expected = self.default_factory()
         else:  # pragma: no cover
             raise ImpossibleCaseError(
@@ -1002,7 +989,7 @@ def _empty_error_str_core(
             return f"any field starting with {key!r}"
         case True, False:
             return f"any field starting with {key!r} (modulo case)"
-        case _ as never:
+        case never:
             assert_never(never)
 
 
@@ -1043,7 +1030,7 @@ def _non_unique_error_str_core(
             head_msg = f"exactly one field starting with {key!r}"
         case True, False:
             head_msg = f"exactly one field starting with {key!r} (modulo case)"
-        case _ as never:
+        case never:
             assert_never(never)
     return f"{head_msg}; got {first!r}, {second!r} and perhaps more"