Flask vary decorators

.github/CODEOWNERS

@@ -0,0 +1 @@
+*       @ahosgood

CHANGELOG.md

@@ -9,6 +9,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added
 
+- Added `vary_by_cookies()` and `vary_by_headers()` decorators for Flask
+
 ### Changed
 
 ### Deprecated

docs/flask.md

@@ -4,7 +4,7 @@
 
 > Added in `v1.5.0`.
 
-A set of decorators to manage the `Cache-Control` header of a route.
+A set of decorators to manage the `Cache-Control` response header of a route.
 
 ### Examples
 
@@ -34,6 +34,31 @@ def private_cache():
     return "Cache me in private caches for up to 2 minutes"
 ```
 
+### Vary
+
+> Added in `v1.6.0`.
+
+Decorators to help set the `Vary` response header of a route.
+
+### Examples
+
+```python
+from flask import Flask
+from tna_utilities.flask import vary_by_cookies, vary_by_headers
+
+app = Flask(__name__)
+
+@app.route("/vary-cookies/")
+@vary_by_cookies()
+def response_varies_based_on_cookies():
+    return "Send different cookies to get a different response"
+
+@app.route("/vary-accept/")
+@vary_by_headers("Accept")
+def response_varies_based_on_accept_header():
+    return "Send a request with a different Accept header to get a different response"
+```
+
 ## `Talisman`
 
 > Added in `v1.4.0`.

tests/test_flask_cache_control.py

@@ -2,7 +2,13 @@
 
 from flask import Flask
 
-from tna_utilities.flask import cacheable_duration, do_not_cache, set_cache_control
+from tna_utilities.flask import (
+    cacheable_duration,
+    do_not_cache,
+    set_cache_control,
+    vary_by_cookies,
+    vary_by_headers,
+)
 
 
 class TestFlaskCacheControl(unittest.TestCase):
@@ -79,3 +85,33 @@ def index():
             rv.headers["Cache-Control"],
             "private, max-age=120",
         )
+
+    def test_vary_by_cookies_route(self):
+        @self.app.route("/")
+        @vary_by_cookies()
+        def index():
+            return "OK"
+
+        rv = self.test_client.get("/")
+
+        self.assertEqual(rv.status_code, 200)
+        self.assertIn("Vary", rv.headers)
+        self.assertEqual(
+            rv.headers["Vary"],
+            "Cookie",
+        )
+
+    def test_vary_by_headers_route(self):
+        @self.app.route("/")
+        @vary_by_headers("Accept-Encoding, User-Agent")
+        def index():
+            return "OK"
+
+        rv = self.test_client.get("/")
+
+        self.assertEqual(rv.status_code, 200)
+        self.assertIn("Vary", rv.headers)
+        self.assertEqual(
+            rv.headers["Vary"],
+            "Accept-Encoding, User-Agent",
+        )

tna_utilities/flask/__init__.py

@@ -2,5 +2,7 @@
     cacheable_duration,
     do_not_cache,
     set_cache_control,
+    vary_by_cookies,
+    vary_by_headers,
 )
 from tna_utilities.flask.talisman import Talisman

tna_utilities/flask/cache_control.py

@@ -8,48 +8,53 @@ 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
+    return set_cache_control("no-store")
 
 
 def cacheable_duration(seconds: int = 3600):
     """
     Decorator to set Cache-Control headers to allow caching of the response for a specified duration.
     """
 
+    return set_cache_control(f"public, max-age={seconds}")
+
+
+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"] = f"public, max-age={seconds}"
+            headers["Cache-Control"] = instructions
             return response
 
         return decorated_function
 
     return decorator
 
 
-def set_cache_control(instructions: str):
+def vary_by_cookies():
     """
-    Decorator to set Cache-Control headers with custom instructions provided as a string.
+    Decorator to set Vary headers to indicate that the response varies based on cookies.
+    """
+
+    return vary_by_headers("Cookie")
+
+
+def vary_by_headers(headers: str):
+    """
+    Decorator to set Vary headers to indicate that the response varies based on specified headers.
     """
 
     def decorator(f):
         @wraps(f)
         def decorated_function(*args, **kwargs):
             response = make_response(f(*args, **kwargs))
-            headers = response.headers
-            headers["Cache-Control"] = instructions
+            response.headers["Vary"] = headers
             return response
 
         return decorated_function