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

utilities/atomicwrites.py

@@ -1,10 +1,11 @@
 from __future__ import annotations
 
+import gzip
 import shutil
 from contextlib import ExitStack, contextmanager
 from dataclasses import dataclass
 from pathlib import Path
-from shutil import rmtree
+from shutil import copyfileobj, rmtree
 from typing import TYPE_CHECKING, assert_never, override
 
 from atomicwrites import replace_atomic
@@ -56,7 +57,7 @@ def move(
             raise ImpossibleCaseError(
                 case=[f"{source.is_file()=}", f"{source.is_dir()=}"]
             )
-        case _ as never:
+        case never:
             assert_never(never)
 
 
@@ -111,21 +112,32 @@ def move_many(*paths: tuple[PathLike, PathLike], overwrite: bool = False) -> Non
 
 
 @contextmanager
-def writer(path: PathLike, /, *, overwrite: bool = False) -> Iterator[Path]:
+def writer(
+    path: PathLike, /, *, compress: bool = False, overwrite: bool = False
+) -> Iterator[Path]:
     """Yield a path for atomically writing files to disk."""
     path = Path(path)
     parent = path.parent
     parent.mkdir(parents=True, exist_ok=True)
     name = path.name
     with TemporaryDirectory(suffix=".tmp", prefix=name, dir=parent) as temp_dir:
-        temp_path = Path(temp_dir, name)
+        temp_path1 = Path(temp_dir, name)
         try:
-            yield temp_path
+            yield temp_path1
         except KeyboardInterrupt:
             rmtree(temp_dir)
         else:
+            if compress:
+                temp_path2 = Path(temp_dir, f"{name}.gz")
+                with (
+                    temp_path1.open("rb") as source,
+                    gzip.open(temp_path2, mode="wb") as dest,
+                ):
+                    copyfileobj(source, dest)
+            else:
+                temp_path2 = temp_path1
             try:
-                move(temp_path, path, overwrite=overwrite)
+                move(temp_path2, path, overwrite=overwrite)
             except _MoveSourceNotFoundError as error:
                 raise _WriterTemporaryPathEmptyError(temp_path=error.source) from None
             except _MoveFileExistsError as error:

utilities/atools.py

@@ -1,44 +1,99 @@
 from __future__ import annotations
 
 from collections.abc import Callable
-from typing import TYPE_CHECKING, ParamSpec, TypeVar
+from functools import partial
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, cast, overload
 
-from atools import memoize
+import atools
+from whenever import TimeDelta
 
-from utilities.datetime import datetime_duration_to_timedelta
-from utilities.types import Coroutine1
+from utilities.types import Coro, PathLike
 
 if TYPE_CHECKING:
-    import datetime as dt
+    from atools._memoize_decorator import Keygen, Pickler
 
-    from utilities.types import Duration
 
+type _Key[**P, T] = tuple[Callable[P, Coro[T]], TimeDelta]
+_MEMOIZED_FUNCS: dict[_Key, Callable[..., Coro[Any]]] = {}
 
-_P = ParamSpec("_P")
-_R = TypeVar("_R")
-_AsyncFunc = Callable[_P, Coroutine1[_R]]
-type _Key = tuple[_AsyncFunc, dt.timedelta]
-_MEMOIZED_FUNCS: dict[_Key, _AsyncFunc] = {}
 
-
-async def call_memoized(
-    func: _AsyncFunc[_P, _R],
-    refresh: Duration | None = None,
+async def call_memoized[**P, T](
+    func: Callable[P, Coro[T]],
+    refresh: TimeDelta | None = None,
     /,
-    *args: _P.args,
-    **kwargs: _P.kwargs,
-) -> _R:
+    *args: P.args,
+    **kwargs: P.kwargs,
+) -> T:
     """Call an asynchronous function, with possible memoization."""
     if refresh is None:
         return await func(*args, **kwargs)
-    timedelta = datetime_duration_to_timedelta(refresh)
-    key: _Key = (func, timedelta)
-    memoized_func: _AsyncFunc[_P, _R]
+    key: _Key = (func, refresh)
+    memoized_func: Callable[P, Coro[T]]
     try:
         memoized_func = _MEMOIZED_FUNCS[key]
     except KeyError:
-        memoized_func = _MEMOIZED_FUNCS[(key)] = memoize(duration=refresh)(func)
+        memoized_func = _MEMOIZED_FUNCS[(key)] = memoize(duration=refresh.in_seconds())(
+            func
+        )
     return await memoized_func(*args, **kwargs)
 
 
+##
+
+
+@overload
+def memoize[F: Callable[..., Coro[Any]]](
+    func: F,
+    /,
+    *,
+    db_path: PathLike | None = None,
+    duration: float | TimeDelta | None = None,
+    keygen: Keygen | None = None,
+    pickler: Pickler | None = None,
+    size: int | None = None,
+) -> F: ...
+@overload
+def memoize[F: Callable[..., Coro[Any]]](
+    func: None = None,
+    /,
+    *,
+    db_path: PathLike | None = None,
+    duration: float | TimeDelta | None = None,
+    keygen: Keygen | None = None,
+    pickler: Pickler | None = None,
+    size: int | None = None,
+) -> Callable[[F], F]: ...
+def memoize[F: Callable[..., Coro[Any]]](
+    func: F | None = None,
+    /,
+    *,
+    db_path: PathLike | None = None,
+    duration: float | TimeDelta | None = None,
+    keygen: Keygen | None = None,
+    pickler: Pickler | None = None,
+    size: int | None = None,
+) -> F | Callable[[F], F]:
+    """Memoize a function."""
+    if func is None:
+        result = partial(
+            memoize,
+            db_path=db_path,
+            duration=duration,
+            keygen=keygen,
+            pickler=pickler,
+            size=size,
+        )
+        return cast("Callable[[F], F]", result)
+    return atools.memoize(
+        db_path=None if db_path is None else Path(db_path),
+        duration=duration.py_timedelta()
+        if isinstance(duration, TimeDelta)
+        else duration,
+        keygen=keygen,
+        pickler=pickler,
+        size=size,
+    )(func)
+
+
 __all__ = ["call_memoized"]

utilities/cachetools.py

@@ -1,39 +1,31 @@
 from __future__ import annotations
 
-from collections.abc import Callable, Iterable, Iterator, MutableSet
+from collections.abc import Callable, Hashable, Iterable, Iterator, MutableSet
 from math import inf
 from time import monotonic
-from typing import TYPE_CHECKING, Any, TypeVar, override
+from typing import TYPE_CHECKING, Any, cast, override
 
 import cachetools
 from cachetools.func import ttl_cache
 
-from utilities.datetime import datetime_duration_to_float
-
 if TYPE_CHECKING:
-    from utilities.types import Duration, TCallable
-
-_K = TypeVar("_K")
-_T = TypeVar("_T")
-_V = TypeVar("_V")
+    from whenever import TimeDelta
 
 
-class TTLCache(cachetools.TTLCache[_K, _V]):
+class TTLCache[K: Hashable, V](cachetools.TTLCache[K, V]):
     """A TTL-cache."""
 
     def __init__(
         self,
         *,
         max_size: int | None = None,
-        max_duration: Duration | None = None,
+        max_duration: TimeDelta | None = None,
         timer: Callable[[], float] = monotonic,
         get_size_of: Callable[[Any], int] | None = None,
     ) -> None:
         super().__init__(
             maxsize=inf if max_size is None else max_size,
-            ttl=inf
-            if max_duration is None
-            else datetime_duration_to_float(max_duration),
+            ttl=inf if max_duration is None else max_duration.in_seconds(),
             timer=timer,
             getsizeof=get_size_of,
         )
@@ -42,19 +34,19 @@ class TTLCache(cachetools.TTLCache[_K, _V]):
 ##
 
 
-class TTLSet(MutableSet[_T]):
+class TTLSet[T: Hashable](MutableSet[T]):
     """A TTL-set."""
 
-    _cache: TTLCache[_T, None]
+    _cache: TTLCache[T, None]
 
     @override
     def __init__(
         self,
-        iterable: Iterable[_T] | None = None,
+        iterable: Iterable[T] | None = None,
         /,
         *,
         max_size: int | None = None,
-        max_duration: Duration | None = None,
+        max_duration: TimeDelta | None = None,
         timer: Callable[[], float] = monotonic,
         get_size_of: Callable[[Any], int] | None = None,
     ) -> None:
@@ -73,7 +65,7 @@ class TTLSet(MutableSet[_T]):
         return self._cache.__contains__(x)
 
     @override
-    def __iter__(self) -> Iterator[_T]:
+    def __iter__(self) -> Iterator[T]:
         return self._cache.__iter__()
 
     @override
@@ -89,30 +81,33 @@ class TTLSet(MutableSet[_T]):
         return set(self._cache).__str__()
 
     @override
-    def add(self, value: _T) -> None:
+    def add(self, value: T) -> None:
         self._cache[value] = None
 
     @override
-    def discard(self, value: _T) -> None:
+    def discard(self, value: T) -> None:
         del self._cache[value]
 
 
 ##
 
 
-def cache(
+def cache[F: Callable](
     *,
     max_size: int | None = None,
-    max_duration: Duration | None = None,
+    max_duration: TimeDelta | None = None,
     timer: Callable[[], float] = monotonic,
     typed_: bool = False,
-) -> Callable[[TCallable], TCallable]:
+) -> Callable[[F], F]:
     """Decorate a function with `max_size` and/or `ttl` settings."""
-    return ttl_cache(
-        maxsize=inf if max_size is None else max_size,
-        ttl=inf if max_duration is None else datetime_duration_to_float(max_duration),
-        timer=timer,
-        typed=typed_,
+    return cast(
+        "F",
+        ttl_cache(
+            maxsize=max_size,
+            ttl=inf if max_duration is None else max_duration.in_seconds(),
+            timer=timer,
+            typed=typed_,
+        ),
     )