# -*- coding: utf-8 -*- import sys import json import os import tempfile from django.conf import settings from django.core.management.base import BaseCommand, CommandError from django.template import loader from django_extensions.management.modelviz import ModelGraph, generate_dot from django_extensions.management.utils import signalcommand try: import pygraphviz HAS_PYGRAPHVIZ = True except ImportError: HAS_PYGRAPHVIZ = False try: try: import pydotplus as pydot except ImportError: import pydot HAS_PYDOT = True except ImportError: HAS_PYDOT = False def retheme(graph_data, app_style={}): if isinstance(app_style, str): if os.path.exists(app_style): try: with open(app_style, "rt") as f: app_style = json.load(f) except Exception as e: print(f"Invalid app style file {app_style}") raise Exception(e) else: return graph_data for gc in graph_data["graphs"]: for g in gc: if "name" in g: for m in g["models"]: app_name = g["app_name"] if app_name in app_style: m["style"] = app_style[app_name] return graph_data class Command(BaseCommand): help = "Creates a GraphViz dot file for the specified app names." " You can pass multiple app names and they will all be combined into a" " single model. Output is usually directed to a dot file." can_import_settings = True def __init__(self, *args, **kwargs): """ Allow defaults for arguments to be set in settings.GRAPH_MODELS. Each argument in self.arguments is a dict where the key is the space-separated args and the value is our kwarg dict. The default from settings is keyed as the long arg name with '--' removed and any '-' replaced by '_'. For example, the default value for --disable-fields can be set in settings.GRAPH_MODELS['disable_fields']. """ self.arguments = { "--app-style": { "action": "store", "help": "Path to style json to configure the style per app", "dest": "app-style", "default": ".app-style.json", }, "--pygraphviz": { "action": "store_true", "default": False, "dest": "pygraphviz", "help": "Output graph data as image using PyGraphViz.", }, "--pydot": { "action": "store_true", "default": False, "dest": "pydot", "help": "Output graph data as image using PyDot(Plus).", }, "--dot": { "action": "store_true", "default": False, "dest": "dot", "help": ( "Output graph data as raw DOT (graph description language) " "text data." ), }, "--json": { "action": "store_true", "default": False, "dest": "json", "help": "Output graph data as JSON", }, "--disable-fields -d": { "action": "store_true", "default": False, "dest": "disable_fields", "help": "Do not show the class member fields", }, "--disable-abstract-fields": { "action": "store_true", "default": False, "dest": "disable_abstract_fields", "help": "Do not show the class member fields that were inherited", }, "--display-field-choices": { "action": "store_true", "default": False, "dest": "display_field_choices", "help": "Display choices instead of field type", }, "--group-models -g": { "action": "store_true", "default": False, "dest": "group_models", "help": "Group models together respective to their application", }, "--all-applications -a": { "action": "store_true", "default": False, "dest": "all_applications", "help": "Automatically include all applications from INSTALLED_APPS", }, "--output -o": { "action": "store", "dest": "outputfile", "help": ( "Render output file. Type of output dependend on file extensions. " "Use png or jpg to render graph to image." ), }, "--layout -l": { "action": "store", "dest": "layout", "default": "dot", "help": "Layout to be used by GraphViz for visualization. Layouts: " "circo dot fdp neato nop nop1 nop2 twopi", }, "--theme -t": { "action": "store", "dest": "theme", "default": "django2018", "help": "Theme to use. Supplied are 'original' and 'django2018'. " "You can create your own by creating dot templates in " "'django_extentions/graph_models/themename/' template directory.", }, "--verbose-names -n": { "action": "store_true", "default": False, "dest": "verbose_names", "help": "Use verbose_name of models and fields", }, "--language -L": { "action": "store", "dest": "language", "help": "Specify language used for verbose_name localization", }, "--exclude-columns -x": { "action": "store", "dest": "exclude_columns", "help": "Exclude specific column(s) from the graph. " "Can also load exclude list from file.", }, "--exclude-models -X": { "action": "store", "dest": "exclude_models", "help": "Exclude specific model(s) from the graph. Can also load " "exclude list from file. Wildcards (*) are allowed.", }, "--include-models -I": { "action": "store", "dest": "include_models", "help": "Restrict the graph to specified models. " "Wildcards (*) are allowed.", }, "--inheritance -e": { "action": "store_true", "default": True, "dest": "inheritance", "help": "Include inheritance arrows (default)", }, "--no-inheritance -E": { "action": "store_false", "default": False, "dest": "inheritance", "help": "Do not include inheritance arrows", }, "--hide-relations-from-fields -R": { "action": "store_false", "default": True, "dest": "relations_as_fields", "help": "Do not show relations as fields in the graph.", }, "--relation-fields-only": { "action": "store", "default": False, "dest": "relation_fields_only", "help": "Only display fields that are relevant for relations", }, "--disable-sort-fields -S": { "action": "store_false", "default": True, "dest": "sort_fields", "help": "Do not sort fields", }, "--hide-edge-labels": { "action": "store_true", "default": False, "dest": "hide_edge_labels", "help": "Do not show relations labels in the graph.", }, "--arrow-shape": { "action": "store", "default": "dot", "dest": "arrow_shape", "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 is 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", "dest": "rankdir", "choices": ["TB", "BT", "LR", "RL"], "help": "Set direction of graph layout. Supported directions: " "TB, LR, BT and RL. Corresponding to directed graphs drawn from " "top to bottom, from left to right, from bottom to top, and from " "right to left, respectively. Default is TB.", }, "--ordering": { "action": "store", "default": None, "dest": "ordering", "choices": ["in", "out"], "help": "Controls how the edges are arranged. Supported orderings: " '"in" (incoming relations first), "out" (outgoing relations first). ' "Default is None.", }, } defaults = getattr(settings, "GRAPH_MODELS", None) if defaults: for argument in self.arguments: arg_split = argument.split(" ") setting_opt = arg_split[0].lstrip("-").replace("-", "_") if setting_opt in defaults: self.arguments[argument]["default"] = defaults[setting_opt] super().__init__(*args, **kwargs) def add_arguments(self, parser): """Unpack self.arguments for parser.add_arguments.""" parser.add_argument("app_label", nargs="*") for argument in self.arguments: parser.add_argument(*argument.split(" "), **self.arguments[argument]) @signalcommand def handle(self, *args, **options): args = options["app_label"] if not args and not options["all_applications"]: default_app_labels = getattr(settings, "GRAPH_MODELS", {}).get("app_labels") if default_app_labels: args = default_app_labels else: raise CommandError("need one or more arguments for appname") # Determine output format based on options, file extension, and library # availability. outputfile = options.get("outputfile") or "" _, outputfile_ext = os.path.splitext(outputfile) outputfile_ext = outputfile_ext.lower() output_opts_names = ["pydot", "pygraphviz", "json", "dot"] output_opts = {k: v for k, v in options.items() if k in output_opts_names} output_opts_count = sum(output_opts.values()) if output_opts_count > 1: raise CommandError( "Only one of %s can be set." % ", ".join(["--%s" % opt for opt in output_opts_names]) ) if output_opts_count == 1: output = next(key for key, val in output_opts.items() if val) elif not outputfile: # When neither outputfile nor a output format option are set, # default to printing .dot format to stdout. Kept for backward # compatibility. output = "dot" elif outputfile_ext == ".dot": output = "dot" elif outputfile_ext == ".json": output = "json" elif HAS_PYGRAPHVIZ: output = "pygraphviz" elif HAS_PYDOT: output = "pydot" else: raise CommandError( "Neither pygraphviz nor pydotplus could be found to generate the image." " To generate text output, use the --json or --dot options." ) if options.get("rankdir") != "TB" and output not in [ "pydot", "pygraphviz", "dot", ]: raise CommandError( "--rankdir is not supported for the chosen output format" ) if options.get("ordering") and output not in ["pydot", "pygraphviz", "dot"]: raise CommandError( "--ordering is not supported for the chosen output format" ) # Consistency check: Abort if --pygraphviz or --pydot options are set # but no outputfile is specified. Before 2.1.4 this silently fell back # to printind .dot format to stdout. if output in ["pydot", "pygraphviz"] and not outputfile: raise CommandError( "An output file (--output) must be specified when --pydot or " "--pygraphviz are set." ) cli_options = " ".join(sys.argv[2:]) graph_models = ModelGraph(args, cli_options=cli_options, **options) graph_models.generate_graph_data() if output == "json": graph_data = graph_models.get_graph_data(as_json=True) return self.render_output_json(graph_data, outputfile) graph_data = graph_models.get_graph_data(as_json=False) theme = options["theme"] template_name = os.path.join( "django_extensions", "graph_models", theme, "digraph.dot" ) template = loader.get_template(template_name) graph_data = retheme(graph_data, app_style=options["app-style"]) dotdata = generate_dot(graph_data, template=template) if output == "pygraphviz": return self.render_output_pygraphviz(dotdata, **options) if output == "pydot": return self.render_output_pydot(dotdata, **options) self.print_output(dotdata, outputfile) def print_output(self, dotdata, output_file=None): """Write model data to file or stdout in DOT (text) format.""" if isinstance(dotdata, bytes): dotdata = dotdata.decode() if output_file: with open(output_file, "wt") as dot_output_f: dot_output_f.write(dotdata) else: self.stdout.write(dotdata) def render_output_json(self, graph_data, output_file=None): """Write model data to file or stdout in JSON format.""" if output_file: with open(output_file, "wt") as json_output_f: json.dump(graph_data, json_output_f) else: self.stdout.write(json.dumps(graph_data)) def render_output_pygraphviz(self, dotdata, **kwargs): """Render model data as image using pygraphviz.""" if not HAS_PYGRAPHVIZ: raise CommandError("You need to install pygraphviz python module") version = pygraphviz.__version__.rstrip("-svn") try: if tuple(int(v) for v in version.split(".")) < (0, 36): # HACK around old/broken AGraph before version 0.36 # (ubuntu ships with this old version) tmpfile = tempfile.NamedTemporaryFile() tmpfile.write(dotdata) tmpfile.seek(0) dotdata = tmpfile.name except ValueError: pass graph = pygraphviz.AGraph(dotdata) graph.layout(prog=kwargs["layout"]) graph.draw(kwargs["outputfile"]) def render_output_pydot(self, dotdata, **kwargs): """Render model data as image using pydot.""" if not HAS_PYDOT: raise CommandError("You need to install pydot python module") graph = pydot.graph_from_dot_data(dotdata) if not graph: raise CommandError("pydot returned an error") if isinstance(graph, (list, tuple)): if len(graph) > 1: sys.stderr.write( "Found more then one graph, rendering only the first one.\n" ) graph = graph[0] output_file = kwargs["outputfile"] formats = [ "bmp", "canon", "cmap", "cmapx", "cmapx_np", "dot", "dia", "emf", "em", "fplus", "eps", "fig", "gd", "gd2", "gif", "gv", "imap", "imap_np", "ismap", "jpe", "jpeg", "jpg", "metafile", "pdf", "pic", "plain", "plain-ext", "png", "pov", "ps", "ps2", "svg", "svgz", "tif", "tiff", "tk", "vml", "vmlz", "vrml", "wbmp", "webp", "xdot", ] ext = output_file[output_file.rfind(".") + 1 :] format_ = ext if ext in formats else "raw" graph.write(output_file, format=format_)