resourcecap 0.2.0__tar.gz → 0.2.2__tar.gz

{resourcecap-0.2.0 → resourcecap-0.2.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: resourcecap
-Version: 0.2.0
+Version: 0.2.2
 Summary: A python package to limit the resource usage of any program
 Author: purushottam-dafure
 Author-email: purushottam-dafure <purushottamdafure@gmail.com>
@@ -97,6 +97,13 @@ Each call to `foo()` logs something like:
 the above code will stop execution at that point and raise BudgetExhaustedError,
 after the 4th call to `foo()` pushes the spend from 30 to 40.
 
+A `Budget` instance can be reused across separate, non-overlapping `with`
+blocks (its spend resets each time), but it can't be open twice at once —
+entering it while it's already open, from a concurrent task/thread or
+otherwise, raises `BudgetReentryError`. If you need several concurrent
+pieces of work to share one budget, open it once around all of them; if
+each needs independent tracking, give each its own `Budget` instance.
+
 ## Multiple resources
 
 `amount` on `costs`/`Budget` tracks a single, default resource. To track

{resourcecap-0.2.0 → resourcecap-0.2.2}/README.md

@@ -83,6 +83,13 @@ Each call to `foo()` logs something like:
 the above code will stop execution at that point and raise BudgetExhaustedError,
 after the 4th call to `foo()` pushes the spend from 30 to 40.
 
+A `Budget` instance can be reused across separate, non-overlapping `with`
+blocks (its spend resets each time), but it can't be open twice at once —
+entering it while it's already open, from a concurrent task/thread or
+otherwise, raises `BudgetReentryError`. If you need several concurrent
+pieces of work to share one budget, open it once around all of them; if
+each needs independent tracking, give each its own `Budget` instance.
+
 ## Multiple resources
 
 `amount` on `costs`/`Budget` tracks a single, default resource. To track

{resourcecap-0.2.0 → resourcecap-0.2.2}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "resourcecap"
-version = "0.2.0"
+version = "0.2.2"
 description = "A python package to limit the resource usage of any program"
 readme = "README.md"
 license = "MIT"
@@ -17,9 +17,6 @@ classifiers = [
 name = "purushottam-dafure"
 email = "purushottamdafure@gmail.com"
 
-[project.scripts]
-resourcecap = "resourcecap:main"
-
 [build-system]
 requires = ["uv_build>=0.12.17,<0.13.0"]
 build-backend = "uv_build"

{resourcecap-0.2.0 → resourcecap-0.2.2}/pyproject.toml.orig

@@ -1,6 +1,6 @@
 [project]
 name = "resourcecap"
-version = "0.2.0"
+version = "0.2.2"
 description = "A python package to limit the resource usage of any program"
 readme = "README.md"
 license = "MIT"
@@ -16,9 +16,6 @@ classifiers = [
     "Typing :: Typed",
 ]
 
-[project.scripts]
-resourcecap = "resourcecap:main"
-
 [build-system]
 requires = ["uv_build>=0.12.17,<0.13.0"]
 build-backend = "uv_build"

resourcecap-0.2.2/src/resourcecap/__init__.py

@@ -0,0 +1,5 @@
+from ._budget import Budget, BudgetExhaustedError, BudgetReentryError
+from ._costs import costs
+from ._resources import Spendable
+
+__all__ = ["Budget", "BudgetExhaustedError", "BudgetReentryError", "Spendable", "costs"]

{resourcecap-0.2.0 → resourcecap-0.2.2}/src/resourcecap/_budget.py

@@ -1,5 +1,6 @@
 import contextvars
 import logging
+import threading
 from collections.abc import Hashable, Mapping
 from types import TracebackType
 
@@ -13,6 +14,18 @@ class BudgetExhaustedError(Exception):
     """Raised when a `Budget`'s limit for some resource is exceeded."""
 
 
+class BudgetReentryError(RuntimeError):
+    """Raised when a `Budget` is entered while it's already open.
+
+    A single `Budget` instance isn't safe to have open twice at once —
+    whether that's one task/thread re-entering it while still inside an
+    outer block on the same instance, or two concurrent tasks/threads
+    each entering it independently. Either enter it once around the
+    concurrent work (so everything inside shares that one open block),
+    or create a separate `Budget` instance per concurrent call.
+    """
+
+
 class Budget:
     """Context manager enforcing a spending limit on `costs`-decorated calls made within it.
 
@@ -29,7 +42,9 @@ class Budget:
     that limit before the decorated call runs rather than after.
 
     Each resource's spend resets to zero every time the block is
-    entered, so a `Budget` instance can be reused across separate blocks.
+    entered, so a `Budget` instance can be reused across separate,
+    non-overlapping blocks. It is **not** safe to have the same
+    instance open twice at once — see `BudgetReentryError`.
     """
 
     def __init__(
@@ -39,7 +54,8 @@ class Budget:
     ) -> None:
         self.limits = merge_resources(amount, limits)
         self.spent: dict[Hashable, float] = {}
-        self._token: contextvars.Token[tuple[_tracking._Chargeable, ...]] | None = None
+        self._lock = threading.Lock()
+        self._token: contextvars.Token[tuple[Budget, ...]] | None = None
 
     def exhausts_at_start(self, key: Hashable) -> bool:
         """Whether `key`'s limit should be enforced before the call runs."""
@@ -53,14 +69,15 @@ class Budget:
         (spend charged against it so far), and `remaining` (`allocated
         - used`, which can go negative for a `warn_only` resource).
         """
-        return {
-            key: {
-                "allocated": limit.amount,
-                "used": self.spent.get(key, 0.0),
-                "remaining": limit.amount - self.spent.get(key, 0.0),
+        with self._lock:
+            return {
+                key: {
+                    "allocated": limit.amount,
+                    "used": self.spent.get(key, 0.0),
+                    "remaining": limit.amount - self.spent.get(key, 0.0),
+                }
+                for key, limit in self.limits.items()
             }
-            for key, limit in self.limits.items()
-        }
 
     def charge(self, key: Hashable, amount: float) -> None:
         """Record spend against this budget's limit for `key`.
@@ -73,8 +90,10 @@ class Budget:
         if limit is None:
             return
 
-        spent = self.spent.get(key, 0.0) + amount
-        self.spent[key] = spent
+        with self._lock:
+            spent = self.spent.get(key, 0.0) + amount
+            self.spent[key] = spent
+
         if spent > limit.amount:
             if limit.warn_only:
                 logger.warning(
@@ -86,8 +105,15 @@ class Budget:
                 )
 
     def __enter__(self) -> "Budget":
-        self.spent = {}
-        self._token = _tracking.push(self)
+        with self._lock:
+            if self._token is not None:
+                raise BudgetReentryError(
+                    "this Budget is already open: a Budget instance can't be entered "
+                    "while it's still open elsewhere. Enter it once around the "
+                    "concurrent work, or use a separate Budget per concurrent call."
+                )
+            self.spent = {}
+            self._token = _tracking.push(self)
         return self
 
     def __exit__(
@@ -96,9 +122,10 @@ class Budget:
         exc: BaseException | None,
         tb: TracebackType | None,
     ) -> None:
-        assert self._token is not None
-        _tracking.pop(self._token)
-        self._token = None
+        with self._lock:
+            assert self._token is not None
+            _tracking.pop(self._token)
+            self._token = None
 
     async def __aenter__(self) -> "Budget":
         return self.__enter__()

{resourcecap-0.2.0 → resourcecap-0.2.2}/src/resourcecap/_costs.py

@@ -85,9 +85,7 @@ def costs[**P, R](
                     resolved.amount,
                     total,
                 )
-                _tracking.charge_active_budgets_after(
-                    key, resolved.amount, partial_amounts.get(key, 0.0)
-                )
+                _tracking.charge_active_budgets_after(key, resolved.amount, partial_amounts[key])
 
         if inspect.iscoroutinefunction(func):
             async_func = cast(Callable[P, Awaitable[R]], func)

{resourcecap-0.2.0 → resourcecap-0.2.2}/src/resourcecap/_resources.py

@@ -64,6 +64,9 @@ class Spendable:
     def coerce(cls, value: "float | Spendable") -> "Spendable":
         return value if isinstance(value, Spendable) else cls(value)
 
+    def _with_amount(self, amount: float) -> "Spendable":
+        return Spendable(amount, warn_only=self.warn_only, exhaust_at_start=self.exhaust_at_start)
+
     def resolve(
         self, args: tuple[Any, ...], kwargs: dict[str, Any], result: Any  # noqa: ANN401
     ) -> "Spendable":
@@ -80,7 +83,7 @@ class Spendable:
             total += self.from_args(*args, **kwargs)
         if self.from_result is not None:
             total += self.from_result(result)
-        return Spendable(total, warn_only=self.warn_only, exhaust_at_start=self.exhaust_at_start)
+        return self._with_amount(total)
 
     def resolve_partial(self, args: tuple[Any, ...], kwargs: dict[str, Any]) -> "Spendable":
         """Like `resolve`, but only the static `amount` and `from_args`.
@@ -91,11 +94,7 @@ class Spendable:
         if self.from_args is None:
             return self
 
-        return Spendable(
-            self.amount + self.from_args(*args, **kwargs),
-            warn_only=self.warn_only,
-            exhaust_at_start=self.exhaust_at_start,
-        )
+        return self._with_amount(self.amount + self.from_args(*args, **kwargs))
 
 
 def merge_resources(

{resourcecap-0.2.0 → resourcecap-0.2.2}/src/resourcecap/_tracking.py

@@ -1,32 +1,26 @@
+from __future__ import annotations
+
 import contextvars
 from collections.abc import Hashable
-from typing import Protocol
-
-
-class _Chargeable(Protocol):
-    def exhausts_at_start(self, key: Hashable) -> bool: ...
+from typing import TYPE_CHECKING
 
-    def charge(self, key: Hashable, amount: float) -> None: ...
+if TYPE_CHECKING:
+    from ._budget import Budget
 
 
-_active_budgets: contextvars.ContextVar[tuple[_Chargeable, ...]] = contextvars.ContextVar(
+_active_budgets: contextvars.ContextVar[tuple[Budget, ...]] = contextvars.ContextVar(
     "resourcecap_active_budgets", default=()
 )
 
 
-def push(budget: _Chargeable) -> contextvars.Token[tuple[_Chargeable, ...]]:
+def push(budget: Budget) -> contextvars.Token[tuple[Budget, ...]]:
     return _active_budgets.set((*_active_budgets.get(), budget))
 
 
-def pop(token: contextvars.Token[tuple[_Chargeable, ...]]) -> None:
+def pop(token: contextvars.Token[tuple[Budget, ...]]) -> None:
     _active_budgets.reset(token)
 
 
-def charge_active_budgets(key: Hashable, amount: float) -> None:
-    for budget in _active_budgets.get():
-        budget.charge(key, amount)
-
-
 def charge_active_budgets_at_start(key: Hashable, partial_amount: float) -> None:
     if not partial_amount:
         return

resourcecap-0.2.0/src/resourcecap/__init__.py

@@ -1,9 +0,0 @@
-from ._budget import Budget, BudgetExhaustedError
-from ._costs import costs
-from ._resources import Spendable
-
-__all__ = ["Budget", "BudgetExhaustedError", "Spendable", "costs", "main"]
-
-
-def main() -> None:
-    print("Hello from resourcecap!")