eth-portfolio 1.1.0__tar.gz

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.

Potentially problematic release.


This version of eth-portfolio might be problematic. Click here for more details.

Files changed (55) hide show
  1. eth_portfolio-1.1.0/PKG-INFO +174 -0
  2. eth_portfolio-1.1.0/README.md +145 -0
  3. eth_portfolio-1.1.0/eth_portfolio/__init__.py +16 -0
  4. eth_portfolio-1.1.0/eth_portfolio/_argspec.py +42 -0
  5. eth_portfolio-1.1.0/eth_portfolio/_cache.py +116 -0
  6. eth_portfolio-1.1.0/eth_portfolio/_config.py +3 -0
  7. eth_portfolio-1.1.0/eth_portfolio/_db/__init__.py +0 -0
  8. eth_portfolio-1.1.0/eth_portfolio/_db/decorators.py +147 -0
  9. eth_portfolio-1.1.0/eth_portfolio/_db/entities.py +204 -0
  10. eth_portfolio-1.1.0/eth_portfolio/_db/utils.py +595 -0
  11. eth_portfolio-1.1.0/eth_portfolio/_decimal.py +122 -0
  12. eth_portfolio-1.1.0/eth_portfolio/_decorators.py +71 -0
  13. eth_portfolio-1.1.0/eth_portfolio/_exceptions.py +67 -0
  14. eth_portfolio-1.1.0/eth_portfolio/_ledgers/__init__.py +0 -0
  15. eth_portfolio-1.1.0/eth_portfolio/_ledgers/address.py +892 -0
  16. eth_portfolio-1.1.0/eth_portfolio/_ledgers/portfolio.py +327 -0
  17. eth_portfolio-1.1.0/eth_portfolio/_loaders/__init__.py +33 -0
  18. eth_portfolio-1.1.0/eth_portfolio/_loaders/balances.py +78 -0
  19. eth_portfolio-1.1.0/eth_portfolio/_loaders/token_transfer.py +214 -0
  20. eth_portfolio-1.1.0/eth_portfolio/_loaders/transaction.py +379 -0
  21. eth_portfolio-1.1.0/eth_portfolio/_loaders/utils.py +59 -0
  22. eth_portfolio-1.1.0/eth_portfolio/_shitcoins.py +212 -0
  23. eth_portfolio-1.1.0/eth_portfolio/_utils.py +286 -0
  24. eth_portfolio-1.1.0/eth_portfolio/_ydb/__init__.py +0 -0
  25. eth_portfolio-1.1.0/eth_portfolio/_ydb/token_transfers.py +136 -0
  26. eth_portfolio-1.1.0/eth_portfolio/address.py +382 -0
  27. eth_portfolio-1.1.0/eth_portfolio/buckets.py +181 -0
  28. eth_portfolio-1.1.0/eth_portfolio/constants.py +58 -0
  29. eth_portfolio-1.1.0/eth_portfolio/portfolio.py +629 -0
  30. eth_portfolio-1.1.0/eth_portfolio/protocols/__init__.py +66 -0
  31. eth_portfolio-1.1.0/eth_portfolio/protocols/_base.py +107 -0
  32. eth_portfolio-1.1.0/eth_portfolio/protocols/convex.py +17 -0
  33. eth_portfolio-1.1.0/eth_portfolio/protocols/dsr.py +31 -0
  34. eth_portfolio-1.1.0/eth_portfolio/protocols/lending/__init__.py +49 -0
  35. eth_portfolio-1.1.0/eth_portfolio/protocols/lending/_base.py +57 -0
  36. eth_portfolio-1.1.0/eth_portfolio/protocols/lending/compound.py +185 -0
  37. eth_portfolio-1.1.0/eth_portfolio/protocols/lending/liquity.py +110 -0
  38. eth_portfolio-1.1.0/eth_portfolio/protocols/lending/maker.py +105 -0
  39. eth_portfolio-1.1.0/eth_portfolio/protocols/lending/unit.py +47 -0
  40. eth_portfolio-1.1.0/eth_portfolio/protocols/liquity.py +16 -0
  41. eth_portfolio-1.1.0/eth_portfolio/py.typed +0 -0
  42. eth_portfolio-1.1.0/eth_portfolio/structs/__init__.py +43 -0
  43. eth_portfolio-1.1.0/eth_portfolio/structs/modified.py +69 -0
  44. eth_portfolio-1.1.0/eth_portfolio/structs/structs.py +637 -0
  45. eth_portfolio-1.1.0/eth_portfolio/typing.py +1460 -0
  46. eth_portfolio-1.1.0/eth_portfolio.egg-info/PKG-INFO +174 -0
  47. eth_portfolio-1.1.0/eth_portfolio.egg-info/SOURCES.txt +53 -0
  48. eth_portfolio-1.1.0/eth_portfolio.egg-info/dependency_links.txt +1 -0
  49. eth_portfolio-1.1.0/eth_portfolio.egg-info/requires.txt +9 -0
  50. eth_portfolio-1.1.0/eth_portfolio.egg-info/top_level.txt +1 -0
  51. eth_portfolio-1.1.0/pyproject.toml +22 -0
  52. eth_portfolio-1.1.0/setup.cfg +4 -0
  53. eth_portfolio-1.1.0/setup.py +31 -0
  54. eth_portfolio-1.1.0/tests/test_portfolio.py +0 -0
  55. eth_portfolio-1.1.0/tests/test_typing.py +635 -0
