add client_class option (#3)

Author: methane

pyproject.toml

@@ -14,7 +14,7 @@ authors = [
 ]
 dependencies = [
     "Flask",
-    "pymemcache>=1.3.4",
+    "pymemcache>=3.4.4",
 ]
 dev-dependencies = [
     "pytest",

src/flask_pymemcache.py

@@ -51,25 +51,25 @@
 ::
 
     memcache.client.set('foo', 'bar')
-
 """
 
+from __future__ import annotations
+
 import flask
-import pymemcache.client
-import pymemcache.client.hash
+import pymemcache
 
 
 class FlaskPyMemcache:
-    def __init__(self, app=None, conf_key=None):
+    def __init__(self, app=None, conf_key=None, client_class=None) -> None:
         """
         :type app: flask.Flask
         :parm str conf_key: Key of flask config.
         """
         self.conf_key = conf_key
         if app is not None:
-            self.init_app(app, conf_key)
+            self.init_app(app, conf_key, client_class)
 
-    def init_app(self, app, conf_key=None):
+    def init_app(self, app, conf_key=None, client_class=None) -> None:
         """
         :type app: flask.Flask
         :parm str conf_key: Key of flask config.
@@ -84,9 +84,9 @@ def init_app(self, app, conf_key=None):
 
         if isinstance(conf["server"], list):
             conf["servers"] = conf.pop("server")
-            client = pymemcache.client.hash.HashClient(**conf)
+            client = (client_class or pymemcache.HashClient)(**conf)
         elif isinstance(conf["server"], tuple):
-            client = pymemcache.client.Client(**conf)
+            client = (client_class or pymemcache.Client)(**conf)
         else:
             raise TypeError(
                 "Flask-PyMemcache conf['server'] should be tuple or list of tuples"
@@ -102,8 +102,5 @@ def close_connection(exc=None):
                 client.close()
 
     @property
-    def client(self):
-        """
-        :rtype: pymemcache.client.Client
-        """
+    def client(self) -> pymemcache.Client:
         return flask.current_app.extensions["pymemcache"][self]

tests/test_flask_pymemcache.py

@@ -24,3 +24,17 @@ def test_simple(self):
             assert memcache.client.get(b"foo") == b"bar"
 
         assert pymc.get(b"pxfoo") == b"bar"
+
+    def test_pool(self):
+        memcache = flask_pymemcache.FlaskPyMemcache()
+        app = flask.Flask(__name__)
+        app.config["PYMEMCACHE"] = {
+            "server": ("localhost", 11211),
+            "key_prefix": b"px",
+            "close_on_teardown": False,
+        }
+        memcache.init_app(app, client_class=pymemcache.PooledClient)
+
+        with app.app_context():
+            assert memcache.client.get(b"foo") is None
+            assert isinstance(memcache.client, pymemcache.PooledClient)