resourcecap 0.2.1__tar.gz → 0.2.2__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.
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: resourcecap
3
- Version: 0.2.1
3
+ Version: 0.2.2
4
4
  Summary: A python package to limit the resource usage of any program
5
5
  Author: purushottam-dafure
6
6
  Author-email: purushottam-dafure <purushottamdafure@gmail.com>
@@ -97,6 +97,13 @@ Each call to `foo()` logs something like:
97
97
  the above code will stop execution at that point and raise BudgetExhaustedError,
98
98
  after the 4th call to `foo()` pushes the spend from 30 to 40.
99
99
 
100
+ A `Budget` instance can be reused across separate, non-overlapping `with`
101
+ blocks (its spend resets each time), but it can't be open twice at once —
102
+ entering it while it's already open, from a concurrent task/thread or
103
+ otherwise, raises `BudgetReentryError`. If you need several concurrent
104
+ pieces of work to share one budget, open it once around all of them; if
105
+ each needs independent tracking, give each its own `Budget` instance.
106
+
100
107
  ## Multiple resources
101
108
 
102
109
  `amount` on `costs`/`Budget` tracks a single, default resource. To track
@@ -83,6 +83,13 @@ Each call to `foo()` logs something like:
83
83
  the above code will stop execution at that point and raise BudgetExhaustedError,
84
84
  after the 4th call to `foo()` pushes the spend from 30 to 40.
85
85
 
86
+ A `Budget` instance can be reused across separate, non-overlapping `with`
87
+ blocks (its spend resets each time), but it can't be open twice at once —
88
+ entering it while it's already open, from a concurrent task/thread or
89
+ otherwise, raises `BudgetReentryError`. If you need several concurrent
90
+ pieces of work to share one budget, open it once around all of them; if
91
+ each needs independent tracking, give each its own `Budget` instance.
92
+
86
93
  ## Multiple resources
87
94
 
88
95
  `amount` on `costs`/`Budget` tracks a single, default resource. To track
@@ -1,6 +1,6 @@
1
1
  [project]
2
2
  name = "resourcecap"
3
- version = "0.2.1"
3
+ version = "0.2.2"
4
4
  description = "A python package to limit the resource usage of any program"
5
5
  readme = "README.md"
6
6
  license = "MIT"
@@ -1,6 +1,6 @@
1
1
  [project]
2
2
  name = "resourcecap"
3
- version = "0.2.1"
3
+ version = "0.2.2"
4
4
  description = "A python package to limit the resource usage of any program"
5
5
  readme = "README.md"
6
6
  license = "MIT"
@@ -0,0 +1,5 @@
1
+ from ._budget import Budget, BudgetExhaustedError, BudgetReentryError
2
+ from ._costs import costs
3
+ from ._resources import Spendable
4
+
5
+ __all__ = ["Budget", "BudgetExhaustedError", "BudgetReentryError", "Spendable", "costs"]
@@ -14,6 +14,18 @@ class BudgetExhaustedError(Exception):
14
14
  """Raised when a `Budget`'s limit for some resource is exceeded."""
15
15
 
16
16
 
17
+ class BudgetReentryError(RuntimeError):
18
+ """Raised when a `Budget` is entered while it's already open.
19
+
20
+ A single `Budget` instance isn't safe to have open twice at once —
21
+ whether that's one task/thread re-entering it while still inside an
22
+ outer block on the same instance, or two concurrent tasks/threads
23
+ each entering it independently. Either enter it once around the
24
+ concurrent work (so everything inside shares that one open block),
25
+ or create a separate `Budget` instance per concurrent call.
26
+ """
27
+
28
+
17
29
  class Budget:
18
30
  """Context manager enforcing a spending limit on `costs`-decorated calls made within it.
19
31
 
@@ -30,7 +42,9 @@ class Budget:
30
42
  that limit before the decorated call runs rather than after.
31
43
 
32
44
  Each resource's spend resets to zero every time the block is
33
- entered, so a `Budget` instance can be reused across separate blocks.
45
+ entered, so a `Budget` instance can be reused across separate,
46
+ non-overlapping blocks. It is **not** safe to have the same
47
+ instance open twice at once — see `BudgetReentryError`.
34
48
  """
35
49
 
