diff --git a/CHANGELOG.md b/CHANGELOG.md index aa5babb..dedad35 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -12,6 +12,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - Added option in Flask Talisman to add Adobe Typekit CSP rules with `allow_typekit_content_security_policy=True` - Added `extra_headers` parameter in `Talisman` to update or add any global response headers - New pagination functions for populating pagination components +- New Flask decorators for specifying `Cache-Control` headers in response ### Changed diff --git a/docs/flask.md b/docs/flask.md index 851bada..b56eb04 100644 --- a/docs/flask.md +++ b/docs/flask.md @@ -1,16 +1,50 @@ # Flask -> Added in `v1.4.0`. +## Cache control + +> Added in `v1.5.0`. + +A set of decorators to manage the `Cache-Control` header of a route. + +### Examples + +```python +from flask import Flask +from tna_utilities.flask import cacheable_duration, do_not_cache, set_cache_control + +app = Flask(__name__) + +@app.route("/default/") +def default(): + return "No cache instructions given" + +@app.route("/cacheable/") +@cacheable_duration(3600) +def cachable_for_up_to_1h(): + return "You can cache me for up to an hour" + +@app.route("/not-cacheable/") +@do_not_cache() +def not_cachable(): + return "Don't cache me!" + +@app.route("/not-cacheable/") +@set_cache_control("private, max-age=120") +def private_cache(): + return "Cache me in private caches for up to 2 minutes" +``` ## `Talisman` +> Added in `v1.4.0`. + A stripped-down and opinionated reproduction of [wntrblm/flask-talisman](https://github.com/wntrblm/flask-talisman) which is a fork of [GoogleCloudPlatform/flask-talisman](https://github.com/GoogleCloudPlatform/flask-talisman). ### Examples ```python from flask import Flask -from tna_utilities.flask.talisman import Talisman +from tna_utilities.flask import Talisman app = Flask(__name__) Talisman(app) diff --git a/tests/test_flask_cache_control.py b/tests/test_flask_cache_control.py new file mode 100644 index 0000000..f7a9f9c --- /dev/null +++ b/tests/test_flask_cache_control.py @@ -0,0 +1,80 @@ +import unittest + +from flask import Flask +from tna_utilities.flask import cacheable_duration, do_not_cache, set_cache_control + + +class TestFlaskCacheControl(unittest.TestCase): + def setUp(self): + self.app = Flask(__name__) + self.test_client = self.app.test_client() + + def test_naked_route(self): + @self.app.route("/") + def index(): + return "OK" + + rv = self.test_client.get("/") + + self.assertEqual(rv.status_code, 200) + self.assertNotIn("Cache-Control", rv.headers) + + def test_do_not_cache_route(self): + @self.app.route("/") + @do_not_cache() + def index(): + return "OK" + + rv = self.test_client.get("/") + + self.assertEqual(rv.status_code, 200) + self.assertIn("Cache-Control", rv.headers) + self.assertEqual( + rv.headers["Cache-Control"], + "no-store", + ) + + def test_cacheable_duration_route(self): + @self.app.route("/") + @cacheable_duration() + def index(): + return "OK" + + rv = self.test_client.get("/") + + self.assertEqual(rv.status_code, 200) + self.assertIn("Cache-Control", rv.headers) + self.assertEqual( + rv.headers["Cache-Control"], + "public, max-age=3600", + ) + + def test_cacheable_duration_custom_duration_route(self): + @self.app.route("/") + @cacheable_duration(60) + def index(): + return "OK" + + rv = self.test_client.get("/") + + self.assertEqual(rv.status_code, 200) + self.assertIn("Cache-Control", rv.headers) + self.assertEqual( + rv.headers["Cache-Control"], + "public, max-age=60", + ) + + def test_set_cache_control_route(self): + @self.app.route("/") + @set_cache_control("private, max-age=120") + def index(): + return "OK" + + rv = self.test_client.get("/") + + self.assertEqual(rv.status_code, 200) + self.assertIn("Cache-Control", rv.headers) + self.assertEqual( + rv.headers["Cache-Control"], + "private, max-age=120", + ) diff --git a/tests/test_flask_talisman.py b/tests/test_flask_talisman.py index 5c91069..1c04b06 100644 --- a/tests/test_flask_talisman.py +++ b/tests/test_flask_talisman.py @@ -1,10 +1,10 @@ import unittest from flask import Flask, session -from tna_utilities.flask.talisman import Talisman +from tna_utilities.flask import Talisman -class TestTalisman(unittest.TestCase): +class TestFlaskTalisman(unittest.TestCase): def setUp(self): self.app = Flask(__name__) self.app.config["SECRET_KEY"] = "my_secret_key" diff --git a/tna_utilities/flask/__init__.py b/tna_utilities/flask/__init__.py new file mode 100644 index 0000000..2627994 --- /dev/null +++ b/tna_utilities/flask/__init__.py @@ -0,0 +1,6 @@ +from tna_utilities.flask.cache_control import ( # noqa: F401 + cacheable_duration, + do_not_cache, + set_cache_control, +) +from tna_utilities.flask.talisman import Talisman # noqa: F401 diff --git a/tna_utilities/flask/cache_control.py b/tna_utilities/flask/cache_control.py new file mode 100644 index 0000000..ad91bcd --- /dev/null +++ b/tna_utilities/flask/cache_control.py @@ -0,0 +1,57 @@ +from functools import wraps + +from flask import make_response + + +def do_not_cache(): + """ + Decorator to set Cache-Control headers to prevent caching of the response. + """ + + def decorator(f): + @wraps(f) + def decorated_function(*args, **kwargs): + response = make_response(f(*args, **kwargs)) + headers = response.headers + headers["Cache-Control"] = "no-store" + return response + + return decorated_function + + return decorator + + +def cacheable_duration(seconds: int = 3600): + """ + Decorator to set Cache-Control headers to allow caching of the response for a specified duration. + """ + + def decorator(f): + @wraps(f) + def decorated_function(*args, **kwargs): + response = make_response(f(*args, **kwargs)) + headers = response.headers + headers["Cache-Control"] = f"public, max-age={seconds}" + return response + + return decorated_function + + return decorator + + +def set_cache_control(instructions: str): + """ + Decorator to set Cache-Control headers with custom instructions provided as a string. + """ + + def decorator(f): + @wraps(f) + def decorated_function(*args, **kwargs): + response = make_response(f(*args, **kwargs)) + headers = response.headers + headers["Cache-Control"] = instructions + return response + + return decorated_function + + return decorator