From 9fc77bcbefbb758ea5c857ba54c90c8b641365c1 Mon Sep 17 00:00:00 2001
From: Inada Naoki <songofacandy@gmail.com>
Date: Wed, 10 Jul 2024 15:47:34 +0900
Subject: [PATCH] add client_class option

---
 pyproject.toml                 |  2 +-
 src/flask_pymemcache.py        | 21 +++++++++------------
 tests/test_flask_pymemcache.py | 14 ++++++++++++++
 3 files changed, 24 insertions(+), 13 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index 7b9037e..96714e1 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -14,7 +14,7 @@ authors = [
 ]
 dependencies = [
     "Flask",
-    "pymemcache>=1.3.4",
+    "pymemcache>=3.4.4",
 ]
 dev-dependencies = [
     "pytest",
diff --git a/src/flask_pymemcache.py b/src/flask_pymemcache.py
index 0d239e4..27bcc6b 100644
--- a/src/flask_pymemcache.py
+++ b/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]
diff --git a/tests/test_flask_pymemcache.py b/tests/test_flask_pymemcache.py
index 8e96340..517a265 100644
--- a/tests/test_flask_pymemcache.py
+++ b/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)