Graphs: save results, options and tests

- can generate only above and only bellow
- save dot files
- option for no color for current
- tests

2 files changed, 137 insertions, 61 deletions

diff --git a/ishtar_common/models.py b/ishtar_common/models.py
index 5e7f6e9b3..e8b6fc518 100644
--- a/ishtar_common/models.py
+++ b/ishtar_common/models.py
@@ -1205,6 +1205,38 @@ class RelationItem(models.Model):
         _("Generated relation image (SVG)"), null=True, blank=True,
         upload_to=get_image_path, help_text=max_size_help()
     )
+    relation_bitmap_image = models.FileField(
+        _("Generated relation image (PNG)"), null=True, blank=True,
+        upload_to=get_image_path, help_text=max_size_help()
+    )
+    relation_dot = models.FileField(
+        _("Generated relation image (DOT)"), null=True, blank=True,
+        upload_to=get_image_path, help_text=max_size_help()
+    )
+    relation_image_above = models.FileField(
+        _("Generated above relation image (SVG)"), null=True, blank=True,
+        upload_to=get_image_path, help_text=max_size_help()
+    )
+    relation_dot_above = models.FileField(
+        _("Generated above relation image (DOT)"), null=True, blank=True,
+        upload_to=get_image_path, help_text=max_size_help()
+    )
+    relation_bitmap_image_above = models.FileField(
+        _("Generated above relation image (PNG)"), null=True, blank=True,
+        upload_to=get_image_path, help_text=max_size_help()
+    )
+    relation_image_bellow = models.FileField(
+        _("Generated bellow relation image (SVG)"), null=True, blank=True,
+        upload_to=get_image_path, help_text=max_size_help()
+    )
+    relation_dot_bellow = models.FileField(
+        _("Generated bellow relation image (DOT)"), null=True, blank=True,
+        upload_to=get_image_path, help_text=max_size_help()
+    )
+    relation_bitmap_image_bellow = models.FileField(
+        _("Generated bellow relation image (PNG)"), null=True, blank=True,
+        upload_to=get_image_path, help_text=max_size_help()
+    )
 
     class Meta:
         abstract = True
diff --git a/ishtar_common/utils.py b/ishtar_common/utils.py
index 83bc56279..4f8dce853 100644
--- a/ishtar_common/utils.py
+++ b/ishtar_common/utils.py
@@ -17,6 +17,7 @@
 
 # See the file COPYING for details.
 
+from cairosvg import svg2png
 from csv import QUOTE_ALL
 import datetime
 from functools import wraps
@@ -1089,7 +1090,8 @@ def create_default_areas(models=None, verbose=False):
 
 
 def get_relations_for_graph(rel_model, obj_pk, above_relations=None,
-                            equal_relations=None, treated=None, styles=None):
+                            equal_relations=None, treated=None, styles=None,
+                            render_above=True, render_bellow=True):
     """
     Get all above and equal relations of an object (get all child and parent
     relations)
@@ -1099,6 +1101,8 @@ def get_relations_for_graph(rel_model, obj_pk, above_relations=None,
     :param equal_relations: list of current equal_relations
     :param treated: treated relation list to prevent circular call
     :param styles: current styles
+    :param render_above: render relation above the current object
+    :param render_bellow: render relation bellow the current object
     :return: above and equal relations list (each containing lists of two
     members)
     """
