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

utilities/postgres.py

@@ -9,6 +9,7 @@ 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
@@ -37,6 +38,7 @@ async def pg_dump(
     path: PathLike,
     /,
     *,
+    docker_container: str | None = None,
     format_: _PGDumpFormat = "plain",
     jobs: int | None = None,
     data_only: bool = False,
@@ -51,7 +53,6 @@ async def pg_dump(
     inserts: bool = False,
     on_conflict_do_nothing: bool = False,
     role: str | None = None,
-    docker: str | None = None,
     dry_run: bool = False,
     logger: LoggerLike | None = None,
 ) -> bool:
@@ -61,6 +62,7 @@ async def pg_dump(
     cmd = _build_pg_dump(
         url,
         path,
+        docker_container=docker_container,
         format_=format_,
         jobs=jobs,
         data_only=data_only,
@@ -75,7 +77,6 @@ async def pg_dump(
         inserts=inserts,
         on_conflict_do_nothing=on_conflict_do_nothing,
         role=role,
-        docker=docker,
     )
     if dry_run:
         if logger is not None:
@@ -111,6 +112,7 @@ def _build_pg_dump(
     path: PathLike,
     /,
     *,
+    docker_container: str | None = None,
     format_: _PGDumpFormat = "plain",
     jobs: int | None = None,
     data_only: bool = False,
@@ -125,12 +127,13 @@ def _build_pg_dump(
     inserts: bool = False,
     on_conflict_do_nothing: bool = False,
     role: str | None = None,
-    docker: str | None = None,
 ) -> str:
     extracted = extract_url(url)
     path = _path_pg_dump(path, format_=format_)
-    parts: list[str] = [
-        "pg_dump",
+    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_}",
@@ -146,7 +149,7 @@ def _build_pg_dump(
         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:
@@ -173,8 +176,6 @@ def _build_pg_dump(
         parts.append("--on-conflict-do-nothing")
     if role is not None:
         parts.append(f"--role={role}")
-    if docker is not None:
-        parts = _wrap_docker(parts, docker)
     return " ".join(parts)
 
 
@@ -213,7 +214,7 @@ async def restore(
     schema_exc: MaybeCollectionStr | None = None,
     table: MaybeCollection[TableOrORMInstOrClass | str] | None = None,
     role: str | None = None,
-    docker: str | None = None,
+    docker_container: str | None = None,
     dry_run: bool = False,
     logger: LoggerLike | None = None,
 ) -> bool:
@@ -230,7 +231,7 @@ async def restore(
         schema_exc=schema_exc,
         table=table,
         role=role,
-        docker=docker,
+        docker_container=docker_container,
     )
     if dry_run:
         if logger is not None:
@@ -276,11 +277,11 @@ def _build_pg_restore_or_psql(
     schema_exc: MaybeCollectionStr | None = None,
     table: MaybeCollection[TableOrORMInstOrClass | str] | None = None,
     role: str | None = None,
-    docker: str | None = None,
+    docker_container: str | None = None,
 ) -> str:
     path = Path(path)
     if (path.suffix == ".sql") or psql:
-        return _build_psql(url, path, docker=docker)
+        return _build_psql(url, path, docker_container=docker_container)
     return _build_pg_restore(
         url,
         path,
@@ -292,7 +293,7 @@ def _build_pg_restore_or_psql(
         schemas_exc=schema_exc,
         tables=table,
         role=role,
-        docker=docker,
+        docker_container=docker_container,
     )
 
 
@@ -309,12 +310,14 @@ def _build_pg_restore(
     schemas_exc: MaybeCollectionStr | None = None,
     tables: MaybeCollection[TableOrORMInstOrClass | str] | None = None,
     role: str | None = None,
-    docker: str | None = None,
+    docker_container: str | None = None,
 ) -> str:
     """Run `pg_restore`."""
     extracted = extract_url(url)
-    parts: list[str] = [
-        "pg_restore",
+    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
@@ -328,7 +331,7 @@ def _build_pg_restore(
         f"--username={extracted.username}",
         f"--dbname={extracted.database}",
         "--no-password",
-    ]
+    ])
     if create:
         parts.append("--create")
     if jobs is not None:
@@ -341,17 +344,19 @@ def _build_pg_restore(
         parts.extend([f"--table={_get_table_name(t)}" for t in always_iterable(tables)])
     if role is not None:
         parts.append(f"--role={role}")
-    if docker is not None:
-        parts = _wrap_docker(parts, docker)
     parts.append(str(path))
     return " ".join(parts)
 
 
-def _build_psql(url: URL, path: PathLike, /, *, docker: str | None = None) -> str:
+def _build_psql(
+    url: URL, path: PathLike, /, *, docker_container: str | None = None
+) -> str:
     """Run `psql`."""
     extracted = extract_url(url)
-    parts: list[str] = [
-        "psql",
+    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}",
@@ -360,9 +365,7 @@ def _build_psql(url: URL, path: PathLike, /, *, docker: str | None = None) -> st
         f"--port={extracted.port}",
         f"--username={extracted.username}",
         "--no-password",
-    ]
-    if docker is not None:
-        parts = _wrap_docker(parts, docker)
+    ])
     return " ".join(parts)
 
 
@@ -402,8 +405,4 @@ class _ResolveDataOnlyAndCleanError(Exception):
         return "Cannot use '--data-only' and '--clean' together"
 
 
-def _wrap_docker(parts: list[str], container: str, /) -> list[str]:
-    return ["docker", "exec", "-it", container, *parts]
-
-
 __all__ = ["pg_dump", "restore"]

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"]

utilities/pydantic.py

@@ -0,0 +1,11 @@
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Annotated
+
+from pydantic import BeforeValidator
+
+ExpandedPath = Annotated[Path, BeforeValidator(lambda p: Path(p).expanduser())]
+
+
+__all__ = ["ExpandedPath"]

utilities/pydantic_settings.py

@@ -2,10 +2,12 @@ from __future__ import annotations
 
 from functools import reduce
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, ClassVar, assert_never, override
+from typing import TYPE_CHECKING, Any, ClassVar, assert_never, cast, override
 
+from pydantic import Field, create_model
 from pydantic_settings import (
     BaseSettings,
+    CliSettingsSource,
     JsonConfigSettingsSource,
     PydanticBaseSettingsSource,
     SettingsConfigDict,
@@ -14,6 +16,7 @@ from pydantic_settings import (
 )
 from pydantic_settings.sources import DEFAULT_PATH
 
+from utilities.errors import ImpossibleCaseError
 from utilities.iterables import always_iterable
 
 if TYPE_CHECKING:
@@ -29,7 +32,7 @@ type PathLikeOrWithSection = PathLike | PathLikeWithSection
 
 
 class CustomBaseSettings(BaseSettings):
-    """Base settings for loading JSON files."""
+    """Base settings for loading JSON/TOML/YAML files."""
 
     # paths
     json_files: ClassVar[Sequence[PathLikeOrWithSection]] = []
@@ -38,7 +41,7 @@ class CustomBaseSettings(BaseSettings):
 
     # config
     model_config: ClassVar[SettingsConfigDict] = SettingsConfigDict(
-        env_nested_delimiter="__"
+        frozen=True, env_nested_delimiter="__"
     )
 
     @classmethod
@@ -76,11 +79,6 @@ class CustomBaseSettings(BaseSettings):
             )
 
 
-def load_settings[T: BaseSettings](cls: type[T], /) -> T:
-    """Load a set of settings."""
-    return cls()
-
-
 class JsonConfigSectionSettingsSource(JsonConfigSettingsSource):
     @override
     def __init__(
@@ -158,8 +156,83 @@ def _get_section(
     return reduce(lambda acc, el: acc.get(el, {}), always_iterable(section), mapping)
 
 
+##
+
+
+class HashableBaseSettings(BaseSettings):
+    """Base settings for loading JSON files."""
+
+    # config
+    model_config: ClassVar[SettingsConfigDict] = SettingsConfigDict(frozen=True)
+
+
+##
+
+
+def load_settings[T: BaseSettings](cls: type[T], /, *, cli: bool = False) -> T:
+    """Load a set of settings."""
+    _ = cls.model_rebuild()
+    if cli:
+        cls_with_defaults = _load_settings_create_model(cls)
+
+        @classmethod
+        def settings_customise_sources(
+            cls: type[BaseSettings],
+            settings_cls: type[BaseSettings],
+            init_settings: PydanticBaseSettingsSource,
+            env_settings: PydanticBaseSettingsSource,
+            dotenv_settings: PydanticBaseSettingsSource,
+            file_secret_settings: PydanticBaseSettingsSource,
+        ) -> tuple[PydanticBaseSettingsSource, ...]:
+            parent = cast(
+                "Any", super(cls_with_defaults, cls)
+            ).settings_customise_sources(
+                settings_cls=settings_cls,
+                init_settings=init_settings,
+                env_settings=env_settings,
+                dotenv_settings=dotenv_settings,
+                file_secret_settings=file_secret_settings,
+            )
+            return (
+                CliSettingsSource(
+                    settings_cls, cli_parse_args=True, case_sensitive=False
+                ),
+                *parent,
+            )
+
+        cls_use = type(
+            cls.__name__,
+            (cls_with_defaults,),
+            {"settings_customise_sources": settings_customise_sources},
+        )
+        cls_use = cast("type[T]", cls_use)
+    else:
+        cls_use = cls
+    return cls_use()
+
+
+def _load_settings_create_model[T: BaseSettings](
+    cls: type[T], /, *, values: T | None = None
+) -> type[T]:
+    values_use = cls() if values is None else values
+    kwargs: dict[str, Any] = {}
+    for name, field in cls.model_fields.items():
+        if (ann := field.annotation) is None:
+            raise ImpossibleCaseError(case=[f"{ann=}"])  # pragma: no cover
+        value = getattr(values_use, name)
+        if (
+            isinstance(cast("Any", ann), type)  # 'ann' is possible not a type
+            and issubclass(ann, BaseSettings)
+        ):
+            kwargs[name] = _load_settings_create_model(ann, values=value)
+        else:
+            kwargs[name] = (field.annotation, Field(default=value))
+    return create_model(cls.__name__, __base__=cls, **kwargs)
+
+
 __all__ = [
     "CustomBaseSettings",
+    "HashableBaseSettings",
     "JsonConfigSectionSettingsSource",
     "TomlConfigSectionSettingsSource",
     "YamlConfigSectionSettingsSource",

utilities/pydantic_settings_sops.py

@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+from logging import Filter, LogRecord, getLogger
+from re import search
 from typing import TYPE_CHECKING, Any, ClassVar, override
 
 from pydantic_settings.sources import DEFAULT_PATH
@@ -21,6 +23,17 @@ if TYPE_CHECKING:
     from utilities.types import MaybeSequenceStr
 
 
+class _SuppressDefaultConfigMessage(Filter):
+    @override
+    def filter(self, record: LogRecord) -> bool:
+        return not search(
+            r"^default config file does not exists '.*'$", record.getMessage()
+        )
+
+
+getLogger("sopsy.utils").addFilter(_SuppressDefaultConfigMessage())
+
+
 class SopsBaseSettings(CustomBaseSettings):
     """Base settings for loading secrets using `sops/age`."""
 

utilities/pytest.py

@@ -1,10 +1,13 @@
 from __future__ import annotations
 
+from collections.abc import Callable
 from dataclasses import dataclass
 from functools import partial, wraps
 from inspect import iscoroutinefunction
 from os import environ
 from pathlib import Path
+from re import sub
+from types import FunctionType
 from typing import TYPE_CHECKING, Any, assert_never, cast, override
 
 from whenever import ZonedDateTime
@@ -20,44 +23,57 @@ from utilities.pathlib import (
     get_tail,
     module_path,
 )
-from utilities.platform import (
-    IS_LINUX,
-    IS_MAC,
-    IS_NOT_LINUX,
-    IS_NOT_MAC,
-    IS_NOT_WINDOWS,
-    IS_WINDOWS,
-)
+from utilities.platform import IS_LINUX, IS_MAC, IS_NOT_LINUX, IS_NOT_MAC
 from utilities.random import bernoulli
 from utilities.text import to_bool
 from utilities.types import MaybeCallableBoolLike, MaybeCoro, Seed
 from utilities.whenever import SECOND, get_now_local
 
 if TYPE_CHECKING:
-    from collections.abc import Callable, Iterable
-
-    from utilities.types import Coro, Delta, PathLike
+    from collections.abc import Iterable
 
-try:  # WARNING: this package cannot use unguarded `pytest` imports
     from _pytest.config import Config
     from _pytest.config.argparsing import Parser
     from _pytest.python import Function
-    from pytest import mark, skip
+
+    from utilities.types import Coro, Delta, PathLike
+
+
+IS_CI = "CI" in environ
+IS_CI_AND_NOT_LINUX = IS_CI and IS_NOT_LINUX
+
+
+try:  # WARNING: this package cannot use unguarded `pytest` imports
+    from pytest import mark
 except ModuleNotFoundError:  # pragma: no cover
-    from typing import Any as Config
-    from typing import Any as Function
-    from typing import Any as Parser
 
-    mark = skip = skipif_windows = skipif_mac = skipif_linux = skipif_not_windows = (
-        skipif_not_mac
-    ) = skipif_not_linux = None
+    def skipif_ci[F: Callable](func: F) -> F:
+        return func
+
+    def skipif_mac[F: Callable](func: F) -> F:
+        return func
+
+    def skipif_linux[F: Callable](func: F) -> F:
+        return func
+
+    def skipif_not_mac[F: Callable](func: F) -> F:
+        return func
+
+    def skipif_not_linux[F: Callable](func: F) -> F:
+        return func
+
+    def skipif_ci_and_not_linux[F: Callable](func: F) -> F:
+        return func
+
 else:
-    skipif_windows = mark.skipif(IS_WINDOWS, reason="Skipped for Windows")
+    skipif_ci = mark.skipif(IS_CI, reason="Skipped for CI")
     skipif_mac = mark.skipif(IS_MAC, reason="Skipped for Mac")
     skipif_linux = mark.skipif(IS_LINUX, reason="Skipped for Linux")
-    skipif_not_windows = mark.skipif(IS_NOT_WINDOWS, reason="Skipped for non-Windows")
     skipif_not_mac = mark.skipif(IS_NOT_MAC, reason="Skipped for non-Mac")
     skipif_not_linux = mark.skipif(IS_NOT_LINUX, reason="Skipped for non-Linux")
+    skipif_ci_and_not_linux = mark.skipif(
+        IS_CI_AND_NOT_LINUX, reason="Skipped for CI/non-Linux"
+    )
 
 
 def add_pytest_addoption(parser: Parser, options: list[str], /) -> None:
@@ -90,6 +106,8 @@ def add_pytest_collection_modifyitems(
         def pytest_collection_modifyitems(config, items):
             add_pytest_collection_modifyitems(config, items, ["slow"])
     """
+    from pytest import mark
+
     options = list(options)
     missing = {opt for opt in options if not config.getoption(f"--{opt}")}
     for item in items:
@@ -119,6 +137,15 @@ def add_pytest_configure(config: Config, options: Iterable[tuple[str, str]], /)
 ##
 
 
+def make_ids(obj: Any, /) -> str:
+    if isinstance(obj, FunctionType):
+        return sub(r"\s+at +0x[0-9a-fA-F]+", "", repr(obj))
+    return repr(obj)
+
+
+##
+
+
 def node_id_path(
     node_id: str, /, *, root: PathLike | None = None, suffix: str | None = None
 ) -> Path:
@@ -217,12 +244,12 @@ def _skipif_frac(
     frac: float = 0.5,
     seed: Seed | None = None,
 ) -> None:
-    if skip is None:
-        return  # pragma: no cover
+    from pytest import skip
+
     if ((predicate is None) or to_bool(predicate)) and bernoulli(
         true=1 - frac, seed=seed
     ):
-        _ = skip(reason=f"{_get_name()} skipped (run {frac:.0%})")
+        skip(reason=f"{_get_name()} skipped (run {frac:.0%})")
 
 
 ##
@@ -291,15 +318,15 @@ def _throttle_inner[F: Callable[..., MaybeCoro[None]]](
 
 
 def _skipif_recent(*, root: PathLike | None = None, delta: Delta = SECOND) -> None:
-    if skip is None:
-        return  # pragma: no cover
+    from pytest import skip
+
     path = _get_path(root)
     try:
         contents = path.read_text()
     except FileNotFoundError:
         return
     try:
-        last = ZonedDateTime.parse_common_iso(contents)
+        last = ZonedDateTime.parse_iso(contents)
     except ValueError:
         return
     now = get_now_local()
@@ -328,21 +355,24 @@ def _get_name() -> str:
 def _write(root: PathLike | None = None, /) -> None:
     path = _get_path(root)
     with writer(path, overwrite=True) as temp:
-        _ = temp.write_text(get_now_local().format_common_iso())
+        _ = temp.write_text(get_now_local().format_iso())
 
 
 __all__ = [
+    "IS_CI",
+    "IS_CI_AND_NOT_LINUX",
     "NodeIdToPathError",
     "add_pytest_addoption",
     "add_pytest_collection_modifyitems",
     "add_pytest_configure",
+    "make_ids",
     "node_id_path",
     "run_frac",
+    "skipif_ci",
+    "skipif_ci_and_not_linux",
     "skipif_linux",
     "skipif_mac",
     "skipif_not_linux",
     "skipif_not_mac",
-    "skipif_not_windows",
-    "skipif_windows",
     "throttle",
 ]

utilities/pytest_regressions.py

@@ -1,15 +1,17 @@
 from __future__ import annotations
 
 from contextlib import suppress
+from dataclasses import dataclass
 from json import loads
 from pathlib import Path
 from shutil import copytree
-from typing import TYPE_CHECKING, Any, assert_never
+from typing import TYPE_CHECKING, Any, assert_never, override
 
 from pytest_regressions.file_regression import FileRegressionFixture
 
 from utilities.functions import ensure_str
 from utilities.operator import is_equal
+from utilities.reprlib import get_repr
 
 if TYPE_CHECKING:
     from polars import DataFrame, Series
@@ -70,10 +72,28 @@ class OrjsonRegressionFixture:
             check_fn=self._check_fn,
         )
 
-    def _check_fn(self, path1: Path, path2: Path, /) -> None:
-        left = loads(path1.read_text())
-        right = loads(path2.read_text())
-        assert is_equal(left, right), f"{left=}, {right=}"
+    def _check_fn(self, path_obtained: Path, path_existing: Path, /) -> None:
+        obtained = loads(path_obtained.read_text())
+        existing = loads(path_existing.read_text())
+        if not is_equal(obtained, existing):
+            raise OrjsonRegressionError(
+                path_obtained=path_obtained,
+                path_existing=path_existing,
+                obtained=obtained,
+                existing=existing,
+            )
+
+
+@dataclass(kw_only=True, slots=True)
+class OrjsonRegressionError(Exception):
+    path_obtained: Path
+    path_existing: Path
+    obtained: Any
+    existing: Any
+
+    @override
+    def __str__(self) -> str:
+        return f"Obtained object (at {str(self.path_obtained)!r}) and existing object (at {str(self.path_existing)!r}) differ; got {get_repr(self.obtained)} and {get_repr(self.existing)}"
 
 
 ##
@@ -97,7 +117,6 @@ class PolarsRegressionFixture:
             "describe": obj.describe(percentiles=[i / 10 for i in range(1, 10)]).rows(
                 named=True
             ),
-            "estimated_size": obj.estimated_size(),
             "is_empty": obj.is_empty(),
             "n_unique": obj.n_unique(),
         }
@@ -115,7 +134,7 @@ class PolarsRegressionFixture:
                             col(column).approx_n_unique()
                         ).item()
                 data["approx_n_unique"] = approx_n_unique
-                data["glimpse"] = df.glimpse(return_as_string=True)
+                data["glimpse"] = df.glimpse(return_type="string")
                 data["null_count"] = df.null_count().row(0, named=True)
             case never:
                 assert_never(never)

utilities/shutil.py

@@ -0,0 +1,25 @@
+from __future__ import annotations
+
+import shutil
+from dataclasses import dataclass
+from pathlib import Path
+from typing import override
+
+
+def which(cmd: str, /) -> Path:
+    path = shutil.which(cmd)
+    if path is None:
+        raise WhichError(cmd=cmd)
+    return Path(path)
+
+
+@dataclass(kw_only=True, slots=True)
+class WhichError(Exception):
+    cmd: str
+
+    @override
+    def __str__(self) -> str:
+        return f"{self.cmd!r} not found"
+
+
+__all__ = ["WhichError", "which"]

utilities/sqlalchemy.py

@@ -331,6 +331,20 @@ async def ensure_database_dropped(super_: URL, database: str, /) -> None:
         _ = await conn.execute(text(f"DROP DATABASE IF EXISTS {database}"))
 
 
+async def ensure_database_users_disconnected(super_: URL, database: str, /) -> None:
+    """Ensure a databases' users are disconnected."""
+    engine = create_async_engine(super_, isolation_level="AUTOCOMMIT")
+    match dialect := _get_dialect(engine):
+        case "postgresql":  # skipif-ci-and-not-linux
+            query = f"SELECT pg_terminate_backend(pid) FROM pg_stat_activity WHERE datname = {database!r} AND pid <> pg_backend_pid()"  # noqa: S608
+        case "mssql" | "mysql" | "oracle" | "sqlite":  # pragma: no cover
+            raise NotImplementedError(dialect)
+        case never:
+            assert_never(never)
+    async with engine.begin() as conn:
+        _ = await conn.execute(text(query))
+
+
 ##
 
 
@@ -1166,6 +1180,7 @@ __all__ = [
     "create_engine",
     "ensure_database_created",
     "ensure_database_dropped",
+    "ensure_database_users_disconnected",
     "ensure_tables_created",
     "ensure_tables_dropped",
     "enum_name",