From 1e1077f891fcc7e621992049843a38c870e74be0 Mon Sep 17 00:00:00 2001
From: Karl Hobley <karlhobley10@gmail.com>
Date: Sat, 3 Oct 2015 12:22:34 +0100
Subject: [PATCH] Added WagtailAPIRouter

---
 wagtail/contrib/wagtailapi/router.py | 40 ++++++++++++++++++++++++++++
 wagtail/contrib/wagtailapi/urls.py   | 15 +++++------
 2 files changed, 47 insertions(+), 8 deletions(-)
 create mode 100644 wagtail/contrib/wagtailapi/router.py

diff --git a/wagtail/contrib/wagtailapi/router.py b/wagtail/contrib/wagtailapi/router.py
new file mode 100644
index 0000000000..13818cb577
--- /dev/null
+++ b/wagtail/contrib/wagtailapi/router.py
@@ -0,0 +1,40 @@
+import functools
+
+from django.conf.urls import url, include
+
+from wagtail.utils.urlpatterns import decorate_urlpatterns
+
+
+class WagtailAPIRouter(object):
+    def __init__(self, url_namespace):
+        self.url_namespace = url_namespace
+        self._endpoints = {}
+
+    def register_endpoint(self, name, class_):
+        self._endpoints[name] = class_
+
+    def wrap_view(self, func):
+        @functools.wraps(func)
+        def wrapped(request, *args, **kwargs):
+            request.wagtailapi_router = self
+            return func(request, *args, **kwargs)
+
+        return wrapped
+
+    def get_urlpatterns(self):
+        urlpatterns = []
+
+        for name, class_ in self._endpoints.items():
+            pattern = url(
+                r'^{}/'.format(name),
+                include(class_.get_urlpatterns(), namespace=name)
+            )
+            urlpatterns.append(pattern)
+
+        decorate_urlpatterns(urlpatterns, self.wrap_view)
+
+        return urlpatterns
+
+    @property
+    def urls(self):
+        return self.get_urlpatterns(), self.url_namespace, self.url_namespace
diff --git a/wagtail/contrib/wagtailapi/urls.py b/wagtail/contrib/wagtailapi/urls.py
index 492772f28e..da7f3f1673 100644
--- a/wagtail/contrib/wagtailapi/urls.py
+++ b/wagtail/contrib/wagtailapi/urls.py
@@ -1,17 +1,16 @@
 from __future__ import absolute_import
 
-from django.conf.urls import url, include
+from django.conf.urls import url
 
 from .endpoints import PagesAPIEndpoint, ImagesAPIEndpoint, DocumentsAPIEndpoint
+from .router import WagtailAPIRouter
 
 
-v1 = [
-    url(r'^pages/', include(PagesAPIEndpoint.get_urlpatterns(), namespace='pages')),
-    url(r'^images/', include(ImagesAPIEndpoint.get_urlpatterns(), namespace='images')),
-    url(r'^documents/', include(DocumentsAPIEndpoint.get_urlpatterns(), namespace='documents'))
-]
-
+v1 = WagtailAPIRouter('wagtailapi_v1')
+v1.register_endpoint('pages', PagesAPIEndpoint)
+v1.register_endpoint('images', ImagesAPIEndpoint)
+v1.register_endpoint('documents', DocumentsAPIEndpoint)
 
 urlpatterns = [
-    url(r'^v1/', include(v1, namespace='wagtailapi_v1')),
+    url(r'^v1/', v1.urls),
 ]