dycw-utilities 0.175.17__py3-none-any.whl → 0.185.8__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- dycw_utilities-0.185.8.dist-info/METADATA +33 -0
- dycw_utilities-0.185.8.dist-info/RECORD +90 -0
- {dycw_utilities-0.175.17.dist-info → dycw_utilities-0.185.8.dist-info}/WHEEL +2 -2
- utilities/__init__.py +1 -1
- utilities/altair.py +8 -6
- utilities/asyncio.py +40 -56
- utilities/atools.py +9 -11
- utilities/cachetools.py +8 -6
- utilities/click.py +4 -3
- utilities/concurrent.py +1 -1
- utilities/constants.py +492 -0
- utilities/contextlib.py +23 -30
- utilities/contextvars.py +1 -23
- utilities/core.py +2581 -0
- utilities/dataclasses.py +16 -119
- utilities/docker.py +139 -45
- utilities/enum.py +1 -1
- utilities/errors.py +2 -16
- utilities/fastapi.py +5 -5
- utilities/fpdf2.py +2 -1
- utilities/functions.py +33 -264
- utilities/http.py +2 -3
- utilities/hypothesis.py +48 -25
- utilities/iterables.py +39 -575
- utilities/jinja2.py +3 -6
- utilities/jupyter.py +5 -3
- utilities/libcst.py +1 -1
- utilities/lightweight_charts.py +4 -6
- utilities/logging.py +17 -15
- utilities/math.py +1 -36
- utilities/more_itertools.py +4 -6
- utilities/numpy.py +2 -1
- utilities/operator.py +2 -2
- utilities/orjson.py +24 -25
- utilities/os.py +4 -185
- utilities/packaging.py +129 -0
- utilities/parse.py +33 -13
- utilities/pathlib.py +2 -136
- utilities/platform.py +8 -90
- utilities/polars.py +34 -31
- utilities/postgres.py +9 -4
- utilities/pottery.py +20 -18
- utilities/pqdm.py +3 -4
- utilities/psutil.py +2 -3
- utilities/pydantic.py +18 -4
- utilities/pydantic_settings.py +7 -9
- utilities/pydantic_settings_sops.py +3 -3
- utilities/pyinstrument.py +4 -4
- utilities/pytest.py +49 -108
- utilities/pytest_plugins/pytest_regressions.py +2 -2
- utilities/pytest_regressions.py +8 -6
- utilities/random.py +2 -8
- utilities/redis.py +98 -94
- utilities/reprlib.py +11 -118
- utilities/shellingham.py +66 -0
- utilities/slack_sdk.py +13 -12
- utilities/sqlalchemy.py +42 -30
- utilities/sqlalchemy_polars.py +16 -25
- utilities/subprocess.py +1166 -148
- utilities/tabulate.py +32 -0
- utilities/testbook.py +8 -8
- utilities/text.py +24 -115
- utilities/throttle.py +159 -0
- utilities/time.py +18 -0
- utilities/timer.py +29 -12
- utilities/traceback.py +15 -22
- utilities/types.py +38 -3
- utilities/typing.py +18 -12
- utilities/uuid.py +1 -1
- utilities/version.py +202 -45
- utilities/whenever.py +22 -150
- dycw_utilities-0.175.17.dist-info/METADATA +0 -34
- dycw_utilities-0.175.17.dist-info/RECORD +0 -103
- utilities/atomicwrites.py +0 -182
- utilities/cryptography.py +0 -41
- utilities/getpass.py +0 -8
- utilities/git.py +0 -19
- utilities/grp.py +0 -28
- utilities/gzip.py +0 -31
- utilities/json.py +0 -70
- utilities/permissions.py +0 -298
- utilities/pickle.py +0 -25
- utilities/pwd.py +0 -28
- utilities/re.py +0 -156
- utilities/sentinel.py +0 -73
- utilities/socket.py +0 -8
- utilities/string.py +0 -20
- utilities/tempfile.py +0 -136
- utilities/tzdata.py +0 -11
- utilities/tzlocal.py +0 -28
- utilities/warnings.py +0 -65
- utilities/zipfile.py +0 -25
- utilities/zoneinfo.py +0 -133
- {dycw_utilities-0.175.17.dist-info → dycw_utilities-0.185.8.dist-info}/entry_points.txt +0 -0
utilities/reprlib.py
CHANGED
|
@@ -1,117 +1,22 @@
|
|
|
1
1
|
from __future__ import annotations
|
|
2
2
|
|
|
3
|
-
import reprlib
|
|
4
3
|
from functools import partial
|
|
5
|
-
from typing import TYPE_CHECKING
|
|
4
|
+
from typing import TYPE_CHECKING
|
|
5
|
+
|
|
6
|
+
from utilities.constants import (
|
|
7
|
+
RICH_EXPAND_ALL,
|
|
8
|
+
RICH_INDENT_SIZE,
|
|
9
|
+
RICH_MAX_DEPTH,
|
|
10
|
+
RICH_MAX_LENGTH,
|
|
11
|
+
RICH_MAX_STRING,
|
|
12
|
+
RICH_MAX_WIDTH,
|
|
13
|
+
)
|
|
6
14
|
|
|
7
15
|
if TYPE_CHECKING:
|
|
8
16
|
from collections.abc import Iterator
|
|
9
17
|
|
|
10
18
|
from utilities.types import StrMapping
|
|
11
19
|
|
|
12
|
-
RICH_MAX_WIDTH: int = 80
|
|
13
|
-
RICH_INDENT_SIZE: int = 4
|
|
14
|
-
RICH_MAX_LENGTH: int | None = 20
|
|
15
|
-
RICH_MAX_STRING: int | None = None
|
|
16
|
-
RICH_MAX_DEPTH: int | None = None
|
|
17
|
-
RICH_EXPAND_ALL: bool = False
|
|
18
|
-
|
|
19
|
-
|
|
20
|
-
##
|
|
21
|
-
|
|
22
|
-
|
|
23
|
-
def get_call_args_mapping(*args: Any, **kwargs: Any) -> StrMapping:
|
|
24
|
-
"""Get the representation of a set of call arguments."""
|
|
25
|
-
return {f"args[{i}]": v for i, v in enumerate(args)} | {
|
|
26
|
-
f"kwargs[{k}]": v for k, v in kwargs.items()
|
|
27
|
-
}
|
|
28
|
-
|
|
29
|
-
|
|
30
|
-
##
|
|
31
|
-
|
|
32
|
-
|
|
33
|
-
def get_repr(
|
|
34
|
-
obj: Any,
|
|
35
|
-
/,
|
|
36
|
-
*,
|
|
37
|
-
max_width: int = RICH_MAX_WIDTH,
|
|
38
|
-
indent_size: int = RICH_INDENT_SIZE,
|
|
39
|
-
max_length: int | None = RICH_MAX_LENGTH,
|
|
40
|
-
max_string: int | None = RICH_MAX_STRING,
|
|
41
|
-
max_depth: int | None = RICH_MAX_DEPTH,
|
|
42
|
-
expand_all: bool = RICH_EXPAND_ALL,
|
|
43
|
-
) -> str:
|
|
44
|
-
"""Get the representation of an object."""
|
|
45
|
-
try:
|
|
46
|
-
from rich.pretty import pretty_repr
|
|
47
|
-
except ModuleNotFoundError: # pragma: no cover
|
|
48
|
-
return reprlib.repr(obj)
|
|
49
|
-
return pretty_repr(
|
|
50
|
-
obj,
|
|
51
|
-
max_width=max_width,
|
|
52
|
-
indent_size=indent_size,
|
|
53
|
-
max_length=max_length,
|
|
54
|
-
max_string=max_string,
|
|
55
|
-
max_depth=max_depth,
|
|
56
|
-
expand_all=expand_all,
|
|
57
|
-
)
|
|
58
|
-
|
|
59
|
-
|
|
60
|
-
##
|
|
61
|
-
|
|
62
|
-
|
|
63
|
-
def get_repr_and_class(
|
|
64
|
-
obj: Any,
|
|
65
|
-
/,
|
|
66
|
-
*,
|
|
67
|
-
max_width: int = RICH_MAX_WIDTH,
|
|
68
|
-
indent_size: int = RICH_INDENT_SIZE,
|
|
69
|
-
max_length: int | None = RICH_MAX_LENGTH,
|
|
70
|
-
max_string: int | None = RICH_MAX_STRING,
|
|
71
|
-
max_depth: int | None = RICH_MAX_DEPTH,
|
|
72
|
-
expand_all: bool = RICH_EXPAND_ALL,
|
|
73
|
-
) -> str:
|
|
74
|
-
"""Get the `reprlib`-representation & class of an object."""
|
|
75
|
-
repr_use = get_repr(
|
|
76
|
-
obj,
|
|
77
|
-
max_width=max_width,
|
|
78
|
-
indent_size=indent_size,
|
|
79
|
-
max_length=max_length,
|
|
80
|
-
max_string=max_string,
|
|
81
|
-
max_depth=max_depth,
|
|
82
|
-
expand_all=expand_all,
|
|
83
|
-
)
|
|
84
|
-
return f"Object {repr_use!r} of type {type(obj).__name__!r}"
|
|
85
|
-
|
|
86
|
-
|
|
87
|
-
##
|
|
88
|
-
|
|
89
|
-
|
|
90
|
-
def yield_call_args_repr(
|
|
91
|
-
*args: Any,
|
|
92
|
-
_max_width: int = RICH_MAX_WIDTH,
|
|
93
|
-
_indent_size: int = RICH_INDENT_SIZE,
|
|
94
|
-
_max_length: int | None = RICH_MAX_LENGTH,
|
|
95
|
-
_max_string: int | None = RICH_MAX_STRING,
|
|
96
|
-
_max_depth: int | None = RICH_MAX_DEPTH,
|
|
97
|
-
_expand_all: bool = RICH_EXPAND_ALL,
|
|
98
|
-
**kwargs: Any,
|
|
99
|
-
) -> Iterator[str]:
|
|
100
|
-
"""Pretty print of a set of positional/keyword arguments."""
|
|
101
|
-
mapping = get_call_args_mapping(*args, **kwargs)
|
|
102
|
-
return yield_mapping_repr(
|
|
103
|
-
mapping,
|
|
104
|
-
_max_width=_max_width,
|
|
105
|
-
_indent_size=_indent_size,
|
|
106
|
-
_max_length=_max_length,
|
|
107
|
-
_max_string=_max_string,
|
|
108
|
-
_max_depth=_max_depth,
|
|
109
|
-
_expand_all=_expand_all,
|
|
110
|
-
)
|
|
111
|
-
|
|
112
|
-
|
|
113
|
-
##
|
|
114
|
-
|
|
115
20
|
|
|
116
21
|
def yield_mapping_repr(
|
|
117
22
|
mapping: StrMapping,
|
|
@@ -143,16 +48,4 @@ def yield_mapping_repr(
|
|
|
143
48
|
yield f"{k} = {repr_use(v)}"
|
|
144
49
|
|
|
145
50
|
|
|
146
|
-
__all__ = [
|
|
147
|
-
"RICH_EXPAND_ALL",
|
|
148
|
-
"RICH_INDENT_SIZE",
|
|
149
|
-
"RICH_MAX_DEPTH",
|
|
150
|
-
"RICH_MAX_LENGTH",
|
|
151
|
-
"RICH_MAX_STRING",
|
|
152
|
-
"RICH_MAX_WIDTH",
|
|
153
|
-
"get_call_args_mapping",
|
|
154
|
-
"get_repr",
|
|
155
|
-
"get_repr_and_class",
|
|
156
|
-
"yield_call_args_repr",
|
|
157
|
-
"yield_mapping_repr",
|
|
158
|
-
]
|
|
51
|
+
__all__ = ["yield_mapping_repr"]
|
utilities/shellingham.py
ADDED
|
@@ -0,0 +1,66 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
from dataclasses import dataclass
|
|
4
|
+
from os import environ, name
|
|
5
|
+
from pathlib import Path
|
|
6
|
+
from typing import Literal, override
|
|
7
|
+
|
|
8
|
+
from shellingham import ShellDetectionFailure, detect_shell
|
|
9
|
+
|
|
10
|
+
from utilities.core import OneEmptyError, one, repr_
|
|
11
|
+
from utilities.typing import get_args
|
|
12
|
+
|
|
13
|
+
type Shell = Literal["bash", "fish", "posix", "sh", "zsh"]
|
|
14
|
+
|
|
15
|
+
|
|
16
|
+
def get_shell() -> Shell:
|
|
17
|
+
"""Get the shell."""
|
|
18
|
+
try:
|
|
19
|
+
shell, _ = detect_shell()
|
|
20
|
+
except ShellDetectionFailure: # pragma: no cover
|
|
21
|
+
if name == "posix":
|
|
22
|
+
shell = environ["SHELL"]
|
|
23
|
+
elif name == "nt":
|
|
24
|
+
shell = environ["COMSPEC"]
|
|
25
|
+
else:
|
|
26
|
+
raise _GetShellOSError(name=name) from None
|
|
27
|
+
shells: tuple[Shell, ...] = get_args(Shell)
|
|
28
|
+
matches: list[Shell] = [s for s in shells if _get_shell_match(shell, s)]
|
|
29
|
+
try:
|
|
30
|
+
return one(matches)
|
|
31
|
+
except OneEmptyError: # pragma: no cover
|
|
32
|
+
raise _GetShellUnsupportedError(shell=shell) from None
|
|
33
|
+
|
|
34
|
+
|
|
35
|
+
def _get_shell_match(shell: str, candidate: Shell, /) -> bool:
|
|
36
|
+
*_, name = Path(shell).parts
|
|
37
|
+
return name == candidate
|
|
38
|
+
|
|
39
|
+
|
|
40
|
+
@dataclass(kw_only=True, slots=True)
|
|
41
|
+
class GetShellError(Exception):
|
|
42
|
+
name: str
|
|
43
|
+
|
|
44
|
+
|
|
45
|
+
@dataclass(kw_only=True, slots=True)
|
|
46
|
+
class _GetShellUnsupportedError(Exception):
|
|
47
|
+
shell: str
|
|
48
|
+
|
|
49
|
+
@override
|
|
50
|
+
def __str__(self) -> str:
|
|
51
|
+
return f"Invalid shell; got {repr_(self.shell)}" # pragma: no cover
|
|
52
|
+
|
|
53
|
+
|
|
54
|
+
@dataclass(kw_only=True, slots=True)
|
|
55
|
+
class _GetShellOSError(GetShellError):
|
|
56
|
+
name: str
|
|
57
|
+
|
|
58
|
+
@override
|
|
59
|
+
def __str__(self) -> str:
|
|
60
|
+
return f"Invalid OS; got {repr_(self.name)}" # pragma: no cover
|
|
61
|
+
|
|
62
|
+
|
|
63
|
+
SHELL = get_shell()
|
|
64
|
+
|
|
65
|
+
|
|
66
|
+
__all__ = ["SHELL", "GetShellError", "get_shell"]
|
utilities/slack_sdk.py
CHANGED
|
@@ -7,24 +7,25 @@ from typing import TYPE_CHECKING, override
|
|
|
7
7
|
from slack_sdk.webhook import WebhookClient
|
|
8
8
|
from slack_sdk.webhook.async_client import AsyncWebhookClient
|
|
9
9
|
|
|
10
|
-
|
|
10
|
+
import utilities.asyncio
|
|
11
|
+
from utilities.constants import MINUTE
|
|
12
|
+
from utilities.functions import in_seconds
|
|
11
13
|
from utilities.functools import cache
|
|
12
|
-
from utilities.
|
|
14
|
+
from utilities.math import safe_round
|
|
13
15
|
|
|
14
16
|
if TYPE_CHECKING:
|
|
15
17
|
from slack_sdk.webhook import WebhookResponse
|
|
16
|
-
from whenever import TimeDelta
|
|
17
18
|
|
|
18
|
-
from utilities.types import
|
|
19
|
+
from utilities.types import Duration, MaybeType
|
|
19
20
|
|
|
20
21
|
|
|
21
|
-
_TIMEOUT:
|
|
22
|
+
_TIMEOUT: Duration = MINUTE
|
|
22
23
|
|
|
23
24
|
|
|
24
25
|
##
|
|
25
26
|
|
|
26
27
|
|
|
27
|
-
def send_to_slack(url: str, text: str, /, *, timeout:
|
|
28
|
+
def send_to_slack(url: str, text: str, /, *, timeout: Duration = _TIMEOUT) -> None:
|
|
28
29
|
"""Send a message via Slack synchronously."""
|
|
29
30
|
client = _get_client(url, timeout=timeout)
|
|
30
31
|
response = client.send(text=text)
|
|
@@ -33,9 +34,9 @@ def send_to_slack(url: str, text: str, /, *, timeout: TimeDelta = _TIMEOUT) -> N
|
|
|
33
34
|
|
|
34
35
|
|
|
35
36
|
@cache
|
|
36
|
-
def _get_client(url: str, /, *, timeout:
|
|
37
|
+
def _get_client(url: str, /, *, timeout: Duration = _TIMEOUT) -> WebhookClient:
|
|
37
38
|
"""Get the Slack client."""
|
|
38
|
-
return WebhookClient(url, timeout=
|
|
39
|
+
return WebhookClient(url, timeout=safe_round(in_seconds(timeout)))
|
|
39
40
|
|
|
40
41
|
|
|
41
42
|
async def send_to_slack_async(
|
|
@@ -43,12 +44,12 @@ async def send_to_slack_async(
|
|
|
43
44
|
text: str,
|
|
44
45
|
/,
|
|
45
46
|
*,
|
|
46
|
-
timeout:
|
|
47
|
+
timeout: Duration = _TIMEOUT,
|
|
47
48
|
error: MaybeType[BaseException] = TimeoutError,
|
|
48
49
|
) -> None:
|
|
49
50
|
"""Send a message via Slack."""
|
|
50
51
|
client = _get_async_client(url, timeout=timeout)
|
|
51
|
-
async with
|
|
52
|
+
async with utilities.asyncio.timeout(timeout, error=error):
|
|
52
53
|
response = await client.send(text=text)
|
|
53
54
|
if response.status_code != HTTPStatus.OK: # pragma: no cover
|
|
54
55
|
raise SendToSlackError(text=text, response=response)
|
|
@@ -56,10 +57,10 @@ async def send_to_slack_async(
|
|
|
56
57
|
|
|
57
58
|
@cache
|
|
58
59
|
def _get_async_client(
|
|
59
|
-
url: str, /, *, timeout:
|
|
60
|
+
url: str, /, *, timeout: Duration = _TIMEOUT
|
|
60
61
|
) -> AsyncWebhookClient:
|
|
61
62
|
"""Get the Slack client."""
|
|
62
|
-
return AsyncWebhookClient(url, timeout=
|
|
63
|
+
return AsyncWebhookClient(url, timeout=safe_round(in_seconds(timeout)))
|
|
63
64
|
|
|
64
65
|
|
|
65
66
|
@dataclass(kw_only=True, slots=True)
|
utilities/sqlalchemy.py
CHANGED
|
@@ -65,33 +65,39 @@ from sqlalchemy.orm import (
|
|
|
65
65
|
from sqlalchemy.orm.exc import UnmappedClassError
|
|
66
66
|
from sqlalchemy.pool import NullPool, Pool
|
|
67
67
|
|
|
68
|
-
|
|
69
|
-
from utilities.
|
|
68
|
+
import utilities.asyncio
|
|
69
|
+
from utilities.core import (
|
|
70
|
+
OneEmptyError,
|
|
71
|
+
OneNonUniqueError,
|
|
72
|
+
chunked,
|
|
73
|
+
get_class_name,
|
|
74
|
+
is_pytest,
|
|
75
|
+
normalize_multi_line_str,
|
|
76
|
+
one,
|
|
77
|
+
repr_,
|
|
78
|
+
snake_case,
|
|
79
|
+
)
|
|
80
|
+
from utilities.functions import ensure_str, yield_object_attributes
|
|
70
81
|
from utilities.iterables import (
|
|
71
82
|
CheckLengthError,
|
|
72
83
|
CheckSubSetError,
|
|
73
|
-
OneEmptyError,
|
|
74
|
-
OneNonUniqueError,
|
|
75
84
|
check_length,
|
|
76
85
|
check_subset,
|
|
77
|
-
chunked,
|
|
78
86
|
merge_sets,
|
|
79
87
|
merge_str_mappings,
|
|
80
|
-
one,
|
|
81
88
|
)
|
|
82
|
-
from utilities.
|
|
83
|
-
from utilities.reprlib import get_repr
|
|
84
|
-
from utilities.text import secret_str, snake_case
|
|
89
|
+
from utilities.text import secret_str
|
|
85
90
|
from utilities.types import (
|
|
86
|
-
|
|
91
|
+
Duration,
|
|
87
92
|
MaybeIterable,
|
|
88
93
|
MaybeType,
|
|
94
|
+
StrDict,
|
|
89
95
|
StrMapping,
|
|
90
96
|
TupleOrStrMapping,
|
|
91
97
|
)
|
|
92
98
|
from utilities.typing import (
|
|
93
99
|
is_sequence_of_tuple_or_str_mapping,
|
|
94
|
-
|
|
100
|
+
is_str_mapping,
|
|
95
101
|
is_tuple,
|
|
96
102
|
is_tuple_or_str_mapping,
|
|
97
103
|
)
|
|
@@ -127,12 +133,15 @@ async def check_connect_async(
|
|
|
127
133
|
engine: AsyncEngine,
|
|
128
134
|
/,
|
|
129
135
|
*,
|
|
130
|
-
timeout:
|
|
136
|
+
timeout: Duration | None = None,
|
|
131
137
|
error: MaybeType[BaseException] = TimeoutError,
|
|
132
138
|
) -> bool:
|
|
133
139
|
"""Check if an engine can connect."""
|
|
134
140
|
try:
|
|
135
|
-
async with
|
|
141
|
+
async with (
|
|
142
|
+
utilities.asyncio.timeout(timeout, error=error),
|
|
143
|
+
engine.connect() as conn,
|
|
144
|
+
):
|
|
136
145
|
return bool((await conn.execute(_SELECT)).scalar_one())
|
|
137
146
|
except (gaierror, ConnectionRefusedError, DatabaseError, TimeoutError):
|
|
138
147
|
return False
|
|
@@ -145,7 +154,7 @@ async def check_engine(
|
|
|
145
154
|
engine: AsyncEngine,
|
|
146
155
|
/,
|
|
147
156
|
*,
|
|
148
|
-
timeout:
|
|
157
|
+
timeout: Duration | None = None,
|
|
149
158
|
error: MaybeType[BaseException] = TimeoutError,
|
|
150
159
|
num_tables: int | tuple[int, float] | None = None,
|
|
151
160
|
) -> None:
|
|
@@ -183,7 +192,7 @@ class CheckEngineError(Exception):
|
|
|
183
192
|
|
|
184
193
|
@override
|
|
185
194
|
def __str__(self) -> str:
|
|
186
|
-
return f"{
|
|
195
|
+
return f"{repr_(self.engine)} must have {self.expected} table(s); got {len(self.rows)}"
|
|
187
196
|
|
|
188
197
|
|
|
189
198
|
##
|
|
@@ -352,7 +361,7 @@ async def ensure_tables_created(
|
|
|
352
361
|
engine: AsyncEngine,
|
|
353
362
|
/,
|
|
354
363
|
*tables_or_orms: TableOrORMInstOrClass,
|
|
355
|
-
timeout:
|
|
364
|
+
timeout: Duration | None = None,
|
|
356
365
|
error: MaybeType[BaseException] = TimeoutError,
|
|
357
366
|
) -> None:
|
|
358
367
|
"""Ensure a table/set of tables is/are created."""
|
|
@@ -381,7 +390,7 @@ async def ensure_tables_created(
|
|
|
381
390
|
async def ensure_tables_dropped(
|
|
382
391
|
engine: AsyncEngine,
|
|
383
392
|
*tables_or_orms: TableOrORMInstOrClass,
|
|
384
|
-
timeout:
|
|
393
|
+
timeout: Duration | None = None,
|
|
385
394
|
error: MaybeType[BaseException] = TimeoutError,
|
|
386
395
|
) -> None:
|
|
387
396
|
"""Ensure a table/set of tables is/are dropped."""
|
|
@@ -615,9 +624,9 @@ async def insert_items(
|
|
|
615
624
|
is_upsert: bool = False,
|
|
616
625
|
chunk_size_frac: float = CHUNK_SIZE_FRAC,
|
|
617
626
|
assume_tables_exist: bool = False,
|
|
618
|
-
timeout_create:
|
|
627
|
+
timeout_create: Duration | None = None,
|
|
619
628
|
error_create: MaybeType[BaseException] = TimeoutError,
|
|
620
|
-
timeout_insert:
|
|
629
|
+
timeout_insert: Duration | None = None,
|
|
621
630
|
error_insert: MaybeType[BaseException] = TimeoutError,
|
|
622
631
|
) -> None:
|
|
623
632
|
"""Insert a set of items into a database.
|
|
@@ -860,9 +869,9 @@ async def migrate_data(
|
|
|
860
869
|
table_or_orm_to: TableOrORMInstOrClass | None = None,
|
|
861
870
|
chunk_size_frac: float = CHUNK_SIZE_FRAC,
|
|
862
871
|
assume_tables_exist: bool = False,
|
|
863
|
-
timeout_create:
|
|
872
|
+
timeout_create: Duration | None = None,
|
|
864
873
|
error_create: MaybeType[BaseException] = TimeoutError,
|
|
865
|
-
timeout_insert:
|
|
874
|
+
timeout_insert: Duration | None = None,
|
|
866
875
|
error_insert: MaybeType[BaseException] = TimeoutError,
|
|
867
876
|
) -> None:
|
|
868
877
|
"""Migrate the contents of a table from one database to another."""
|
|
@@ -894,7 +903,7 @@ def selectable_to_string(
|
|
|
894
903
|
com = selectable.compile(
|
|
895
904
|
dialect=engine_or_conn.dialect, compile_kwargs={"literal_binds": True}
|
|
896
905
|
)
|
|
897
|
-
return str(com)
|
|
906
|
+
return normalize_multi_line_str(str(com))
|
|
898
907
|
|
|
899
908
|
|
|
900
909
|
##
|
|
@@ -916,12 +925,15 @@ async def yield_connection(
|
|
|
916
925
|
engine: AsyncEngine,
|
|
917
926
|
/,
|
|
918
927
|
*,
|
|
919
|
-
timeout:
|
|
928
|
+
timeout: Duration | None = None,
|
|
920
929
|
error: MaybeType[BaseException] = TimeoutError,
|
|
921
930
|
) -> AsyncIterator[AsyncConnection]:
|
|
922
931
|
"""Yield an async connection."""
|
|
923
932
|
try:
|
|
924
|
-
async with
|
|
933
|
+
async with (
|
|
934
|
+
utilities.asyncio.timeout(timeout, error=error),
|
|
935
|
+
engine.begin() as conn,
|
|
936
|
+
):
|
|
925
937
|
yield conn
|
|
926
938
|
except GeneratorExit: # pragma: no cover
|
|
927
939
|
if not is_pytest():
|
|
@@ -1020,7 +1032,7 @@ def _is_pair_of_str_mapping_and_table(
|
|
|
1020
1032
|
obj: Any, /
|
|
1021
1033
|
) -> TypeGuard[_PairOfStrMappingAndTable]:
|
|
1022
1034
|
"""Check if an object is a pair of a string mapping and a table."""
|
|
1023
|
-
return _is_pair_with_predicate_and_table(obj,
|
|
1035
|
+
return _is_pair_with_predicate_and_table(obj, is_str_mapping)
|
|
1024
1036
|
|
|
1025
1037
|
|
|
1026
1038
|
def _is_pair_of_tuple_and_table(obj: Any, /) -> TypeGuard[_PairOfTupleAndTable]:
|
|
@@ -1063,7 +1075,7 @@ def _map_mapping_to_table(
|
|
|
1063
1075
|
mapping=mapping, columns=columns, extra=error.extra
|
|
1064
1076
|
) from None
|
|
1065
1077
|
return {k: v for k, v in mapping.items() if k in columns}
|
|
1066
|
-
out:
|
|
1078
|
+
out: StrDict = {}
|
|
1067
1079
|
for key, value in mapping.items():
|
|
1068
1080
|
try:
|
|
1069
1081
|
col = one(c for c in columns if snake_case(c) == snake_case(key))
|
|
@@ -1096,7 +1108,7 @@ class _MapMappingToTableExtraColumnsError(_MapMappingToTableError):
|
|
|
1096
1108
|
|
|
1097
1109
|
@override
|
|
1098
1110
|
def __str__(self) -> str:
|
|
1099
|
-
return f"Mapping {
|
|
1111
|
+
return f"Mapping {repr_(self.mapping)} must be a subset of table columns {repr_(self.columns)}; got extra {self.extra}"
|
|
1100
1112
|
|
|
1101
1113
|
|
|
1102
1114
|
@dataclass(kw_only=True, slots=True)
|
|
@@ -1105,7 +1117,7 @@ class _MapMappingToTableSnakeMapEmptyError(_MapMappingToTableError):
|
|
|
1105
1117
|
|
|
1106
1118
|
@override
|
|
1107
1119
|
def __str__(self) -> str:
|
|
1108
|
-
return f"Mapping {
|
|
1120
|
+
return f"Mapping {repr_(self.mapping)} must be a subset of table columns {repr_(self.columns)}; cannot find column to map to {self.key!r} modulo snake casing"
|
|
1109
1121
|
|
|
1110
1122
|
|
|
1111
1123
|
@dataclass(kw_only=True, slots=True)
|
|
@@ -1116,7 +1128,7 @@ class _MapMappingToTableSnakeMapNonUniqueError(_MapMappingToTableError):
|
|
|
1116
1128
|
|
|
1117
1129
|
@override
|
|
1118
1130
|
def __str__(self) -> str:
|
|
1119
|
-
return f"Mapping {
|
|
1131
|
+
return f"Mapping {repr_(self.mapping)} must be a subset of table columns {repr_(self.columns)}; found columns {self.first!r}, {self.second!r} and perhaps more to map to {self.key!r} modulo snake casing"
|
|
1120
1132
|
|
|
1121
1133
|
|
|
1122
1134
|
##
|
|
@@ -1156,7 +1168,7 @@ def _orm_inst_to_dict_predicate(
|
|
|
1156
1168
|
|
|
1157
1169
|
def _tuple_to_mapping(
|
|
1158
1170
|
values: tuple[Any, ...], table_or_orm: TableOrORMInstOrClass, /
|
|
1159
|
-
) ->
|
|
1171
|
+
) -> StrDict:
|
|
1160
1172
|
columns = get_column_names(table_or_orm)
|
|
1161
1173
|
mapping = dict(zip(columns, tuple(values), strict=False))
|
|
1162
1174
|
return {k: v for k, v in mapping.items() if v is not None}
|
utilities/sqlalchemy_polars.py
CHANGED
|
@@ -26,17 +26,11 @@ from polars import (
|
|
|
26
26
|
from sqlalchemy import Column, Select, select
|
|
27
27
|
from sqlalchemy.exc import DuplicateColumnError
|
|
28
28
|
|
|
29
|
-
|
|
30
|
-
from utilities.
|
|
31
|
-
from utilities.
|
|
32
|
-
|
|
33
|
-
OneError,
|
|
34
|
-
check_duplicates,
|
|
35
|
-
chunked,
|
|
36
|
-
one,
|
|
37
|
-
)
|
|
29
|
+
import utilities.asyncio
|
|
30
|
+
from utilities.constants import UTC
|
|
31
|
+
from utilities.core import OneError, chunked, identity, one, repr_, snake_case
|
|
32
|
+
from utilities.iterables import CheckDuplicatesError, check_duplicates
|
|
38
33
|
from utilities.polars import zoned_date_time_dtype
|
|
39
|
-
from utilities.reprlib import get_repr
|
|
40
34
|
from utilities.sqlalchemy import (
|
|
41
35
|
CHUNK_SIZE_FRAC,
|
|
42
36
|
TableOrORMInstOrClass,
|
|
@@ -45,9 +39,7 @@ from utilities.sqlalchemy import (
|
|
|
45
39
|
get_columns,
|
|
46
40
|
insert_items,
|
|
47
41
|
)
|
|
48
|
-
from utilities.text import snake_case
|
|
49
42
|
from utilities.typing import is_subclass_gen
|
|
50
|
-
from utilities.zoneinfo import UTC
|
|
51
43
|
|
|
52
44
|
if TYPE_CHECKING:
|
|
53
45
|
from collections.abc import (
|
|
@@ -62,9 +54,8 @@ if TYPE_CHECKING:
|
|
|
62
54
|
from sqlalchemy.ext.asyncio import AsyncEngine
|
|
63
55
|
from sqlalchemy.sql import ColumnCollection
|
|
64
56
|
from sqlalchemy.sql.base import ReadOnlyColumnCollection
|
|
65
|
-
from whenever import TimeDelta
|
|
66
57
|
|
|
67
|
-
from utilities.types import
|
|
58
|
+
from utilities.types import Duration, MaybeType, TimeZoneLike
|
|
68
59
|
|
|
69
60
|
|
|
70
61
|
async def insert_dataframe(
|
|
@@ -77,9 +68,9 @@ async def insert_dataframe(
|
|
|
77
68
|
is_upsert: bool = False,
|
|
78
69
|
chunk_size_frac: float = CHUNK_SIZE_FRAC,
|
|
79
70
|
assume_tables_exist: bool = False,
|
|
80
|
-
timeout_create:
|
|
71
|
+
timeout_create: Duration | None = None,
|
|
81
72
|
error_create: type[Exception] = TimeoutError,
|
|
82
|
-
timeout_insert:
|
|
73
|
+
timeout_insert: Duration | None = None,
|
|
83
74
|
error_insert: type[Exception] = TimeoutError,
|
|
84
75
|
) -> None:
|
|
85
76
|
"""Insert/upsert a DataFrame into a database."""
|
|
@@ -169,7 +160,7 @@ class _InsertDataFrameMapDFColumnToTableColumnAndTypeError(Exception):
|
|
|
169
160
|
|
|
170
161
|
@override
|
|
171
162
|
def __str__(self) -> str:
|
|
172
|
-
return f"Unable to map DataFrame column {self.df_col_name!r} into table schema {
|
|
163
|
+
return f"Unable to map DataFrame column {self.df_col_name!r} into table schema {repr_(self.table_schema)} with snake={self.snake}"
|
|
173
164
|
|
|
174
165
|
|
|
175
166
|
def _insert_dataframe_check_df_and_db_types(
|
|
@@ -200,7 +191,7 @@ async def select_to_dataframe(
|
|
|
200
191
|
in_clauses: tuple[Column[Any], Iterable[Any]] | None = None,
|
|
201
192
|
in_clauses_chunk_size: int | None = None,
|
|
202
193
|
chunk_size_frac: float = CHUNK_SIZE_FRAC,
|
|
203
|
-
timeout:
|
|
194
|
+
timeout: Duration | None = None,
|
|
204
195
|
error: MaybeType[BaseException] = TimeoutError,
|
|
205
196
|
**kwargs: Any,
|
|
206
197
|
) -> DataFrame: ...
|
|
@@ -216,7 +207,7 @@ async def select_to_dataframe(
|
|
|
216
207
|
in_clauses: None = None,
|
|
217
208
|
in_clauses_chunk_size: int | None = None,
|
|
218
209
|
chunk_size_frac: float = CHUNK_SIZE_FRAC,
|
|
219
|
-
timeout:
|
|
210
|
+
timeout: Duration | None = None,
|
|
220
211
|
error: MaybeType[BaseException] = TimeoutError,
|
|
221
212
|
**kwargs: Any,
|
|
222
213
|
) -> Iterable[DataFrame]: ...
|
|
@@ -232,7 +223,7 @@ async def select_to_dataframe(
|
|
|
232
223
|
in_clauses: tuple[Column[Any], Iterable[Any]],
|
|
233
224
|
in_clauses_chunk_size: int | None = None,
|
|
234
225
|
chunk_size_frac: float = CHUNK_SIZE_FRAC,
|
|
235
|
-
timeout:
|
|
226
|
+
timeout: Duration | None = None,
|
|
236
227
|
error: MaybeType[BaseException] = TimeoutError,
|
|
237
228
|
**kwargs: Any,
|
|
238
229
|
) -> AsyncIterable[DataFrame]: ...
|
|
@@ -248,7 +239,7 @@ async def select_to_dataframe(
|
|
|
248
239
|
in_clauses: tuple[Column[Any], Iterable[Any]] | None = None,
|
|
249
240
|
in_clauses_chunk_size: int | None = None,
|
|
250
241
|
chunk_size_frac: float = CHUNK_SIZE_FRAC,
|
|
251
|
-
timeout:
|
|
242
|
+
timeout: Duration | None = None,
|
|
252
243
|
error: MaybeType[BaseException] = TimeoutError,
|
|
253
244
|
**kwargs: Any,
|
|
254
245
|
) -> DataFrame | Iterable[DataFrame] | AsyncIterable[DataFrame]: ...
|
|
@@ -263,7 +254,7 @@ async def select_to_dataframe(
|
|
|
263
254
|
in_clauses: tuple[Column[Any], Iterable[Any]] | None = None,
|
|
264
255
|
in_clauses_chunk_size: int | None = None,
|
|
265
256
|
chunk_size_frac: float = CHUNK_SIZE_FRAC,
|
|
266
|
-
timeout:
|
|
257
|
+
timeout: Duration | None = None,
|
|
267
258
|
error: MaybeType[BaseException] = TimeoutError,
|
|
268
259
|
**kwargs: Any,
|
|
269
260
|
) -> DataFrame | Iterable[DataFrame] | AsyncIterable[DataFrame]:
|
|
@@ -272,7 +263,7 @@ async def select_to_dataframe(
|
|
|
272
263
|
sel = _select_to_dataframe_apply_snake(sel)
|
|
273
264
|
schema = _select_to_dataframe_map_select_to_df_schema(sel, time_zone=time_zone)
|
|
274
265
|
if in_clauses is None:
|
|
275
|
-
async with
|
|
266
|
+
async with utilities.asyncio.timeout(timeout, error=error):
|
|
276
267
|
return read_database(
|
|
277
268
|
sel,
|
|
278
269
|
cast("Any", engine),
|
|
@@ -289,7 +280,7 @@ async def select_to_dataframe(
|
|
|
289
280
|
chunk_size_frac=chunk_size_frac,
|
|
290
281
|
)
|
|
291
282
|
if batch_size is None:
|
|
292
|
-
async with
|
|
283
|
+
async with utilities.asyncio.timeout(timeout, error=error):
|
|
293
284
|
dfs = [
|
|
294
285
|
await select_to_dataframe(
|
|
295
286
|
sel,
|
|
@@ -310,7 +301,7 @@ async def select_to_dataframe(
|
|
|
310
301
|
return DataFrame(schema=schema)
|
|
311
302
|
|
|
312
303
|
async def yield_dfs() -> AsyncIterator[DataFrame]:
|
|
313
|
-
async with
|
|
304
|
+
async with utilities.asyncio.timeout(timeout, error=error):
|
|
314
305
|
for sel_i in sels:
|
|
315
306
|
for df in await select_to_dataframe(
|
|
316
307
|
sel_i,
|