Repository: django-extensions/django-extensions (python)

Issue:
[feature request] Optional colorization for graph edges based on on_delete criteria
(Admin: Feel free to remove this line, the brackets in the title, and add the actual "feature request" tag)

I've using the visualization command to study the effects of `on_delete` in one of our projects. I generate a simple graph filtering using my apps and then I color the edges in the graph using Gimp. Red is for `CASCADE` and yellow for `PROTECTED`. This allow to study deletion criteria for relations with just a quick look.

I tought that patching the mainstream project was worth it, and here I am.

Do you think it will be a good addition to the model graph command? I don't have time right now for preparing the PR but if someone think this will be useful I will try to find time in the future.

Ofc, is someone wants to do it it will be nice. A quick inspections shown that the model graph class is a good place to start. There is already a simple method a that class that prepares de json for the edges. I think extending the signature of that definition and adding the 'fillcolor' param there will be enough.

Pull request description:
Flag for graph_models to color code relations based on on_delete. (#1604)
I implemented a new flag (`--color-code-deletions`) for the graph_models command. It adds custom colors for 'ForeignKey's and 'OneToOneField's edges, based on their on_delete setting. That allows a quick visual understanding of the effects of a model's deletion, as proposed on issue #1604.

This is my first contribution on this project, so, if _anything_ can/should be improved, let me know!

Closes #1604.

Patch:
diff --git a/django_extensions/management/commands/graph_models.py b/django_extensions/management/commands/graph_models.py
index c70e1696..1f2f651c 100644
--- a/django_extensions/management/commands/graph_models.py
+++ b/django_extensions/management/commands/graph_models.py
@@ -163,7 +163,7 @@ class Command(BaseCommand):
                 'action': 'store_true',
                 'default': False,
                 'dest': 'hide_edge_labels',
-                'help': 'Do not showrelations labels in the graph.',
+                'help': 'Do not show relations labels in the graph.',
             },
             '--arrow-shape': {
                 'action': 'store',
@@ -172,6 +172,12 @@ class Command(BaseCommand):
                 'choices': ['box', 'crow', 'curve', 'icurve', 'diamond', 'dot', 'inv', 'none', 'normal', 'tee', 'vee'],
                 'help': 'Arrow shape to use for relations. Default is dot. Available shapes: box, crow, curve, icurve, diamond, dot, inv, none, normal, tee, vee.',
             },
+            '--color-code-deletions': {
+                'action': 'store_true',
+                'default': False,
+                'dest': 'color_code_deletions',
+                'help': 'Color the relations according to their on_delete setting, where it it applicable. The colors are: red (CASCADE), orange (SET_NULL), green (SET_DEFAULT), yellow (SET), blue (PROTECT), grey (DO_NOTHING) and purple (RESTRICT).',
+            },
             '--rankdir': {
                 'action': 'store',
                 'default': 'TB',
diff --git a/django_extensions/management/modelviz.py b/django_extensions/management/modelviz.py
index 15552d07..fc49da6c 100644
--- a/django_extensions/management/modelviz.py
+++ b/django_extensions/management/modelviz.py
@@ -13,6 +13,7 @@ import os
 import re
 
 from django.apps import apps
+from django.db.models import deletion
 from django.db.models.fields.related import (
     ForeignKey, ManyToManyField, OneToOneField, RelatedField,
 )
@@ -44,6 +45,17 @@ __contributors__ = [
 ]
 
 
+ON_DELETE_COLORS = {
+    deletion.CASCADE: 'red',
+    deletion.PROTECT: 'blue',
+    deletion.SET_NULL: 'orange',
+    deletion.SET_DEFAULT: 'green',
+    deletion.SET: 'yellow',
+    deletion.DO_NOTHING: 'grey',
+    deletion.RESTRICT: 'purple',
+}
+
+
 def parse_file_or_list(arg):
     if not arg:
         return []
@@ -80,6 +92,7 @@ class ModelGraph:
         )
         self.hide_edge_labels = kwargs.get('hide_edge_labels', False)
         self.arrow_shape = kwargs.get("arrow_shape")
+        self.color_code_deletions = kwargs.get("color_code_deletions", False)
         if self.all_applications:
             self.app_labels = [app.label for app in apps.get_app_configs()]
         else:
@@ -152,7 +165,7 @@ class ModelGraph:
             'primary_key': field.primary_key,
         }
 
-    def add_relation(self, field, model, extras=""):
+    def add_relation(self, field, model, extras="", color=None):
         if self.verbose_names and field.verbose_name:
             label = force_str(field.verbose_name)
             if label.islower():
@@ -183,6 +196,9 @@ class ModelGraph:
         else:
             target_model = field.remote_field.model
 