@@ -0,0 +1,174 @@
1
+ Metadata-Version: 2.2
2
+ Name: eth-portfolio
3
+ Version: 1.1.0
4
+ Summary: eth-portfolio makes it easy to analyze your portfolio.
5
+ Home-page: https://github.com/BobTheBuidler/eth-portfolio
6
+ Author: BobTheBuidler
7
+ Author-email: bobthebuidlerdefi@gmail.com
8
+ Project-URL: Homepage, https://github.com/BobTheBuidler/eth-portfolio
9
+ Project-URL: Documentation, https://bobthebuidler.github.io/eth-portfolio
10
+ Project-URL: Source Code, https://github.com/BobTheBuidler/eth-portfolio
11
+ Description-Content-Type: text/markdown
12
+ Requires-Dist: checksum_dict>=1.1.2
13
+ Requires-Dist: dank_mids>=4.20.109
14
+ Requires-Dist: eth-brownie<1.21,>=1.19.3
15
+ Requires-Dist: eth_retry<1,>=0.1.15
16
+ Requires-Dist: evmspec>=0.2.0
17
+ Requires-Dist: ez-a-sync>=0.24.43
18
+ Requires-Dist: numpy<2
19
+ Requires-Dist: pandas<1.6,>=1.4.3
20
+ Requires-Dist: ypricemagic<5,>=4.0.57
21
+ Dynamic: author
22
+ Dynamic: author-email
23
+ Dynamic: description
24
+ Dynamic: description-content-type
25
+ Dynamic: home-page
26
+ Dynamic: project-url
27
+ Dynamic: requires-dist
28
+ Dynamic: summary
29
+
30
+ # eth-portfolio
31
+ Use `eth-portfolio` to output information about your portfolio in a streamlined, speed-optimized way.
32
+
33
+ ### NOTES:
34
+ - This lib is still a WIP and the provided API is subject to change without notice.
35
+
36
+ ### INSTALLATION:
37
+ - You should start with a fresh virtual environment, and then just...
38
+ ```
39
+ pip install git+https://github.com/BobTheBuidler/eth-portfolio
40
+ ```
41
+
42
+ **Known Issues**
43
+ Make sure you are using Python >= 3.8 and < 3.13
44
+ If you have a PyYaml Issue with 3.4.1 not installing due to an issue with cython, try the following:
45
+ ```
46
+ pip install wheel
47
+ pip install --no-build-isolation "Cython<3" "pyyaml==5.4.1"
48
+ ```
49
+ then try again
50
+ `
51
+ pip install git+https://github.com/BobTheBuidler/eth-portfolio
52
+ `
53
+
54
+
55
+ ### USE:
56
+ For basic use, input each of your addresses as environment variables using the following pattern:
57
+ ```
58
+ PORTFOLIO_ADDRESS_0=0x123...
59
+ PORTFOLIO_ADDRESS_1=0x234...
60
+ PORTFOLIO_ADDRESS_2=0x345...
61
+ ```
62
+
63
+ Then do...
64
+ ```
65
+ from eth_portfolio import portfolio
66
+ portfolio.eth_balance(block)
67
+
68
+ >>> {
69
+ 0xaddress0: _BalanceItem(balance=1234, usd_value=5678)
70
+ 0xaddress1: _BalanceItem(balance=1234, usd_value=5678)
71
+ 0xaddress2: _BalanceItem(balance=1234, usd_value=5678)
72
+ }
73
+ ```
74
+
75
+ Getting token transfers:
76
+ ```
77
+ from eth_portfolio import portfolio
78
+ token_transfers = portfolio.token_transfers.get(start_block, end_block)
79
+ token_transfers.df()
80
+
81
+ >>> {
82
+ 0xaddress0: AddressTokenTransfersLedger(...) # Each of these contains the token transfers for the specified address
83
+ 0xaddress1: AddressTokenTransfersLedger(...)
84
+ 0xaddress2: AddressTokenTransfersLedger(...)
85
+ }
86
+ ```
87
+
88
+ Getting transactions as a DataFrame:
89
+ ```
90
+ from eth_portfolio import portfolio
91
+ txs = portfolio.transactions.get(start_block, end_block)
92
+ txs.df()
93
+
94
+ >>> [I am a pretend DataFrame]
95
+ ```
96
+
97
+ Getting assets:
98
+ ```
99
+ from eth_portfolio import portfolio
100
+ portfolio.describe(start_block, end_block)
101
+
102
+ >>> {
103
+ 'assets': {
104
+ 'wallet0_address': {
105
+ 'token0': {
106
+ 'amount': 123,
107
+ 'value usd: 456,
108
+ },
109
+ 'token1': {
110
+ 'amount': 123,
111
+ 'value usd: 456,
112
+ },
113
+ },
114
+ 'wallet0_address': {
115
+ 'token0': {
116
+ 'amount': 123,
117
+ 'value usd: 456,
118
+ },
119
+ 'token1': {
120
+ 'amount': 123,
121
+ 'value usd: 456,
122
+ },
123
+ },
124
+ },
125
+ 'debt': {
126
+ 'wallet0_address': {
127
+ 'token0': {
128
+ 'amount': 123,
129
+ 'value usd: 456,
130
+ },
131
+ 'token1': {
132
+ 'amount': 123,
133
+ 'value usd: 456,
134
+ },
135
+ },
136
+ 'wallet1_address': {
137
+ 'token0': {
138
+ 'amount': 123,
139
+ 'value usd: 456,
140
+ },
141
+ 'token1': {
142
+ 'amount': 123,
143
+ 'value usd: 456,
144
+ },
145
+ },
146
+ },
147
+ }
148
+ ```
149
+
150
+ Getting net worth:
151
+ ```
152
+ from eth_portfolio import portfolio
153
+ desc = portfolio.describe(block)
154
+ assets = desc['assets'] # OR you can do `assets = portfolio.assets(block)`
155
+ debt = desc['debt'] # OR you can do `debt = portfolio.debt(block)`
156
+ assets = sum(assets.values())
157
+ debt = sum(debt.values())
158
+ net = assets - debt
159
+ net.sum_usd()
160
+
161
+ >>> Decimal("123456.78900")
162
+ ```
163
+
164
+ ### ADVANCED:
165
+ You also have more granular control available using the Portfolio object:
166
+ ```
167
+ from eth_portfolio import Portfolio
168
+ port = Portfolio([0xaddress0, 0xaddress1, 0xaddress2])
169
+ port.describe(chain.height)
170
+
171
+ # Or for async code
172
+ async_port = Portfolio([0xaddress0, 0xaddress1, 0xaddress2], asynchronous=True)
173
+ await port.describe(chain.height)
174
+ ```
@@ -0,0 +1,145 @@
1
+ # eth-portfolio
2
+ Use `eth-portfolio` to output information about your portfolio in a streamlined, speed-optimized way.
3
+
4
+ ### NOTES:
5
+ - This lib is still a WIP and the provided API is subject to change without notice.
6
+
7
+ ### INSTALLATION:
8
+ - You should start with a fresh virtual environment, and then just...
9
+ ```
10
+ pip install git+https://github.com/BobTheBuidler/eth-portfolio
11
+ ```
12
+
13
+ **Known Issues**
14
+ Make sure you are using Python >= 3.8 and < 3.13
15
+ If you have a PyYaml Issue with 3.4.1 not installing due to an issue with cython, try the following:
16
+ ```
17
+ pip install wheel
18
+ pip install --no-build-isolation "Cython<3" "pyyaml==5.4.1"
19
+ ```
20
+ then try again
21
+ `
22
+ pip install git+https://github.com/BobTheBuidler/eth-portfolio
23
+ `
24
+
25
+
26
+ ### USE:
27
+ For basic use, input each of your addresses as environment variables using the following pattern:
28
+ ```
29
+ PORTFOLIO_ADDRESS_0=0x123...
30
+ PORTFOLIO_ADDRESS_1=0x234...
31
+ PORTFOLIO_ADDRESS_2=0x345...
32
+ ```
33
+
34
+ Then do...
35
+ ```
36
+ from eth_portfolio import portfolio
37
+ portfolio.eth_balance(block)
38
+
39
+ >>> {
40
+ 0xaddress0: _BalanceItem(balance=1234, usd_value=5678)
41
+ 0xaddress1: _BalanceItem(balance=1234, usd_value=5678)
42
+ 0xaddress2: _BalanceItem(balance=1234, usd_value=5678)
43
+ }
44
+ ```
45
+
46
+ Getting token transfers:
47
+ ```
48
+ from eth_portfolio import portfolio
49
+ token_transfers = portfolio.token_transfers.get(start_block, end_block)
50
+ token_transfers.df()
51
+
52
+ >>> {
53
+ 0xaddress0: AddressTokenTransfersLedger(...) # Each of these contains the token transfers for the specified address
54
+ 0xaddress1: AddressTokenTransfersLedger(...)
55
+ 0xaddress2: AddressTokenTransfersLedger(...)
56
+ }
57
+ ```
58
+
59
+ Getting transactions as a DataFrame:
60
+ ```
61
+ from eth_portfolio import portfolio
62
+ txs = portfolio.transactions.get(start_block, end_block)
63
+ txs.df()
64
+
65
+ >>> [I am a pretend DataFrame]
66
+ ```
67
+
68
+ Getting assets:
69
+ ```
70
+ from eth_portfolio import portfolio
71
+ portfolio.describe(start_block, end_block)
72
+
73
+ >>> {
74
+ 'assets': {
75
+ 'wallet0_address': {
76
+ 'token0': {
77
+ 'amount': 123,
78
+ 'value usd: 456,
79
+ },
80
+ 'token1': {
81
+ 'amount': 123,
82
+ 'value usd: 456,
83
+ },
84
+ },
85
+ 'wallet0_address': {
86
+ 'token0': {
87
+ 'amount': 123,
88
+ 'value usd: 456,
89
+ },
90
+ 'token1': {
91
+ 'amount': 123,
92
+ 'value usd: 456,
93
+ },
94
+ },
95
+ },
96
+ 'debt': {
97
+ 'wallet0_address': {
98
+ 'token0': {
99
+ 'amount': 123,
100
+ 'value usd: 456,
101
+ },
102
+ 'token1': {
103
+ 'amount': 123,
104
+ 'value usd: 456,
105
+ },
106
+ },
107
+ 'wallet1_address': {
108
+ 'token0': {
109
+ 'amount': 123,
110
+ 'value usd: 456,
111
+ },
112
+ 'token1': {
113
+ 'amount': 123,
114
+ 'value usd: 456,
115
+ },
116
+ },
117
+ },
118
+ }
119
+ ```
120
+
121
+ Getting net worth:
122
+ ```
123
+ from eth_portfolio import portfolio
124
+ desc = portfolio.describe(block)
125
+ assets = desc['assets'] # OR you can do `assets = portfolio.assets(block)`
126
+ debt = desc['debt'] # OR you can do `debt = portfolio.debt(block)`
127
+ assets = sum(assets.values())
128
+ debt = sum(debt.values())
129
+ net = assets - debt
130
+ net.sum_usd()
131
+
132
+ >>> Decimal("123456.78900")
133
+ ```
134
+
135
+ ### ADVANCED:
136
+ You also have more granular control available using the Portfolio object:
137
+ ```
138
+ from eth_portfolio import Portfolio
139
+ port = Portfolio([0xaddress0, 0xaddress1, 0xaddress2])
140
+ port.describe(chain.height)
141
+
142
+ # Or for async code
143
+ async_port = Portfolio([0xaddress0, 0xaddress1, 0xaddress2], asynchronous=True)
144
+ await port.describe(chain.height)
145
+ ```
@@ -0,0 +1,16 @@
1
+ import a_sync._smart
2
+
3
+ a_sync._smart.set_smart_task_factory()
4
+
5
+
6
+ from eth_portfolio.portfolio import Portfolio, portfolio
7
+
8
+ # make sure we init the extended db before we init ypm somewhere
9
+ from eth_portfolio._db import utils
10
+ from eth_portfolio._shitcoins import SHITCOINS
11
+
12
+ __all__ = [
13
+ "Portfolio",
14
+ "portfolio",
15
+ "SHITCOINS",
16
+ ]
@@ -0,0 +1,42 @@
1
+ from inspect import getfullargspec
2
+ from typing import Any, Callable, List, Tuple, Type
3
+
4
+ # WIP:
5
+
6
+
7
+ def get_args_type(sample: Callable) -> Tuple[Type, ...]:
8
+ argspec = getfullargspec(sample)
9
+ args = {
10
+ arg_name: argspec.annotations[arg_name] if arg_name in argspec.annotations else Any
11
+ for arg_name in argspec.args
12
+ if arg_name != "self"
13
+ }
14
+ return tuple(*args.values())
15
+
16
+
17
+ def get_kwargs_type(sample: Callable) -> Tuple[Type, ...]:
18
+ argspec = getfullargspec(sample)
19
+ num_kwargs = len(argspec.args) - len(argspec.defaults or [])
20
+ kwarg_names = argspec.args[num_kwargs:]
21
+ kwargs = {
22
+ kwarg_name: argspec.annotations[kwarg_name] if kwarg_name in argspec.annotations else Any
23
+ for kwarg_name in kwarg_names
24
+ }
25
+ _kwarg_types: List[Type[object]] = list(*kwargs.values())
26
+ if num_kwargs == 1:
27
+ return Tuple[_kwarg_types[0]] # type: ignore [valid-type,return-value]
28
+ elif num_kwargs == 2:
29
+ return Tuple[_kwarg_types[0], _kwarg_types[1]] # type: ignore [misc,return-value]
30
+ elif num_kwargs == 3:
31
+ return Tuple[_kwarg_types[0], _kwarg_types[1], _kwarg_types[2]] # type: ignore [misc,return-value]
32
+ else:
33
+ return Any # type: ignore [misc,return-value]
34
+
35
+
36
+ def get_return_type(sample: Callable) -> Type:
37
+ argspec = getfullargspec(sample)
38
+ return argspec.annotations["return"] if "return" in argspec.annotations else Any
39
+
40
+
41
+ def get_types(sample: Callable) -> Tuple[Type, Type, Type]:
42
+ return get_args_type(sample), get_kwargs_type(sample), get_return_type(sample) # type: ignore [return-value]
@@ -0,0 +1,116 @@
1
+ import functools
2
+ import inspect
3
+ from asyncio import PriorityQueue, Task, current_task, get_event_loop
4
+ from concurrent.futures import Executor
5
+ from hashlib import md5
6
+ from logging import getLogger
7
+ from os import makedirs
8
+ from os.path import exists, join
9
+ from pickle import dumps, load, loads
10
+ from random import random
11
+ from typing import Any, Callable, List, NoReturn
12
+
13
+ from a_sync import PruningThreadPoolExecutor
14
+ from a_sync._typing import P, T
15
+ from a_sync.asyncio import create_task
16
+ from a_sync.primitives.queue import log_broken
17
+ from aiofiles import open as _aio_open
18
+ from brownie import chain
19
+
20
+ BASE_PATH = f"./cache/{chain.id}/"
21
+ _THREAD_NAME_PREFIX = "eth-portfolio-cache-decorator"
22
+ _EXISTS_EXECUTOR = PruningThreadPoolExecutor(8, f"{_THREAD_NAME_PREFIX}-exists")
23
+
24
+
25
+ def cache_to_disk(fn: Callable[P, T]) -> Callable[P, T]:
26
+ # sourcery skip: use-contextlib-suppress
27
+ name = fn.__name__
28
+ cache_path_for_fn = f"{BASE_PATH}{fn.__module__.replace('.', '/')}/{name}"
29
+ logger = getLogger(f"eth_portfolio.cache_to_disk.{name}")
30
+
31
+ def get_cache_file_path(args, kwargs):
32
+ # Create a unique filename based on the function arguments
33
+ key = md5(dumps((args, sorted(kwargs.items())))).hexdigest()
34
+ return join(cache_path_for_fn, f"{key}.json")
35
+
36
+ write_executor = PruningThreadPoolExecutor(8, f"{_THREAD_NAME_PREFIX}-{fn.__qualname__}-write")
37
+
38
+ makedirs(cache_path_for_fn, exist_ok=True)
39
+
40
+ if inspect.iscoroutinefunction(fn):
41
+ read_executor = PruningThreadPoolExecutor(
42
+ 8, f"{_THREAD_NAME_PREFIX}-{fn.__qualname__}-read"
43
+ )
44
+
45
+ queue = PriorityQueue()
46
+
47
+ async def cache_deco_worker_coro(func) -> NoReturn:
48
+ try:
49
+ while True:
50
+ _, fut, cache_path, args, kwargs = await queue.get()
51
+ try:
52
+ async with _aio_open(cache_path, "rb", executor=read_executor) as f:
53
+ fut.set_result(loads(await f.read()))
54
+ except Exception as e:
55
+ fut.set_exception(e)
56
+ except Exception as e:
57
+ logger.error("%s for %s is broken!!!", current_task(), func)
58
+ logger.exception(e)
59
+ raise
60
+
61
+ workers: List[Task[NoReturn]] = []
62
+
63
+ @functools.wraps(fn)
64
+ async def disk_cache_wrap(*args: P.args, **kwargs: P.kwargs) -> T:
65
+ cache_path = get_cache_file_path(args, kwargs)
66
+ if await _EXISTS_EXECUTOR.run(exists, cache_path):
67
+ fut = get_event_loop().create_future()
68
+ # we intentionally mix up the order to break up heavy load block ranges
69
+ queue.put_nowait((random(), fut, cache_path, args, kwargs))
70
+ if not workers:
71
+ workers.extend(create_task(cache_deco_worker_coro(fn)) for _ in range(100))
72
+ try:
73
+ return await fut
74
+ except EOFError:
75
+ pass
76
+
77
+ async_result: T = await fn(*args, **kwargs)
78
+ try:
79
+ await __cache_write(cache_path, async_result, write_executor)
80
+ except OSError as e:
81
+ # I was having some weird issues in docker that I don't want to debug,
82
+ # so I'm going to assume you have another means to let you know you're
83
+ # out of disk space and will pass right on through here so my script
84
+ # can continue
85
+ if e.strerror != "No space left on device":
86
+ raise
87
+ return async_result
88
+
89
+ else:
90
+
91
+ @functools.wraps(fn)
92
+ def disk_cache_wrap(*args: P.args, **kwargs: P.kwargs) -> T:
93
+ cache_path = get_cache_file_path(args, kwargs)
94
+ try:
95
+ with open(cache_path, "rb") as f:
96
+ return load(f)
97
+ except (FileNotFoundError, EOFError):
98
+ pass
99
+
100
+ sync_result: T = fn(*args, **kwargs) # type: ignore [assignment, return-value]
101
+ try:
102
+ create_task(
103
+ coro=__cache_write(cache_path, sync_result, write_executor),
104
+ skip_gc_until_done=True,
105
+ )
106
+ except RuntimeError:
107
+ pass
108
+ return sync_result
109
+
110
+ return disk_cache_wrap
111
+
112
+
113
+ async def __cache_write(cache_path: str, result: Any, executor: Executor) -> None:
114
+ result = dumps(result)
115
+ async with _aio_open(cache_path, "wb", executor=executor) as f:
116
+ await f.write(result)
@@ -0,0 +1,3 @@
1
+ import os
2
+
3
+ REORG_BUFFER = int(os.environ.get("REORG_BUFFER", 30))
File without changes
@@ -0,0 +1,147 @@
1
+ from asyncio import iscoroutinefunction
2
+ from asyncio import sleep as aio_sleep
3
+ from functools import wraps
4
+ from logging import DEBUG, getLogger
5
+ from random import random
6
+ from time import sleep as time_sleep
7
+ from typing import Callable, TypeVar
8
+
9
+ from a_sync._typing import AnyFn
10
+ from pony.orm import OperationalError, TransactionError
11
+ from typing_extensions import ParamSpec
12
+
13
+
14
+ P = ParamSpec("P")
15
+ T = TypeVar("T")
16
+
17
+
18
+ logger = getLogger(__name__)
19
+ __logger_is_enabled_for = logger.isEnabledFor
20
+ __logger_warning = logger.warning
21
+ __logger_log = logger._log
22
+
23
+
24
+ def break_locks(fn: AnyFn[P, T]) -> AnyFn[P, T]:
25
+ """
26
+ Decorator to handle database lock errors by retrying the function.
27
+
28
+ This decorator is designed to wrap functions that interact with a database
29
+ and may encounter `OperationalError` due to database locks. It will retry
30
+ the function indefinitely if a "database is locked" error occurs. After
31
+ 5 attempts, a warning is logged, but the function will continue to retry
32
+ until it succeeds or a non-lock-related error occurs.
33
+
34
+ Args:
35
+ fn: The function to be wrapped, which may be a coroutine or a regular function.
36
+
37
+ Examples:
38
+ Basic usage with a regular function:
39
+
40
+ >>> @break_locks
41
+ ... def my_function():
42
+ ... # Function logic that may encounter a database lock
43
+ ... pass
44
+
45
+ Basic usage with an asynchronous function:
46
+
47
+ >>> @break_locks
48
+ ... async def my_async_function():
49
+ ... # Async function logic that may encounter a database lock
50
+ ... pass
51
+
52
+ See Also:
53
+ - :func:`pony.orm.db_session`: For managing database sessions.
54
+ """
55
+ if iscoroutinefunction(fn):
56
+
57
+ @wraps(fn)
58
+ async def break_locks_wrap(*args: P.args, **kwargs: P.kwargs) -> T:
59
+ debug_logs_enabled = None
60
+ tries = 0
61
+ while True:
62
+ try:
63
+ return await fn(*args, **kwargs)
64
+ except OperationalError as e:
65
+ if str(e) != "database is locked":
66
+ raise e
67
+
68
+ if debug_logs_enabled is None:
69
+ debug_logs_enabled = __logger_is_enabled_for(DEBUG)
70
+
71
+ if debug_logs_enabled is True:
72
+ __logger_log(DEBUG, "%s.%s %s", (fn.__module__, fn.__name__, e))
73
+
74
+ await aio_sleep(tries * random())
75
+ tries += 1
76
+ if tries > 5:
77
+ __logger_warning("%s caught in err loop with %s", fn, e)
78
+
79
+ else:
80
+
81
+ @wraps(fn)
82
+ def break_locks_wrap(*args: P.args, **kwargs: P.kwargs) -> T:
83
+ debug_logs_enabled = None
84
+ tries = 0
85
+ while True:
86
+ try:
87
+ return fn(*args, **kwargs) # type: ignore [return-value]
88
+ except OperationalError as e:
89
+ if str(e) != "database is locked":
90
+ raise e
91
+
92
+ if debug_logs_enabled is None:
93
+ debug_logs_enabled = __logger_is_enabled_for(DEBUG)
94
+
95
+ if debug_logs_enabled is True:
96
+ __logger_log(DEBUG, "%s.%s %s", (fn.__module__, fn.__name__, e))
97
+
98
+ time_sleep(tries * random())
99
+ tries += 1
100
+ if tries > 5:
101
+ __logger_warning("%s caught in err loop with %s", fn, e)
102
+
103
+ return break_locks_wrap
104
+
105
+
106
+ def requery_objs_on_diff_tx_err(fn: Callable[P, T]) -> Callable[P, T]:
107
+ """
108
+ Decorator to handle transaction errors by retrying the function.
109
+
110
+ This decorator is designed to wrap functions that may encounter
111
+ `TransactionError` due to mixing objects from different transactions.
112
+ It will retry the function until it succeeds or a non-transaction-related
113
+ error occurs.
114
+
115
+ Args:
116
+ fn: The function to be wrapped, which must not be a coroutine.
117
+
118
+ Raises:
119
+ TypeError: If the function is a coroutine.
120
+
121
+ Examples:
122
+ Basic usage with a function that may encounter transaction errors:
123
+
124
+ >>> @requery_objs_on_diff_tx_err
125
+ ... def my_function():
126
+ ... # Function logic that may encounter a transaction error
127
+ ... pass
128
+
129
+ See Also:
130
+ - :func:`pony.orm.db_session`: For managing database sessions.
131
+ """
132
+ if iscoroutinefunction(fn):
133
+ raise TypeError(f"{fn} must not be async")
134
+
135
+ @wraps(fn)
136
+ def requery_wrap(*args: P.args, **kwargs: P.kwargs) -> T:
137
+ while True:
138
+ try:
139
+ return fn(*args, **kwargs)
140
+ except TransactionError as e:
141
+ if str(e) != "An attempt to mix objects belonging to different transactions":
142
+ raise e
143
+ # The error occurs if you committed new objs to the db and started a new transaction while still inside of a `db_session`,
144
+ # and then tried to use the newly committed objects in the next transaction. Now that the objects are in the db this will
145
+ # not reoccur. The next iteration will be successful.
146
+
147
+ return requery_wrap