Migration Vis updates

.gitignore

@@ -133,3 +133,5 @@ dmypy.json
 
 # Custom
 migration-dep-tree*
+migration_vis/*
+!migration_vis/.gitkeep

apps/visualize/admin.py

@@ -6,3 +6,4 @@
 class MigrationSnapshot(admin.ModelAdmin):
     list_display = ["id", "output_format", "created_at", "modified_at"]
     list_filter = ["output_format"]
+    readonly_fields = ("graph_source", "output_file", "created_at", "modified_at")

apps/visualize/migrations/0001_initial.py

@@ -1,4 +1,4 @@
-# Generated by Django 3.2.7 on 2022-09-10 19:55
+# Generated by Django 3.2.7 on 2022-09-10 20:15
 
 from django.db import migrations, models
 
@@ -16,6 +16,7 @@ class Migration(migrations.Migration):
             fields=[
                 ('id', models.BigAutoField(auto_created=True, primary_key=True, serialize=False, verbose_name='ID')),
                 ('output_format', models.CharField(choices=[('bmp', 'BMP'), ('cgimage', 'CGIMAGE'), ('canon', 'CANON'), ('dot', 'DOT'), ('gv', 'GV'), ('xdot', 'XDOT'), ('xdot1.2', 'XDOT1.2'), ('xdot1.4', 'XDOT1.4'), ('eps', 'EPS'), ('exr', 'EXR'), ('fig', 'FIG'), ('gd', 'GD'), ('gd2', 'GD2'), ('gif', 'GIF'), ('gtk', 'GTK'), ('ico', 'ICO'), ('cmap', 'CMAP'), ('ismap', 'ISMAP'), ('imap', 'IMAP'), ('cmapx', 'CMAPX'), ('imap_np', 'IMAP_NP'), ('cmapx_np', 'CMAPX_NP'), ('jpg', 'JPG'), ('jpeg', 'JPEG'), ('jpe', 'JPE'), ('jp2', 'JP2'), ('json', 'JSON'), ('json0', 'JSON0'), ('dot_json', 'DOT_JSON'), ('xdot_json', 'XDOT_JSON'), ('pdf', 'PDF'), ('pic', 'PIC'), ('pct', 'PCT'), ('pict', 'PICT'), ('plain', 'PLAIN'), ('plain-ext', 'PLAIN-EXT'), ('png', 'PNG'), ('pov', 'POV'), ('ps2', 'PS2'), ('psd', 'PSD'), ('sgi', 'SGI'), ('svg', 'SVG'), ('svgz', 'SVGZ'), ('tga', 'TGA'), ('tif', 'TIF'), ('tiff', 'TIFF'), ('tk', 'TK'), ('vml', 'VML'), ('vmlz', 'VMLZ'), ('vrml', 'VRML'), ('wbmp', 'WBMP'), ('webp', 'WEBP'), ('xlib', 'XLIB'), ('x11', 'X11')], default='gv', max_length=10, verbose_name='Visualization File Output Format')),
+                ('graph_source', models.TextField(blank=True, null=True)),
                 ('output_file', models.FileField(blank=True, null=True, upload_to='migration_vis/')),
                 ('created_at', models.DateField(auto_now_add=True)),
                 ('modified_at', models.DateField(auto_now=True)),

apps/visualize/models.py

@@ -201,6 +201,7 @@ class MigrationSnapshot(models.Model):
         choices=FORMAT_CHOICES,
         default=GV,
     )
+    graph_source = models.TextField(blank=True, null=True)
     output_file = models.FileField(upload_to="migration_vis/", blank=True, null=True)
     created_at = models.DateField(auto_now_add=True)
     modified_at = models.DateField(auto_now=True)
@@ -210,23 +211,26 @@ def __str__(self):
 
     def save(self, *args, **kwargs):
         if not self.output_file:
-            file_loc = f"migrate_output"
-            file_name = f"{file_loc}.{self.output_format}"
-
-            try:
-                visualizer = MigrationVisualizer(
-                    None, filename=None, output_format=self.output_format
-                )
-                visualizer.render(save_loc=file_loc)
-
-                with open(file_name, "rb") as f:
-                    self.output_file.save(file_name, File(f))
-
-            finally:
-                os.remove(file_name)
-
+            self._record_snapshot()
         super().save(*args, **kwargs)
 
+    def _record_snapshot(self):
+        file_loc = f"migrate_output"
+        file_name = f"{file_loc}.{self.output_format}"
+
+        try:
+            visualizer = MigrationVisualizer(
+                None, filename=None, output_format=self.output_format
+            )
+            visualizer.render(save_loc=file_loc)
+            self.graph_source = str(visualizer.source)
+            with open(file_name, "rb") as f:
+                self.output_file.save(file_name, File(f))
+
+        finally:
+            os.remove(file_loc)
+            os.remove(file_name)
+
 
 """
 MigrationSnapshot.objects.create(output_format='gv')

apps/visualize/visualizer.py

@@ -197,16 +197,16 @@ def _add_dependencies(self, node):
             else:
                 self._add_edges(dep, self._get_tuple(node))
 
-    def render(self, save_loc=None):
+    def render(self, save_loc=None, view=False, **kwargs):
         if save_loc is None:
             save_loc = self.filename
 
         self._create_digraph()
         if save_loc:
-            self.picture.render(save_loc, view=False)
+            self.picture.render(save_loc, view=view, **kwargs)
         else:
             with NamedTemporaryFile() as temp:
-                self.picture.render(temp.name, view=True)
+                self.picture.render(temp.name, view=True, **kwargs)
 
         # temp = NamedTemporaryFile()
         # self.picture.render(temp.name, view=False)