Flag for graph_models to color code relations based on on_delete. (#1604

) (#1664)

* Add flag to color code relations with on_delete on the model graph render

* Add more model relations on the test app

* Create test for the model graph relations color coding

* Add documentation example for the --color-code-deletions flag

* Removed django 3.1 support

* Removed unused import

Co-authored-by: Camilo Nova <camilo.nova@gmail.com>

django_extensions/management/commands/graph_models.py

@@ -163,7 +163,7 @@ def __init__(self, *args, **kwargs):
                 '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 @@ def __init__(self, *args, **kwargs):
                 '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',

django_extensions/management/modelviz.py

@@ -13,6 +13,7 @@
 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 @@
 ]
 
 
+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 @@ def __init__(self, app_labels, **kwargs):
         )
         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 @@ def add_attributes(self, field, abstract_fields):
             '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 @@ def add_relation(self, field, model, extras=""):
         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 @@ def process_local_fields(self, field, model, abstract_fields):
             # 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 @@ def process_local_fields(self, field, model, abstract_fields):
                 '[arrowhead=none, arrowtail={}, dir=both]'.format(
                     self.arrow_shape
                 ),
+                color
             )
         else:
             relation = None

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,

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 @@ def test_render_unicode_field_label(self):
             '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'])

tests/testapp/models.py

@@ -43,6 +43,20 @@ class Meta:
         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'