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

utilities/polars_ols.py

@@ -6,8 +6,8 @@ from polars import Expr, Series, struct
 from polars_ols import RollingKwargs, compute_rolling_least_squares
 
 from utilities.errors import ImpossibleCaseError
-from utilities.functions import is_sequence_of
 from utilities.polars import concat_series, ensure_expr_or_series
+from utilities.typing import is_sequence_of
 
 if TYPE_CHECKING:
     from polars._typing import IntoExprColumn

utilities/postgres.py

@@ -0,0 +1,408 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from pathlib import Path
+from shutil import rmtree
+from typing import TYPE_CHECKING, Literal, assert_never, override
+
+from sqlalchemy import Table
+from sqlalchemy.orm import DeclarativeBase
+
+from utilities.asyncio import stream_command
+from utilities.docker import docker_exec_cmd
+from utilities.iterables import always_iterable
+from utilities.logging import to_logger
+from utilities.os import temp_environ
+from utilities.pathlib import ensure_suffix
+from utilities.sqlalchemy import extract_url, get_table_name
+from utilities.timer import Timer
+from utilities.types import PathLike
+
+if TYPE_CHECKING:
+    from sqlalchemy import URL
+
+    from utilities.sqlalchemy import TableOrORMInstOrClass
+    from utilities.types import (
+        LoggerLike,
+        MaybeCollection,
+        MaybeCollectionStr,
+        PathLike,
+    )
+
+
+type _PGDumpFormat = Literal["plain", "custom", "directory", "tar"]
+
+
+async def pg_dump(
+    url: URL,
+    path: PathLike,
+    /,
+    *,
+    docker_container: str | None = None,
+    format_: _PGDumpFormat = "plain",
+    jobs: int | None = None,
+    data_only: bool = False,
+    clean: bool = False,
+    create: bool = False,
+    extension: MaybeCollectionStr | None = None,
+    extension_exc: MaybeCollectionStr | None = None,
+    schema: MaybeCollectionStr | None = None,
+    schema_exc: MaybeCollectionStr | None = None,
+    table: MaybeCollection[TableOrORMInstOrClass | str] | None = None,
+    table_exc: MaybeCollection[TableOrORMInstOrClass | str] | None = None,
+    inserts: bool = False,
+    on_conflict_do_nothing: bool = False,
+    role: str | None = None,
+    dry_run: bool = False,
+    logger: LoggerLike | None = None,
+) -> bool:
+    """Run `pg_dump`."""
+    path = _path_pg_dump(path, format_=format_)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    cmd = _build_pg_dump(
+        url,
+        path,
+        docker_container=docker_container,
+        format_=format_,
+        jobs=jobs,
+        data_only=data_only,
+        clean=clean,
+        create=create,
+        extension=extension,
+        extension_exc=extension_exc,
+        schema=schema,
+        schema_exc=schema_exc,
+        table=table,
+        table_exc=table_exc,
+        inserts=inserts,
+        on_conflict_do_nothing=on_conflict_do_nothing,
+        role=role,
+    )
+    if dry_run:
+        if logger is not None:
+            to_logger(logger).info("Would run:\n\t%r", str(cmd))
+        return True
+    with temp_environ(PGPASSWORD=url.password), Timer() as timer:  # pragma: no cover
+        try:
+            output = await stream_command(cmd)
+        except KeyboardInterrupt:
+            if logger is not None:
+                to_logger(logger).info(
+                    "Cancelled backup to %r after %s", str(path), timer
+                )
+            rmtree(path, ignore_errors=True)
+            return False
+        if output.return_code != 0:
+            if logger is not None:
+                to_logger(logger).exception(
+                    "Backup to %r failed after %s\nstderr:\n%s",
+                    str(path),
+                    timer,
+                    output.stderr,
+                )
+            rmtree(path, ignore_errors=True)
+            return False
+    if logger is not None:  # pragma: no cover
+        to_logger(logger).info("Backup to %r finished after %s", str(path), timer)
+    return True  # pragma: no cover
+
+
+def _build_pg_dump(
+    url: URL,
+    path: PathLike,
+    /,
+    *,
+    docker_container: str | None = None,
+    format_: _PGDumpFormat = "plain",
+    jobs: int | None = None,
+    data_only: bool = False,
+    clean: bool = False,
+    create: bool = False,
+    extension: MaybeCollectionStr | None = None,
+    extension_exc: MaybeCollectionStr | None = None,
+    schema: MaybeCollectionStr | None = None,
+    schema_exc: MaybeCollectionStr | None = None,
+    table: MaybeCollection[TableOrORMInstOrClass | str] | None = None,
+    table_exc: MaybeCollection[TableOrORMInstOrClass | str] | None = None,
+    inserts: bool = False,
+    on_conflict_do_nothing: bool = False,
+    role: str | None = None,
+) -> str:
+    extracted = extract_url(url)
+    path = _path_pg_dump(path, format_=format_)
+    parts: list[str] = ["pg_dump"]
+    if docker_container is not None:
+        parts = docker_exec_cmd(docker_container, *parts, PGPASSWORD=extracted.password)
+    parts.extend([
+        # general options
+        f"--file={str(path)!r}",
+        f"--format={format_}",
+        "--verbose",
+        # output options
+        *_resolve_data_only_and_clean(data_only=data_only, clean=clean),
+        "--large-objects",
+        "--no-owner",
+        "--no-privileges",
+        # connection options
+        f"--dbname={extracted.database}",
+        f"--host={extracted.host}",
+        f"--port={extracted.port}",
+        f"--username={extracted.username}",
+        "--no-password",
+    ])
+    if (format_ == "directory") and (jobs is not None):
+        parts.append(f"--jobs={jobs}")
+    if create:
+        parts.append("--create")
+    if extension is not None:
+        parts.extend([f"--extension={e}" for e in always_iterable(extension)])
+    if extension_exc is not None:
+        parts.extend([
+            f"--exclude-extension={e}" for e in always_iterable(extension_exc)
+        ])
+    if schema is not None:
+        parts.extend([f"--schema={s}" for s in always_iterable(schema)])
+    if schema_exc is not None:
+        parts.extend([f"--exclude-schema={s}" for s in always_iterable(schema_exc)])
+    if table is not None:
+        parts.extend([f"--table={_get_table_name(t)}" for t in always_iterable(table)])
+    if table_exc is not None:
+        parts.extend([
+            f"--exclude-table={_get_table_name(t)}" for t in always_iterable(table_exc)
+        ])
+    if inserts:
+        parts.append("--inserts")
+    if on_conflict_do_nothing:
+        parts.append("--on-conflict-do-nothing")
+    if role is not None:
+        parts.append(f"--role={role}")
+    return " ".join(parts)
+
+
+def _path_pg_dump(path: PathLike, /, *, format_: _PGDumpFormat = "plain") -> Path:
+    match format_:
+        case "plain":
+            suffix = ".sql"
+        case "custom":
+            suffix = ".pgdump"
+        case "directory":
+            suffix = None
+        case "tar":
+            suffix = ".tar"
+        case never:
+            assert_never(never)
+    path = Path(path)
+    if suffix is not None:
+        path = ensure_suffix(path, suffix)
+    return path
+
+
+##
+
+
+async def restore(
+    url: URL,
+    path: PathLike,
+    /,
+    *,
+    psql: bool = False,
+    data_only: bool = False,
+    clean: bool = False,
+    create: bool = False,
+    jobs: int | None = None,
+    schema: MaybeCollectionStr | None = None,
+    schema_exc: MaybeCollectionStr | None = None,
+    table: MaybeCollection[TableOrORMInstOrClass | str] | None = None,
+    role: str | None = None,
+    docker_container: str | None = None,
+    dry_run: bool = False,
+    logger: LoggerLike | None = None,
+) -> bool:
+    """Run `pg_restore`/`psql`."""
+    cmd = _build_pg_restore_or_psql(
+        url,
+        path,
+        psql=psql,
+        data_only=data_only,
+        clean=clean,
+        create=create,
+        jobs=jobs,
+        schema=schema,
+        schema_exc=schema_exc,
+        table=table,
+        role=role,
+        docker_container=docker_container,
+    )
+    if dry_run:
+        if logger is not None:
+            to_logger(logger).info("Would run:\n\t%r", str(cmd))
+        return True
+    with temp_environ(PGPASSWORD=url.password), Timer() as timer:  # pragma: no cover
+        try:
+            output = await stream_command(cmd)
+        except KeyboardInterrupt:
+            if logger is not None:
+                to_logger(logger).info(
+                    "Cancelled restore from %r after %s", str(path), timer
+                )
+            return False
+        if output.return_code != 0:
+            if logger is not None:
+                to_logger(logger).exception(
+                    "Restore from %r failed after %s\nstderr:\n%s",
+                    str(path),
+                    timer,
+                    output.stderr,
+                )
+            return False
+    if logger is not None:  # pragma: no cover
+        to_logger(logger).info("Restore from %r finished after %s", str(path), timer)
+    return True  # pragma: no cover
+
+
+##
+
+
+def _build_pg_restore_or_psql(
+    url: URL,
+    path: PathLike,
+    /,
+    *,
+    psql: bool = False,
+    data_only: bool = False,
+    clean: bool = False,
+    create: bool = False,
+    jobs: int | None = None,
+    schema: MaybeCollectionStr | None = None,
+    schema_exc: MaybeCollectionStr | None = None,
+    table: MaybeCollection[TableOrORMInstOrClass | str] | None = None,
+    role: str | None = None,
+    docker_container: str | None = None,
+) -> str:
+    path = Path(path)
+    if (path.suffix == ".sql") or psql:
+        return _build_psql(url, path, docker_container=docker_container)
+    return _build_pg_restore(
+        url,
+        path,
+        data_only=data_only,
+        clean=clean,
+        create=create,
+        jobs=jobs,
+        schemas=schema,
+        schemas_exc=schema_exc,
+        tables=table,
+        role=role,
+        docker_container=docker_container,
+    )
+
+
+def _build_pg_restore(
+    url: URL,
+    path: PathLike,
+    /,
+    *,
+    data_only: bool = False,
+    clean: bool = False,
+    create: bool = False,
+    jobs: int | None = None,
+    schemas: MaybeCollectionStr | None = None,
+    schemas_exc: MaybeCollectionStr | None = None,
+    tables: MaybeCollection[TableOrORMInstOrClass | str] | None = None,
+    role: str | None = None,
+    docker_container: str | None = None,
+) -> str:
+    """Run `pg_restore`."""
+    extracted = extract_url(url)
+    parts: list[str] = ["pg_restore"]
+    if docker_container is not None:
+        parts = docker_exec_cmd(docker_container, *parts, PGPASSWORD=extracted.password)
+    parts.extend([
+        # general options
+        "--verbose",
+        # restore options
+        *_resolve_data_only_and_clean(data_only=data_only, clean=clean),
+        "--exit-on-error",
+        "--no-owner",
+        "--no-privileges",
+        # connection options
+        f"--host={extracted.host}",
+        f"--port={extracted.port}",
+        f"--username={extracted.username}",
+        f"--dbname={extracted.database}",
+        "--no-password",
+    ])
+    if create:
+        parts.append("--create")
+    if jobs is not None:
+        parts.append(f"--jobs={jobs}")
+    if schemas is not None:
+        parts.extend([f"--schema={s}" for s in always_iterable(schemas)])
+    if schemas_exc is not None:
+        parts.extend([f"--exclude-schema={s}" for s in always_iterable(schemas_exc)])
+    if tables is not None:
+        parts.extend([f"--table={_get_table_name(t)}" for t in always_iterable(tables)])
+    if role is not None:
+        parts.append(f"--role={role}")
+    parts.append(str(path))
+    return " ".join(parts)
+
+
+def _build_psql(
+    url: URL, path: PathLike, /, *, docker_container: str | None = None
+) -> str:
+    """Run `psql`."""
+    extracted = extract_url(url)
+    parts: list[str] = ["psql"]
+    if docker_container is not None:
+        parts = docker_exec_cmd(docker_container, *parts, PGPASSWORD=extracted.password)
+    parts.extend([
+        # general options
+        f"--dbname={extracted.database}",
+        f"--file={str(path)!r}",
+        # connection options
+        f"--host={extracted.host}",
+        f"--port={extracted.port}",
+        f"--username={extracted.username}",
+        "--no-password",
+    ])
+    return " ".join(parts)
+
+
+##
+
+
+def _get_table_name(obj: TableOrORMInstOrClass | str, /) -> str:
+    match obj:
+        case Table() | DeclarativeBase() | type() as table_or_orm:
+            return get_table_name(table_or_orm)
+        case str() as name:
+            return name
+        case never:
+            assert_never(never)
+
+
+def _resolve_data_only_and_clean(
+    *, data_only: bool = False, clean: bool = False
+) -> list[str]:
+    match data_only, clean:
+        case False, False:
+            return []
+        case True, False:
+            return ["--data-only"]
+        case False, True:
+            return ["--clean", "--if-exists"]
+        case True, True:
+            raise _ResolveDataOnlyAndCleanError
+        case never:
+            assert_never(never)
+
+
+@dataclass(kw_only=True, slots=True)
+class _ResolveDataOnlyAndCleanError(Exception):
+    @override
+    def __str__(self) -> str:
+        return "Cannot use '--data-only' and '--clean' together"
+
+
+__all__ = ["pg_dump", "restore"]