+        if color:
+            extras = '[{}, color={}]'.format(extras[1:-1], color)
+
         _rel = self.get_relation_context(target_model, field, label, extras)
 
         if _rel not in model['relations'] and self.use_model(_rel['target']):
@@ -337,9 +353,15 @@ class ModelGraph:
             # excluding field redundant with inheritance relation
             # excluding fields inherited from abstract classes. they too show as local_fields
             return newmodel
+
+        color = None
+        if self.color_code_deletions and isinstance(field, (OneToOneField, ForeignKey)):
+            field_on_delete = getattr(field.remote_field, 'on_delete', None)
+            color = ON_DELETE_COLORS.get(field_on_delete)
+
         if isinstance(field, OneToOneField):
             relation = self.add_relation(
-                field, newmodel, '[arrowhead=none, arrowtail=none, dir=both]'
+                field, newmodel, '[arrowhead=none, arrowtail=none, dir=both]', color
             )
         elif isinstance(field, ForeignKey):
             relation = self.add_relation(
@@ -348,6 +370,7 @@ class ModelGraph:
                 '[arrowhead=none, arrowtail={}, dir=both]'.format(
                     self.arrow_shape
                 ),
+                color
             )
         else:
             relation = None
diff --git a/docs/graph_models.rst b/docs/graph_models.rst
index b1443f5b..3315d2e3 100644
--- a/docs/graph_models.rst
+++ b/docs/graph_models.rst
@@ -133,6 +133,11 @@ image by using the *graph_models* command::
   # Create a graph with 'normal' arrow shape for relations
   $ ./manage.py graph_models -a --arrow-shape normal -o my_project_sans_foo_bar.png
 
+::
+
+  # Create a graph with colored edges for relations with on_delete settings
+  $ ./manage.py graph_models -a --color-code-deletions -o my_project_colored.png
+
 ::
 
   # Create a graph with different layout direction,


Test patch:
diff --git a/tests/management/test_modelviz.py b/tests/management/test_modelviz.py
index db498b86..2ba86024 100644
--- a/tests/management/test_modelviz.py
+++ b/tests/management/test_modelviz.py
@@ -1,6 +1,6 @@
 # -*- coding: utf-8 -*-
 from django.test import SimpleTestCase
-from django_extensions.management.modelviz import generate_graph_data
+from django_extensions.management.modelviz import generate_graph_data, ON_DELETE_COLORS
 
 
 class ModelVizTests(SimpleTestCase):
@@ -25,3 +25,31 @@ class ModelVizTests(SimpleTestCase):
             'parent_cafe': u'Café latte',
         }
         self.assertEqual(expected, fields)
+
+    def test_on_delete_color_coding(self):
+        app_labels = ['django_extensions']
+        data = generate_graph_data(app_labels, color_code_deletions=True)
+
+        models = data['graphs'][0]['models']
+
+        for model in models:
+            relations = [x for x in model['relations'] if x['type'] in ('ForeignKey', 'OneToOneField')]
+
+            for relation in relations:
+                field = [x['field'] for x in model['fields'] if x['name'] == relation['name']][0]
+                on_delete = getattr(field.remote_field, 'on_delete', None)
+                expected_color = ON_DELETE_COLORS[on_delete]
+
+                self.assertIn('color={}'.format(expected_color), relation['arrows'])
+
+    def test_disabled_on_delete_color_coding(self):
+        app_labels = ['django_extensions']
+        data = generate_graph_data(app_labels)
+
+        models = data['graphs'][0]['models']
+
+        for model in models:
+            relations = [x for x in model['relations'] if x['type'] in ('ForeignKey', 'OneToOneField')]
+
+            for relation in relations:
+                self.assertNotIn('color=', relation['arrows'])
diff --git a/tests/testapp/models.py b/tests/testapp/models.py
index 141580da..e044babc 100644
--- a/tests/testapp/models.py
+++ b/tests/testapp/models.py
@@ -43,6 +43,20 @@ class Note(models.Model):
         app_label = 'django_extensions'
 
 
+class Neighborhood(models.Model):
+    name = models.CharField(max_length=50)
+
+    class Meta:
+        app_label = 'django_extensions'
+
+
+class Bank(models.Model):
+    name = models.CharField(max_length=50)
+
+    class Meta:
+        app_label = 'django_extensions'
+
+
 class Person(models.Model):
     name = models.ForeignKey(Name, on_delete=models.CASCADE)
     age = models.PositiveIntegerField()
@@ -54,6 +68,8 @@ class Person(models.Model):
         on_delete=models.CASCADE,
     )
     clubs = models.ManyToManyField(Club, through='testapp.Membership')
+    neighborhood = models.ForeignKey(Neighborhood, on_delete=models.SET_NULL, null=True)
+    current_bank = models.ForeignKey(Bank, on_delete=models.PROTECT, null=True)
 
     class Meta:
         app_label = 'django_extensions'