diff --git a/flask_pymemcache.py b/flask_pymemcache.py index 0fd18bb..7dd8497 100644 --- a/flask_pymemcache.py +++ b/flask_pymemcache.py @@ -54,6 +54,7 @@ memcache.client.set('foo', 'bar') """ + from __future__ import absolute_import, division, print_function import flask import pymemcache.client @@ -61,7 +62,6 @@ class FlaskPyMemcache(object): - def __init__(self, app=None, conf_key=None): """ :type app: flask.Flask @@ -76,26 +76,29 @@ def init_app(self, app, conf_key=None): :type app: flask.Flask :parm str conf_key: Key of flask config. """ - conf_key = conf_key or self.conf_key or 'PYMEMCACHE' + conf_key = conf_key or self.conf_key or "PYMEMCACHE" self.conf_key = conf_key conf = app.config[conf_key] if not isinstance(conf, dict): raise TypeError("Flask-PyMemcache conf should be dict") - close_on_teardown = conf.pop('close_on_teardown', False) + close_on_teardown = conf.pop("close_on_teardown", False) - if isinstance(conf['server'], list): - conf['servers'] = conf.pop('server') + if isinstance(conf["server"], list): + conf["servers"] = conf.pop("server") client = pymemcache.client.hash.HashClient(**conf) - elif isinstance(conf['server'], tuple): + elif isinstance(conf["server"], tuple): client = pymemcache.client.Client(**conf) else: - raise TypeError("Flask-PyMemcache conf['server'] should be tuple or list of tuples") + raise TypeError( + "Flask-PyMemcache conf['server'] should be tuple or list of tuples" + ) - app.extensions.setdefault('pymemcache', {}) - app.extensions['pymemcache'][self] = client + app.extensions.setdefault("pymemcache", {}) + app.extensions["pymemcache"][self] = client if close_on_teardown: + @app.teardown_appcontext def close_connection(exc=None): client.close() @@ -105,4 +108,4 @@ def client(self): """ :rtype: pymemcache.client.Client """ - return flask.current_app.extensions['pymemcache'][self] + return flask.current_app.extensions["pymemcache"][self] diff --git a/test_flask_pymemcache.py b/test_flask_pymemcache.py index 0d0d63b..9871409 100644 --- a/test_flask_pymemcache.py +++ b/test_flask_pymemcache.py @@ -9,19 +9,20 @@ class TestFlaskPyMemcache(TestCase): def test_simple(self): - pymc = pymemcache.client.Client(('localhost', 11211)) + pymc = pymemcache.client.Client(("localhost", 11211)) memcache = flask_pymemcache.FlaskPyMemcache() app = flask.Flask(__name__) - app.config['PYMEMCACHE'] = { - 'server': ('localhost', 11211), - 'key_prefix': b'px', - 'close_on_teardown': False} + app.config["PYMEMCACHE"] = { + "server": ("localhost", 11211), + "key_prefix": b"px", + "close_on_teardown": False, + } memcache.init_app(app) with app.app_context(): - memcache.client.set(b'foo', b'bar') + memcache.client.set(b"foo", b"bar") with app.app_context(): - assert memcache.client.get(b'foo') == b'bar' + assert memcache.client.get(b"foo") == b"bar" - assert pymc.get(b'pxfoo') == b'bar' + assert pymc.get(b"pxfoo") == b"bar"