utilities/pottery.py

@@ -1,7 +1,8 @@
 from __future__ import annotations
 
-from contextlib import asynccontextmanager, suppress
+from contextlib import suppress
 from dataclasses import dataclass
+from sys import maxsize
 from typing import TYPE_CHECKING, override
 
 from pottery import AIORedlock
@@ -9,30 +10,50 @@ from pottery.exceptions import ReleaseUnlockedLock
 from redis.asyncio import Redis
 
 from utilities.asyncio import sleep_td, timeout_td
+from utilities.contextlib import enhanced_async_context_manager
 from utilities.iterables import always_iterable
-from utilities.whenever import MILLISECOND, SECOND
+from utilities.whenever import MILLISECOND, SECOND, to_nanoseconds
 
 if TYPE_CHECKING:
     from collections.abc import AsyncIterator, Iterable
 
-    from whenever import TimeDelta
+    from whenever import Delta
 
     from utilities.types import MaybeIterable
 
+_NUM: int = 1
+_TIMEOUT_RELEASE: Delta = 10 * SECOND
+_SLEEP: Delta = MILLISECOND
 
-@asynccontextmanager
+
+##
+
+
+async def extend_lock(
+    *, lock: AIORedlock | None = None, raise_on_redis_errors: bool | None = None
+) -> None:
+    """Extend a lock."""
+    if lock is not None:
+        await lock.extend(raise_on_redis_errors=raise_on_redis_errors)
+
+
+##
+
+
+@enhanced_async_context_manager
 async def yield_access(
     redis: MaybeIterable[Redis],
     key: str,
     /,
     *,
-    num: int = 1,
-    timeout_acquire: TimeDelta | None = None,
-    timeout_release: TimeDelta = 10 * SECOND,
-    sleep: TimeDelta = MILLISECOND,
-    throttle: TimeDelta | None = None,
-) -> AsyncIterator[None]:
-    """Acquire access to a locked resource, amongst 1 of multiple connections."""
+    num: int = _NUM,
+    timeout_release: Delta = _TIMEOUT_RELEASE,
+    num_extensions: int | None = None,
+    timeout_acquire: Delta | None = None,
+    sleep: Delta = _SLEEP,
+    throttle: Delta | None = None,
+) -> AsyncIterator[AIORedlock]:
+    """Acquire access to a locked resource."""
     if num <= 0:
         raise _YieldAccessNumLocksError(key=key, num=num)
     masters = (  # skipif-ci-and-not-linux
@@ -42,7 +63,8 @@ async def yield_access(
         AIORedlock(
             key=f"{key}_{i}_of_{num}",
             masters=masters,
-            auto_release_time=timeout_release.in_seconds(),
+            auto_release_time=to_nanoseconds(timeout_release) / 1e9,
+            num_extensions=maxsize if num_extensions is None else num_extensions,
         )
         for i in range(1, num + 1)
     ]
@@ -51,7 +73,7 @@ async def yield_access(
         lock = await _get_first_available_lock(
             key, locks, num=num, timeout=timeout_acquire, sleep=sleep
         )
-        yield
+        yield lock
     finally:  # skipif-ci-and-not-linux
         await sleep_td(throttle)
         if lock is not None:
@@ -64,9 +86,9 @@ async def _get_first_available_lock(
     locks: Iterable[AIORedlock],
     /,
     *,
-    num: int = 1,
-    timeout: TimeDelta | None = None,
-    sleep: TimeDelta | None = None,
+    num: int = _NUM,
+    timeout: Delta | None = None,
+    sleep: Delta | None = _SLEEP,
 ) -> AIORedlock:
     locks = list(locks)  # skipif-ci-and-not-linux
     error = _YieldAccessUnableToAcquireLockError(  # skipif-ci-and-not-linux
@@ -91,11 +113,12 @@ async def _get_first_available_lock_if_any(
 @dataclass(kw_only=True, slots=True)
 class YieldAccessError(Exception):
     key: str
-    num: int
 
 
 @dataclass(kw_only=True, slots=True)
 class _YieldAccessNumLocksError(YieldAccessError):
+    num: int
+
     @override
     def __str__(self) -> str:
         return f"Number of locks for {self.key!r} must be positive; got {self.num}"
@@ -103,11 +126,12 @@ class _YieldAccessNumLocksError(YieldAccessError):
 
 @dataclass(kw_only=True, slots=True)
 class _YieldAccessUnableToAcquireLockError(YieldAccessError):
-    timeout: TimeDelta | None
+    num: int
+    timeout: Delta | None
 
     @override
     def __str__(self) -> str:
         return f"Unable to acquire any 1 of {self.num} locks for {self.key!r} after {self.timeout}"  # skipif-ci-and-not-linux
 
 
-__all__ = ["YieldAccessError", "yield_access"]
+__all__ = ["YieldAccessError", "extend_lock", "yield_access"]

utilities/pqdm.py

@@ -9,7 +9,7 @@ from tqdm.auto import tqdm as tqdm_auto
 from utilities.functions import get_func_name
 from utilities.iterables import apply_to_varargs
 from utilities.os import get_cpu_use
-from utilities.sentinel import Sentinel, sentinel
+from utilities.sentinel import Sentinel, is_sentinel, sentinel
 
 if TYPE_CHECKING:
     from collections.abc import Callable, Iterable
@@ -90,7 +90,7 @@ def pqdm_starmap[T](
                 **_get_desc(desc, func),
                 **kwargs,
             )
-        case _ as never:
+        case never:
             assert_never(never)
     return list(result)
 
@@ -98,7 +98,7 @@ def pqdm_starmap[T](
 def _get_desc(
     desc: str | None | Sentinel, func: Callable[..., Any], /
 ) -> dict[str, str]:
-    desc_use = get_func_name(func) if isinstance(desc, Sentinel) else desc
+    desc_use = get_func_name(func) if is_sentinel(desc) else desc
     return {} if desc_use is None else {"desc": desc_use}
 
 

utilities/psutil.py

@@ -1,71 +1,19 @@
 from __future__ import annotations
 
 from dataclasses import dataclass, field
-from json import dumps
-from logging import getLogger
 from math import isclose, nan
-from pathlib import Path
-from typing import TYPE_CHECKING, Self, override
+from typing import TYPE_CHECKING, Self
 
 from psutil import swap_memory, virtual_memory
 
-from utilities.asyncio import Looper
 from utilities.contextlib import suppress_super_object_attribute_error
-from utilities.whenever import SECOND, get_now
+from utilities.whenever import get_now
 
 if TYPE_CHECKING:
-    from logging import Logger
+    from whenever import ZonedDateTime
 
-    from whenever import TimeDelta, ZonedDateTime
 
-    from utilities.types import PathLike
-
-
-@dataclass(kw_only=True)
-class MemoryMonitorService(Looper[None]):
-    """Service to monitor memory usage."""
-
-    # base
-    freq: TimeDelta = field(default=10 * SECOND, repr=False)
-    backoff: TimeDelta = field(default=10 * SECOND, repr=False)
-    # self
-    console: str | None = field(default=None, repr=False)
-    path: PathLike = "memory.txt"
-    _console: Logger | None = field(init=False, repr=False)
-    _path: Path = field(init=False, repr=False)
-
-    @override
-    def __post_init__(self) -> None:
-        super().__post_init__()
-        if self.console is not None:
-            self._console = getLogger(self.console)
-        self._path = Path(self.path)
-        self._path.parent.mkdir(parents=True, exist_ok=True)
-
-    @override
-    async def core(self) -> None:
-        await super().core()
-        memory = MemoryUsage.new()
-        mapping = {
-            "datetime": memory.datetime.format_common_iso(),
-            "virtual used (mb)": memory.virtual_used_mb,
-            "virtual total (mb)": memory.virtual_total_mb,
-            "virtual (%)": memory.virtual_pct,
-            "swap used (mb)": memory.swap_used_mb,
-            "swap total (mb)": memory.swap_total_mb,
-            "swap (%)": memory.swap_pct,
-        }
-        ser = dumps(mapping)
-        with self._path.open(mode="a") as fh:
-            _ = fh.write(f"{ser}\n")
-        if self._console is not None:
-            self._console.info("%s", mapping)
-
-
-##
-
-
-@dataclass(kw_only=True)
+@dataclass(order=True, unsafe_hash=True, kw_only=True)
 class MemoryUsage:
     """A memory usage."""
 
@@ -113,4 +61,4 @@ class MemoryUsage:
         return round(bytes_ / (1024**2))
 
 
-__all__ = ["MemoryMonitorService", "MemoryUsage"]
+__all__ = ["MemoryUsage"]

utilities/pwd.py

@@ -0,0 +1,28 @@
+from __future__ import annotations
+
+from typing import assert_never
+
+from utilities.os import EFFECTIVE_USER_ID
+from utilities.platform import SYSTEM
+
+
+def get_uid_name(uid: int, /) -> str | None:
+    """Get the name of a user ID."""
+    match SYSTEM:
+        case "windows":  # skipif-not-windows
+            return None
+        case "mac" | "linux":  # skipif-windows
+            from pwd import getpwuid
+
+            return getpwuid(uid).pw_name
+        case never:
+            assert_never(never)
+
+
+ROOT_USER_NAME = get_uid_name(0)
+EFFECTIVE_USER_NAME = (
+    None if EFFECTIVE_USER_ID is None else get_uid_name(EFFECTIVE_USER_ID)
+)
+
+
+__all__ = ["EFFECTIVE_USER_NAME", "ROOT_USER_NAME", "get_uid_name"]