resourcecap 0.2.0__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.0
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.0"
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"
@@ -17,9 +17,6 @@ classifiers = [
17
17
  name = "purushottam-dafure"
18
18
  email = "purushottamdafure@gmail.com"
19
19
 
20
- [project.scripts]
21
- resourcecap = "resourcecap:main"
22
-
23
20
  [build-system]
24
21
  requires = ["uv_build>=0.12.17,<0.13.0"]
25
22
  build-backend = "uv_build"
@@ -1,6 +1,6 @@
1
1
  [project]
2
2
  name = "resourcecap"
3
- version = "0.2.0"
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"
@@ -16,9 +16,6 @@ classifiers = [
16
16
  "Typing :: Typed",
17
17
  ]
18
18
 
19
- [project.scripts]
20
- resourcecap = "resourcecap:main"
21
-
22
19
  [build-system]
23
20
  requires = ["uv_build>=0.12.17,<0.13.0"]
24
21
  build-backend = "uv_build"
@@ -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"]
@@ -1,5 +1,6 @@
1
1
  import contextvars
2
2
  import logging
3
+ import threading
3
4
  from collections.abc import Hashable, Mapping
4
5
  from types import TracebackType
5
6
 
@@ -13,6 +14,18 @@ class BudgetExhaustedError(Exception):
13
14
  """Raised when a `Budget`'s limit for some resource is exceeded."""
14
15
 
15
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
+
16
29
  class Budget:
17
30
  """Context manager enforcing a spending limit on `costs`-decorated calls made within it.
18
31
 
@@ -29,7 +42,9 @@ class Budget:
29
42
  that limit before the decorated call runs rather than after.
30
43
 
31
44
  Each resource's spend resets to zero every time the block is
32
- 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`.
33
48
  """
34
49
 
35
50
  def __init__(
@@ -39,7 +54,8 @@ class Budget:
39
54
  ) -> None:
40
55
  self.limits = merge_resources(amount, limits)
41
56
  self.spent: dict[Hashable, float] = {}
42
- self._token: contextvars.Token[tuple[_tracking._Chargeable, ...]] | None = None
57
+ self._lock = threading.Lock()
58
+ self._token: contextvars.Token[tuple[Budget, ...]] | None = None
43
59
 
44
60
  def exhausts_at_start(self, key: Hashable) -> bool:
45
61
  """Whether `key`'s limit should be enforced before the call runs."""
@@ -53,14 +69,15 @@ class Budget:
53
69
  (spend charged against it so far), and `remaining` (`allocated
54
70
  - used`, which can go negative for a `warn_only` resource).
55
71
  """
56
- return {
57
- key: {
58
- "allocated": limit.amount,
59
- "used": self.spent.get(key, 0.0),
60
- "remaining": limit.amount - self.spent.get(key, 0.0),
72
+ with self._lock:
73
+ return {
74
+ key: {
75
+ "allocated": limit.amount,
76
+ "used": self.spent.get(key, 0.0),
77
+ "remaining": limit.amount - self.spent.get(key, 0.0),
78
+ }
79
+ for key, limit in self.limits.items()
61
80
  }
62
- for key, limit in self.limits.items()
63
- }
64
81
 
65
82
  def charge(self, key: Hashable, amount: float) -> None:
66
83
  """Record spend against this budget's limit for `key`.
@@ -73,8 +90,10 @@ class Budget:
73
90
  if limit is None:
74
91
  return
75
92
 
76
- spent = self.spent.get(key, 0.0) + amount
77
- self.spent[key] = spent
93
+ with self._lock:
94
+ spent = self.spent.get(key, 0.0) + amount
95
+ self.spent[key] = spent
96
+
78
97
  if spent > limit.amount:
79
98
  if limit.warn_only:
80
99
  logger.warning(
@@ -86,8 +105,15 @@ class Budget:
86
105
  )
87
106
 
88
107
  def __enter__(self) -> "Budget":
89
- self.spent = {}
90
- self._token = _tracking.push(self)
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
+ )
115
+ self.spent = {}
116
+ self._token = _tracking.push(self)
91
117
  return self
92
118
 
93
119
  def __exit__(
@@ -96,9 +122,10 @@ class Budget:
96
122
  exc: BaseException | None,
97
123
  tb: TracebackType | None,
98
124
  ) -> None:
99
- assert self._token is not None
100
- _tracking.pop(self._token)
101
- self._token = None
125
+ with self._lock:
126
+ assert self._token is not None
127
+ _tracking.pop(self._token)
128
+ self._token = None
102
129
 
103
130
  async def __aenter__(self) -> "Budget":
104
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,9 +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", "main"]
6
-
7
-
8
- def main() -> None:
9
- print("Hello from resourcecap!")