From 34e050fb8a14cee280bcc9de2e79dee94658f037 Mon Sep 17 00:00:00 2001
From: Jiri Kuncar <jiri.kuncar@cern.ch>
Date: Mon, 10 Nov 2014 18:16:42 +0100
Subject: [PATCH] sitemap: configurable decorators for endpoints

* Adds an option SITEMAP_VIEW_DECORATORS for specifying list of view
  decorators.  (closes #4)

Signed-off-by: Jiri Kuncar <jiri.kuncar@cern.ch>
---
 flask_sitemap/__init__.py | 22 +++++++++++++++++++++-
 flask_sitemap/config.py   |  9 ++++++++-
 tests/helpers.py          |  4 ++++
 tests/test_ext.py         | 34 +++++++++++++++++++++++++++++++++-
 4 files changed, 66 insertions(+), 3 deletions(-)

diff --git a/flask_sitemap/__init__.py b/flask_sitemap/__init__.py
index 6834630..5636d51 100644
--- a/flask_sitemap/__init__.py
+++ b/flask_sitemap/__init__.py
@@ -30,6 +30,9 @@
 
 from collections import Mapping
 from flask import current_app, request, Blueprint, render_template, url_for
+from functools import wraps
+from werkzeug.utils import import_string
+
 
 from . import config
 from .version import __version__
@@ -47,6 +50,7 @@ class Sitemap(object):
 
     def __init__(self, app=None):
         """Initialize login callback."""
+        self.decorators = []
         self.url_generators = [self._routes_without_params]
 
         if app is not None:
@@ -67,6 +71,12 @@ def init_app(self, app):
             if k.startswith('SITEMAP_'):
                 self.app.config.setdefault(k, getattr(config, k))
 
+        # Set decorators from configuration
+        for decorator in app.config.get('SITEMAP_VIEW_DECORATORS'):
+            if isinstance(decorator, string_types):
+                decorator = import_string(decorator)
+            self.decorators.append(decorator)
+
         # Create and register Blueprint
         if app.config.get('SITEMAP_BLUEPRINT'):
             # Add custom `template_folder`
@@ -76,13 +86,23 @@ def init_app(self, app):
             self.blueprint.add_url_rule(
                 app.config.get('SITEMAP_ENDPOINT_URL'),
                 'sitemap',
-                self.sitemap
+                self._decorate(self.sitemap)
             )
             app.register_blueprint(
                 self.blueprint,
                 url_prefix=app.config.get('SITEMAP_BLUEPRINT_URL_PREFIX')
             )
 
+    def _decorate(self, view):
+        """Decorate view with given decorators."""
+        @wraps(view)
+        def wrapper(*args, **kwargs):
+            new_view = view
+            for decorator in self.decorators:
+                new_view = decorator(new_view)
+            return new_view(*args, **kwargs)
+        return wrapper
+
     def sitemap(self):
         """Generate sitemap.xml."""
         return render_template('flask_sitemap/sitemap.xml',
diff --git a/flask_sitemap/config.py b/flask_sitemap/config.py
index 9fac8b6..334672c 100644
--- a/flask_sitemap/config.py
+++ b/flask_sitemap/config.py
@@ -41,16 +41,23 @@
 ------------------------
 
 Default: ``None``.
+
+SITEMAP_VIEW_DECORAROS
+----------------------
+
+Default: ``[]``.
 """
 
 SITEMAP_BLUEPRINT = 'flask_sitemap'
 
 SITEMAP_BLUEPRINT_URL_PREFIX = '/'
 
-SITEMAP_ENDPOINT_URL = '/sitemap.xml'
+SITEMAP_ENDPOINT_URL = 'sitemap.xml'
 
 SITEMAP_URL_SCHEME = 'http'
 
 SITEMAP_INCLUDE_RULES_WITHOUT_PARAMS = False
 
 SITEMAP_IGNORE_ENDPOINTS = None
+
+SITEMAP_VIEW_DECORATORS = []
diff --git a/tests/helpers.py b/tests/helpers.py
index 9da5ac1..2809fdc 100644
--- a/tests/helpers.py
+++ b/tests/helpers.py
@@ -22,3 +22,7 @@ def setUp(self):
         app.config['TESTING'] = True
         app.logger.disabled = True
         self.app = app
+
+
+def dummy_decorator(dummy):
+    return lambda *args, **kwargs: 'dummy'
diff --git a/tests/test_ext.py b/tests/test_ext.py
index 369dff0..426c145 100644
--- a/tests/test_ext.py
+++ b/tests/test_ext.py
@@ -10,14 +10,21 @@
 from __future__ import absolute_import
 
 import os
+import sys
 
 from contextlib import contextmanager
 from datetime import datetime
-from flask import request_started, request
+from flask import request_started, request, url_for
 from flask_sitemap import Sitemap, config as default_config
 
 from .helpers import FlaskTestCase
 
+# PY2/3 compatibility
+if sys.version_info[0] == 3:
+    b = lambda s: s.encode("latin-1")
+else:
+    b = lambda s: s
+
 
 class TestSitemap(FlaskTestCase):
 
@@ -155,3 +162,28 @@ def user():
         assert 'http://www.example.com/second' in results
         assert 'http://www.example.com/third' not in results
         assert 'http://www.example.com/fourth' not in results
+
+    def test_decorators_order(self):
+        def first(dummy):
+            return lambda *args, **kwargs: 'first'
+
+        def second(dummy):
+            return lambda *args, **kwargs: 'second'
+
+        def third(dummy):
+            return lambda *args, **kwargs: 'third'
+
+        self.app.config['SITEMAP_VIEW_DECORATORS'] = [
+            first, second, 'tests.helpers.dummy_decorator']
+        sitemap = Sitemap(app=self.app)
+
+        assert first in sitemap.decorators
+        assert second in sitemap.decorators
+
+        with self.app.test_client() as c:
+            assert b('dummy') == c.get('/sitemap.xml').data
+
+        sitemap.decorators.append(third)
+
+        with self.app.test_client() as c:
+            assert b('third') == c.get('/sitemap.xml').data