Add pooling helpers.

pylibmc.py

@@ -112,6 +112,86 @@ class Client(_pylibmc.client):
 
     behaviors = property(get_behaviors, set_behaviors)
 
+from contextlib import contextmanager
+
+class ClientPool(list):
+    """Client pooling helper.
+
+    This is mostly useful in threaded environments, because a client isn't
+    thread-safe at all. Instead, what you want to do is have each thread use
+    its own client, but you don't want to reconnect these all the time.
+
+    The solution is a pool, and this class is a helper for that.
+
+    >>> mc = Client(["127.0.0.1"])
+    >>> pool = ClientPool()
+    >>> pool.fill(mc, 4)
+    >>> with pool.reserve() as mc:
+    ...     mc.set("hi", "ho")
+    ...     mc.delete("hi")
+    ... 
+    """
+
+    @contextmanager
+    def reserve(self):
+        """Reserve a client, and put it back into the pool when done."""
+        mc = self.pop()
+        try:
+            yield mc
+        finally:
+            self.append(mc)
+
+    def fill(self, mc, n_slots):
+        """Fill *n_slots* of the pool with clones of *mc*."""
+        for i in xrange(n_slots):
+            self.append(mc.clone())
+
+class ThreadMappedPool(dict):
+    """Much like the *ClientPool*, helps you with pooling.
+
+    In a threaded environment, you'd most likely want to have a client per
+    thread. And there'd be no harm in one thread keeping the same client at all
+    times. So, why not map threads to clients? That's what this class does.
+
+    If a client is reserved, this class checks for a key based on the current
+    thread, and if none exists, clones the master client and inserts that key.
+
+    >>> mc = Client(["127.0.0.1"])
+    >>> pool = ThreadMappedPool(mc)
+    >>> with pool.reserve() as mc:
+    ...     mc.set("hi", "ho")
+    ...     mc.delete("hi")
+    ... 
+    """
+
+    def __new__(cls, master):
+        return super(ThreadMappedPool, cls).__new__(cls)
+
+    def __init__(self, master):
+        self.master = master
+
+    @contextmanager
+    def reserve(self):
+        """Reserve a client.
+
+        Creates a new client based on the master client if none exists for the
+        current thread.
+        """
+        key = threading.current_thread().name
+        mc = self.pop(key, None)
+        if mc is None:
+            mc = self.master.clone()
+        try:
+            yield mc
+        finally:
+            self[key] = mc
+
+# This makes sure ThreadMappedPool doesn't exist with non-thread Pythons.
+try:
+    import threading
+except ImportError:
+    del ThreadMappedPool
+
 if __name__ == "__main__":
     import doctest
     doctest.testmod()