@@ -1114,63 +1118,81 @@ def get_relations_for_graph(rel_model, obj_pk, above_relations=None,
         return above_relations, equal_relations, styles
 
     treated.append(obj_pk)
-    q = rel_model.objects.filter(
-        left_record_id=obj_pk,
-        relation_type__logical_relation__isnull=False
-    ).values('right_record_id', 'relation_type__logical_relation')
-    if not q.count():
-        return [], []
-    for relation in q.all():
-        logical_relation = relation['relation_type__logical_relation']
-        right_record = relation['right_record_id']
-        if not logical_relation:
-            continue
-        elif logical_relation == 'above'and \
-                (obj_pk, right_record) not in above_relations:
-            above_relations.append((obj_pk, right_record))
-        elif logical_relation == 'bellow' and \
-                (right_record, obj_pk) not in above_relations:
-            above_relations.append((right_record, obj_pk))
-        elif logical_relation == 'equal' and \
-                (right_record, obj_pk) not in equal_relations and \
-                (obj_pk, right_record) not in equal_relations:
-            equal_relations.append((obj_pk, right_record))
-        else:
-            continue
-        ar, er, substyles = get_relations_for_graph(
-            rel_model, right_record, above_relations, equal_relations, treated,
-            styles
-        )
-        styles.update(substyles)
-        error_style = "color=red"
-        for r in ar:
-            if r not in above_relations:
-                above_relations.append(r)
-            reverse_rel = tuple(reversed(r))
-            if reverse_rel in above_relations:
-                # circular
-                if r not in styles:
-                    styles[r] = []
-                if reverse_rel not in styles:
-                    styles[reverse_rel] = []
-
-                if error_style not in styles[r]:
-                    styles[r].append(error_style)
-                if error_style not in styles[reverse_rel]:
-                    styles[reverse_rel].append(error_style)
-            if r[0] == r[1]:
-                # same entity
-                if r not in styles:
-                    styles[r] = []
-                if error_style not in styles[r]:
-                    styles[r].append("color=red")
-        for r in er:
-            if r not in equal_relations:
-                equal_relations.append(r)
+
+    for q, inverse in (
+            (rel_model.objects.filter(
+                left_record_id=obj_pk,
+                relation_type__logical_relation__isnull=False), False),
+            (rel_model.objects.filter(
+                right_record_id=obj_pk,
+                relation_type__logical_relation__isnull=False), True)):
+        q = q.values("left_record_id",'right_record_id',
+                     'relation_type__logical_relation')
+        get_above, get_bellow = render_above, render_bellow
+        if inverse and (not render_above or not render_bellow):
+            get_above, get_bellow = not render_above, not render_bellow
+
+        for relation in q.all():
+            logical_relation = relation['relation_type__logical_relation']
+            left_record = relation['left_record_id']
+            right_record = relation['right_record_id']
+            if not logical_relation:
+                continue
+            elif get_bellow and logical_relation == 'above' and \
+                    (left_record, right_record) not in above_relations:
+                above_relations.append((left_record, right_record))
+            elif get_above and logical_relation == 'bellow' and \
+                    (right_record, left_record) not in above_relations:
+                above_relations.append((right_record, left_record))
+            elif logical_relation == 'equal' and \
+                    (right_record, left_record) not in equal_relations and \
+                    (left_record, right_record) not in equal_relations:
+                equal_relations.append((left_record, right_record))
+            else:
+                continue
+
+            if right_record == obj_pk:
+                other_record = left_record
+            else:
+                other_record = right_record
+
+            ar, er, substyles = get_relations_for_graph(
+                rel_model, other_record, above_relations, equal_relations,
+                treated, styles, render_above=render_above,
+                render_bellow=render_bellow
+            )
+            styles.update(substyles)
+            error_style = "color=red"
+            for r in ar:
+                if r not in above_relations:
+                    above_relations.append(r)
+                reverse_rel = tuple(reversed(r))
+                if reverse_rel in above_relations:
+                    # circular
+                    if r not in styles:
+                        styles[r] = []
+                    if reverse_rel not in styles:
+                        styles[reverse_rel] = []
+
+                    if error_style not in styles[r]:
+                        styles[r].append(error_style)
+                    if error_style not in styles[reverse_rel]:
+                        styles[reverse_rel].append(error_style)
+                if r[0] == r[1]:
+                    # same entity
+                    if r not in styles:
+                        styles[r] = []
+                    if error_style not in styles[r]:
+                        styles[r].append("color=red")
+            for r in er:
+                if r not in equal_relations:
+                    equal_relations.append(r)
     return above_relations, equal_relations, styles
 
 
-def generate_relation_graph(obj, debug=False):
+def generate_relation_graph(obj, highlight_current=True,
+                            render_above=True, render_bellow=True,
+                            debug=False):
     if not settings.DOT_BINARY:
         return
 
@@ -1179,8 +1201,8 @@ def generate_relation_graph(obj, debug=False):
 
     # get relations
     above_relations, equal_relations, styles = get_relations_for_graph(
-        rel_model, obj.pk)
-    print(styles)
+        rel_model, obj.pk, render_above=render_above,
+        render_bellow=render_bellow)
     if not above_relations and not equal_relations:
         obj.relation_image = None
         obj.save()
@@ -1201,14 +1223,14 @@ def generate_relation_graph(obj, debug=False):
                 described.append(left_pk)
                 left = model.objects.get(pk=left_pk)
                 style = 'label="{}"'.format(left.relation_label)
-                if left.pk == obj.pk:
+                if left.pk == obj.pk and highlight_current:
                     style += ',style=filled,fillcolor="#C6C0C0"'
                 dot_str += u'item{}[{}];\n'.format(left.pk, style)
             if right_pk not in described:
                 described.append(right_pk)
                 right = model.objects.get(pk=right_pk)
                 style = 'label="{}"'.format(right.relation_label)
-                if right.pk == obj.pk:
+                if right.pk == obj.pk and highlight_current:
                     style += ',style=filled,fillcolor="#C6C0C0"'
                 dot_str += u'item{}[{}];\n'.format(right.pk, style)
             if not directed:  # on the same level
@@ -1226,6 +1248,18 @@ def generate_relation_graph(obj, debug=False):
     with open(dot_name, 'w') as dot_file:
         dot_file.write(dot_str)
 
+    if not render_above:
+        suffix = "_above"
+    elif not render_bellow:
+        suffix = "_bellow"
+    else:
+        suffix = ""
+
+    with open(dot_name, "r") as dot_file:
+        django_file = File(dot_file)
+        attr = "relation_dot" + suffix
+        getattr(obj, attr).save("relations.dot", django_file, save=True)
+
     # execute dot program
     args = (settings.DOT_BINARY, "-Tsvg", dot_name)
 
@@ -1234,10 +1268,20 @@ def generate_relation_graph(obj, debug=False):
     with open(svg_tmp_name, "w") as svg_file:
         popen = subprocess.Popen(args, stdout=svg_file)
         popen.wait()
-
     with open(svg_tmp_name, "r") as svg_file:
         django_file = File(svg_file)
-        obj.relation_image.save("relations.svg", django_file, save=True)
+        attr = "relation_image" + suffix
+        getattr(obj, attr).save("relations.svg", django_file, save=True)
+
+    png_name = tempdir + os.path.sep + "relations.png"
+
+    with open(png_name, "wb") as png_file:
+        svg2png(open(svg_tmp_name, 'rb').read(),
+                write_to=png_file)
+    with open(png_name, "rb") as png_file:
+        django_file = File(png_file)
+        attr = "relation_bitmap_image" + suffix
+        getattr(obj, attr).save("relations.png", django_file, save=True)
 
     if debug:
         print(u"DOT file: {}. Tmp SVG file: {}.".format(dot_name, svg_tmp_name))