Initial commit

venv/lib/python3.10/site-packages/pyrate_limiter/buckets/mp_bucket.py

@@ -0,0 +1,53 @@
+"""multiprocessing In-memory Bucket using a multiprocessing.Manager.ListProxy
+     and a multiprocessing.Lock.
+"""
+from multiprocessing import Manager
+from multiprocessing import RLock
+from multiprocessing.managers import ListProxy
+from multiprocessing.synchronize import RLock as LockType
+from typing import List
+from typing import Optional
+
+from ..abstracts import Rate
+from ..abstracts import RateItem
+from pyrate_limiter.buckets import InMemoryBucket
+
+
+class MultiprocessBucket(InMemoryBucket):
+
+    items: List[RateItem]  # ListProxy
+    mp_lock: LockType
+
+    def __init__(self, rates: List[Rate], items: List[RateItem], mp_lock: LockType):
+
+        if not isinstance(items, ListProxy):
+            raise ValueError("items must be a ListProxy")
+
+        self.rates = sorted(rates, key=lambda r: r.interval)
+        self.items = items
+        self.mp_lock = mp_lock
+
+    def put(self, item: RateItem) -> bool:
+        with self.mp_lock:
+            return super().put(item)
+
+    def leak(self, current_timestamp: Optional[int] = None) -> int:
+        with self.mp_lock:
+            return super().leak(current_timestamp)
+
+    def limiter_lock(self):
+        return self.mp_lock
+
+    @classmethod
+    def init(
+        cls,
+        rates: List[Rate],
+    ):
+        """
+        Creates a single ListProxy so that this bucket can be shared across multiple processes.
+        """
+        shared_items: List[RateItem] = Manager().list()  # type: ignore[assignment]
+
+        mp_lock: LockType = RLock()
+
+        return cls(rates=rates, items=shared_items, mp_lock=mp_lock)