36
50
  def __init__(
@@ -41,7 +55,7 @@ class Budget:
41
55
  self.limits = merge_resources(amount, limits)
42
56
  self.spent: dict[Hashable, float] = {}
43
57
  self._lock = threading.Lock()
44
- self._token: contextvars.Token[tuple[_tracking._Chargeable, ...]] | None = None
58
+ self._token: contextvars.Token[tuple[Budget, ...]] | None = None
45
59
 
46
60
  def exhausts_at_start(self, key: Hashable) -> bool:
47
61
  """Whether `key`'s limit should be enforced before the call runs."""
@@ -92,8 +106,14 @@ class Budget:
92
106
 
93
107
  def __enter__(self) -> "Budget":
94
108
  with self._lock:
109
+ if self._token is not None:
110
+ raise BudgetReentryError(
111
+ "this Budget is already open: a Budget instance can't be entered "
112
+ "while it's still open elsewhere. Enter it once around the "
113
+ "concurrent work, or use a separate Budget per concurrent call."
114
+ )
95
115
  self.spent = {}
96
- self._token = _tracking.push(self)
116
+ self._token = _tracking.push(self)
97
117
  return self
98
118
 
99
119
  def __exit__(
@@ -102,9 +122,10 @@ class Budget:
102
122
  exc: BaseException | None,
103
123
  tb: TracebackType | None,
104
124
  ) -> None:
105
- assert self._token is not None
106
- _tracking.pop(self._token)
107
- self._token = None
125
+ with self._lock:
126
+ assert self._token is not None
127
+ _tracking.pop(self._token)
128
+ self._token = None
108
129
 
109
130
  async def __aenter__(self) -> "Budget":
110
131
  return self.__enter__()
@@ -85,9 +85,7 @@ def costs[**P, R](
85
85
  resolved.amount,
86
86
  total,
87
87
  )
88
- _tracking.charge_active_budgets_after(
89
- key, resolved.amount, partial_amounts.get(key, 0.0)
90
- )
88
+ _tracking.charge_active_budgets_after(key, resolved.amount, partial_amounts[key])
91
89
 
92
90
  if inspect.iscoroutinefunction(func):
93
91
  async_func = cast(Callable[P, Awaitable[R]], func)
@@ -64,6 +64,9 @@ class Spendable:
64
64
  def coerce(cls, value: "float | Spendable") -> "Spendable":
65
65
  return value if isinstance(value, Spendable) else cls(value)
66
66
 
67
+ def _with_amount(self, amount: float) -> "Spendable":
68
+ return Spendable(amount, warn_only=self.warn_only, exhaust_at_start=self.exhaust_at_start)
69
+
67
70
  def resolve(
68
71
  self, args: tuple[Any, ...], kwargs: dict[str, Any], result: Any # noqa: ANN401
69
72
  ) -> "Spendable":
@@ -80,7 +83,7 @@ class Spendable:
80
83
  total += self.from_args(*args, **kwargs)
81
84
  if self.from_result is not None:
82
85
  total += self.from_result(result)
83
- return Spendable(total, warn_only=self.warn_only, exhaust_at_start=self.exhaust_at_start)
86
+ return self._with_amount(total)
84
87
 
85
88
  def resolve_partial(self, args: tuple[Any, ...], kwargs: dict[str, Any]) -> "Spendable":
86
89
  """Like `resolve`, but only the static `amount` and `from_args`.
@@ -91,11 +94,7 @@ class Spendable:
91
94
  if self.from_args is None:
92
95
  return self
93
96
 
94
- return Spendable(
95
- self.amount + self.from_args(*args, **kwargs),
96
- warn_only=self.warn_only,
97
- exhaust_at_start=self.exhaust_at_start,
98
- )
97
+ return self._with_amount(self.amount + self.from_args(*args, **kwargs))
99
98
 
100
99
 
101
100
  def merge_resources(
@@ -1,32 +1,26 @@
1
+ from __future__ import annotations
2
+
1
3
  import contextvars
2
4
  from collections.abc import Hashable
3
- from typing import Protocol
4
-
5
-
6
- class _Chargeable(Protocol):
7
- def exhausts_at_start(self, key: Hashable) -> bool: ...
5
+ from typing import TYPE_CHECKING
8
6
 
9
- def charge(self, key: Hashable, amount: float) -> None: ...
7
+ if TYPE_CHECKING:
8
+ from ._budget import Budget
10
9
 
11
10
 
12
- _active_budgets: contextvars.ContextVar[tuple[_Chargeable, ...]] = contextvars.ContextVar(
11
+ _active_budgets: contextvars.ContextVar[tuple[Budget, ...]] = contextvars.ContextVar(
13
12
  "resourcecap_active_budgets", default=()
14
13
  )
15
14
 
16
15
 
17
- def push(budget: _Chargeable) -> contextvars.Token[tuple[_Chargeable, ...]]:
16
+ def push(budget: Budget) -> contextvars.Token[tuple[Budget, ...]]:
18
17
  return _active_budgets.set((*_active_budgets.get(), budget))
19
18
 
20
19
 
21
- def pop(token: contextvars.Token[tuple[_Chargeable, ...]]) -> None:
20
+ def pop(token: contextvars.Token[tuple[Budget, ...]]) -> None:
22
21
  _active_budgets.reset(token)
23
22
 
24
23
 
25
- def charge_active_budgets(key: Hashable, amount: float) -> None:
26
- for budget in _active_budgets.get():
27
- budget.charge(key, amount)
28
-
29
-
30
24
  def charge_active_budgets_at_start(key: Hashable, partial_amount: float) -> None:
31
25
  if not partial_amount:
32
26
  return
@@ -1,5 +0,0 @@
1
- from ._budget import Budget, BudgetExhaustedError
2
- from ._costs import costs
3
- from ._resources import Spendable
4
-
5
- __all__ = ["Budget", "BudgetExhaustedError", "Spendable", "costs"]