Graph generation: refactor view - add an option in profile

2 files changed, 16 insertions, 8 deletions

diff --git a/archaeological_context_records/urls.py b/archaeological_context_records/urls.py
index fc917f542..48d8e54a3 100644
--- a/archaeological_context_records/urls.py
+++ b/archaeological_context_records/urls.py
@@ -53,7 +53,8 @@ urlpatterns = [
         views.show_contextrecord, name='show-historized-contextrecord'),
     url(r'revert-contextrecord/(?P<pk>.+)/(?P<date>.+)$',
         views.revert_contextrecord, name='revert-contextrecord'),
-    url(r'generate-relation-image-contextrecord/(?P<pk>.+)/',
+    url(r'generate-relation-image-contextrecord/(?P<pk>.+)/('
+        r'?P<type>(bellow)|(above))?',
         views.GenerateRelationImage.as_view(),
         name='generate-relation-image-contextrecord'),
     url(r'get-contextrecord/own/(?P<type>.+)?$', views.get_contextrecord,
diff --git a/archaeological_context_records/views.py b/archaeological_context_records/views.py
index 73c77fa48..fbe382623 100644
--- a/archaeological_context_records/views.py
+++ b/archaeological_context_records/views.py
@@ -164,18 +164,25 @@ def reset_wizards(request):
 
 
 class GenerateRelationImage(IshtarMixin, LoginRequiredMixin, RedirectView):
+    model = models.ContextRecord
+
     def get_redirect_url(self, *args, **kwargs):
-        return reverse('display-' + models.ContextRecord.SLUG,
-                       args=[self.context_record.pk])
+        return reverse('display-item',
+                       args=[self.model.SLUG, self.item.pk])
 
     def get(self, request, *args, **kwargs):
         try:
-            self.context_record = models.ContextRecord.objects.get(
-                pk=kwargs['pk']
-            )
-        except models.ContextRecord.DoesNotExist:
+            self.item = self.model.objects.get(pk=kwargs['pk'])
+        except self.model.DoesNotExist:
             raise Http404()
-        self.context_record.generate_relation_image()
+        render_above, render_bellow = True, True
+        render_type = kwargs.get("type", None)
+        if render_type == "above":
+            render_bellow = False
+        elif render_type == "bellow":
+            render_above = False
+        self.item.generate_relation_image(
+            render_bellow=render_bellow, render_above=render_above)
         return super(GenerateRelationImage, self).get(request, *args, **kwargs)