527 lines
19 KiB
Python
527 lines
19 KiB
Python
# -*- coding: utf-8 -*-
|
|
"""
|
|
modelviz.py - DOT file generator for Django Models
|
|
|
|
Based on:
|
|
Django model to DOT (Graphviz) converter
|
|
by Antonio Cavedoni <antonio@cavedoni.org>
|
|
Adapted to be used with django-extensions
|
|
"""
|
|
|
|
import datetime
|
|
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,
|
|
)
|
|
from django.db.models.fields.reverse_related import (
|
|
OneToOneRel,
|
|
ManyToOneRel,
|
|
)
|
|
from django.contrib.contenttypes.fields import GenericRelation
|
|
from django.template import Context, Template, loader
|
|
from django.utils.encoding import force_str
|
|
from django.utils.safestring import mark_safe
|
|
from django.utils.translation import activate as activate_language
|
|
|
|
|
|
__version__ = "1.1"
|
|
__license__ = "Python"
|
|
__author__ = ("Bas van Oostveen <v.oostveen@gmail.com>",)
|
|
__contributors__ = [
|
|
"Antonio Cavedoni <https://cavedoni.com/>Stefano J. Attardi <https://attardi.org/>",
|
|
"limodou",
|
|
"Carlo C8E Miron",
|
|
"Andre Campos <cahenan@gmail.com>",
|
|
"Justin Findlay <jfindlay@gmail.com>",
|
|
"Alexander Houben <alexander@houben.ch>",
|
|
"Joern Hees <gitdev@joernhees.de>",
|
|
"Kevin Cherepski <cherepski@gmail.com>",
|
|
"Jose Tomas Tocino <theom3ga@gmail.com>",
|
|
"Adam Dobrawy <naczelnik@jawnosc.tk>",
|
|
"Mikkel Munch Mortensen <https://www.detfalskested.dk/>",
|
|
"Andrzej Bistram <andrzej.bistram@gmail.com>",
|
|
"Daniel Lipsitt <danlipsitt@gmail.com>",
|
|
"Florian Anceau <flow.gunso@gmail.com>",
|
|
]
|
|
|
|
|
|
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 []
|
|
if isinstance(arg, (list, tuple, set)):
|
|
return arg
|
|
if "," not in arg and os.path.isfile(arg):
|
|
return [e.strip() for e in open(arg).readlines()]
|
|
return [e.strip() for e in arg.split(",")]
|
|
|
|
|
|
class ModelGraph:
|
|
def __init__(self, app_labels, **kwargs):
|
|
self.graphs = []
|
|
self.cli_options = kwargs.get("cli_options", None)
|
|
self.disable_fields = kwargs.get("disable_fields", False)
|
|
self.disable_abstract_fields = kwargs.get("disable_abstract_fields", False)
|
|
self.include_models = parse_file_or_list(kwargs.get("include_models", ""))
|
|
self.all_applications = kwargs.get("all_applications", False)
|
|
self.use_subgraph = kwargs.get("group_models", False)
|
|
self.verbose_names = kwargs.get("verbose_names", False)
|
|
self.inheritance = kwargs.get("inheritance", True)
|
|
self.relations_as_fields = kwargs.get("relations_as_fields", True)
|
|
self.relation_fields_only = kwargs.get("relation_fields_only", False)
|
|
self.sort_fields = kwargs.get("sort_fields", True)
|
|
self.language = kwargs.get("language", None)
|
|
if self.language is not None:
|
|
activate_language(self.language)
|
|
self.exclude_columns = parse_file_or_list(kwargs.get("exclude_columns", ""))
|
|
self.exclude_models = parse_file_or_list(kwargs.get("exclude_models", ""))
|
|
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:
|
|
self.app_labels = app_labels
|
|
self.rankdir = kwargs.get("rankdir")
|
|
self.display_field_choices = kwargs.get("display_field_choices", False)
|
|
self.ordering = kwargs.get("ordering")
|
|
|
|
def generate_graph_data(self):
|
|
self.process_apps()
|
|
|
|
nodes = []
|
|
for graph in self.graphs:
|
|
nodes.extend([e["name"] for e in graph["models"]])
|
|
|
|
for graph in self.graphs:
|
|
for model in graph["models"]:
|
|
for relation in model["relations"]:
|
|
if relation is not None:
|
|
if relation["target"] in nodes:
|
|
relation["needs_node"] = False
|
|
|
|
def get_graph_data(self, as_json=False):
|
|
now = datetime.datetime.now()
|
|
graph_data = {
|
|
"created_at": now.strftime("%Y-%m-%d %H:%M"),
|
|
"cli_options": self.cli_options,
|
|
"disable_fields": self.disable_fields,
|
|
"disable_abstract_fields": self.disable_abstract_fields,
|
|
"display_field_choices": self.display_field_choices,
|
|
"use_subgraph": self.use_subgraph,
|
|
"rankdir": self.rankdir,
|
|
"ordering": self.ordering,
|
|
}
|
|
|
|
if as_json:
|
|
# We need to remove the model and field class
|
|
# because it is not JSON serializable
|
|
graphs = [context.flatten() for context in self.graphs]
|
|
for context in graphs:
|
|
for model_data in context["models"]:
|
|
model_data.pop("model")
|
|
for field_data in model_data["fields"]:
|
|
field_data.pop("field")
|
|
graph_data["graphs"] = graphs
|
|
else:
|
|
graph_data["graphs"] = self.graphs
|
|
|
|
return graph_data
|
|
|
|
def add_attributes(self, field, abstract_fields):
|
|
if self.verbose_names and field.verbose_name:
|
|
label = force_str(field.verbose_name)
|
|
if label.islower():
|
|
label = label.capitalize()
|
|
else:
|
|
label = field.name
|
|
|
|
t = type(field).__name__
|
|
if isinstance(field, (OneToOneField, ForeignKey)):
|
|
t += " ({0})".format(field.remote_field.field_name)
|
|
if self.display_field_choices and field.choices is not None:
|
|
choices = {c for c, _ in field.choices}
|
|
t = str(choices)
|
|
# TODO: ManyToManyField, GenericRelation
|
|
|
|
return {
|
|
"field": field,
|
|
"name": field.name,
|
|
"label": label,
|
|
"type": t,
|
|
"blank": field.blank,
|
|
"abstract": any(
|
|
field.creation_counter == abstract_field.creation_counter
|
|
for abstract_field in abstract_fields
|
|
),
|
|
"relation": isinstance(field, RelatedField),
|
|
"primary_key": field.primary_key,
|
|
}
|
|
|
|
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():
|
|
label = label.capitalize()
|
|
else:
|
|
label = field.name
|
|
|
|
# show related field name
|
|
if hasattr(field, "related_query_name"):
|
|
related_query_name = field.related_query_name()
|
|
if self.verbose_names and related_query_name.islower():
|
|
related_query_name = related_query_name.replace("_", " ").capitalize()
|
|
label = "{} ({})".format(label, force_str(related_query_name))
|
|
if self.hide_edge_labels:
|
|
label = ""
|
|
|
|
# handle self-relationships and lazy-relationships
|
|
if isinstance(field.remote_field.model, str):
|
|
if field.remote_field.model == "self":
|
|
target_model = field.model
|
|
else:
|
|
if "." in field.remote_field.model:
|
|
app_label, model_name = field.remote_field.model.split(".", 1)
|
|
else:
|
|
app_label = field.model._meta.app_label
|
|
model_name = field.remote_field.model
|
|
target_model = apps.get_model(app_label, model_name)
|
|
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"]):
|
|
return _rel
|
|
|
|
def get_abstract_models(self, appmodels):
|
|
abstract_models = []
|
|
for appmodel in appmodels:
|
|
abstract_models += [
|
|
abstract_model
|
|
for abstract_model in appmodel.__bases__
|
|
if hasattr(abstract_model, "_meta") and abstract_model._meta.abstract
|
|
]
|
|
abstract_models = list(set(abstract_models)) # remove duplicates
|
|
return abstract_models
|
|
|
|
def get_app_context(self, app):
|
|
return Context(
|
|
{
|
|
"name": '"%s"' % app.name,
|
|
"app_name": "%s" % app.name,
|
|
"cluster_app_name": "cluster_%s" % app.name.replace(".", "_"),
|
|
"models": [],
|
|
}
|
|
)
|
|
|
|
def get_appmodel_attributes(self, appmodel):
|
|
if self.relations_as_fields:
|
|
attributes = [field for field in appmodel._meta.local_fields]
|
|
else:
|
|
# Find all the 'real' attributes. Relations are depicted as graph edges
|
|
# instead of attributes
|
|
attributes = [
|
|
field
|
|
for field in appmodel._meta.local_fields
|
|
if not isinstance(field, RelatedField)
|
|
]
|
|
return attributes
|
|
|
|
def get_appmodel_abstracts(self, appmodel):
|
|
return [
|
|
abstract_model.__name__
|
|
for abstract_model in appmodel.__bases__
|
|
if hasattr(abstract_model, "_meta") and abstract_model._meta.abstract
|
|
]
|
|
|
|
def get_appmodel_context(self, appmodel, appmodel_abstracts):
|
|
context = {
|
|
"model": appmodel,
|
|
"app_name": appmodel.__module__.replace(".", "_"),
|
|
"name": appmodel.__name__,
|
|
"abstracts": appmodel_abstracts,
|
|
"fields": [],
|
|
"relations": [],
|
|
}
|
|
|
|
if self.verbose_names and appmodel._meta.verbose_name:
|
|
context["label"] = force_str(appmodel._meta.verbose_name)
|
|
else:
|
|
context["label"] = context["name"]
|
|
|
|
return context
|
|
|
|
def get_bases_abstract_fields(self, c):
|
|
_abstract_fields = []
|
|
for e in c.__bases__:
|
|
if hasattr(e, "_meta") and e._meta.abstract:
|
|
_abstract_fields.extend(e._meta.fields)
|
|
_abstract_fields.extend(self.get_bases_abstract_fields(e))
|
|
return _abstract_fields
|
|
|
|
def get_inheritance_context(self, appmodel, parent):
|
|
label = "multi-table"
|
|
if parent._meta.abstract:
|
|
label = "abstract"
|
|
if appmodel._meta.proxy:
|
|
label = "proxy"
|
|
label += r"\ninheritance"
|
|
if self.hide_edge_labels:
|
|
label = ""
|
|
return {
|
|
"target_app": parent.__module__.replace(".", "_"),
|
|
"target": parent.__name__,
|
|
"type": "inheritance",
|
|
"name": "inheritance",
|
|
"label": label,
|
|
"arrows": "[arrowhead=empty, arrowtail=none, dir=both]",
|
|
"needs_node": True,
|
|
}
|
|
|
|
def get_models(self, app):
|
|
appmodels = list(app.get_models())
|
|
return appmodels
|
|
|
|
def get_relation_context(self, target_model, field, label, extras):
|
|
return {
|
|
"target_app": target_model.__module__.replace(".", "_"),
|
|
"target": target_model.__name__,
|
|
"type": type(field).__name__,
|
|
"name": field.name,
|
|
"label": label,
|
|
"arrows": extras,
|
|
"needs_node": True,
|
|
}
|
|
|
|
def process_attributes(self, field, model, pk, abstract_fields):
|
|
newmodel = model.copy()
|
|
if self.skip_field(field) or pk and field == pk:
|
|
return newmodel
|
|
newmodel["fields"].append(self.add_attributes(field, abstract_fields))
|
|
return newmodel
|
|
|
|
def process_apps(self):
|
|
for app_label in self.app_labels:
|
|
app = apps.get_app_config(app_label)
|
|
if not app:
|
|
continue
|
|
app_graph = self.get_app_context(app)
|
|
app_models = self.get_models(app)
|
|
abstract_models = self.get_abstract_models(app_models)
|
|
app_models = abstract_models + app_models
|
|
|
|
for appmodel in app_models:
|
|
if not self.use_model(appmodel._meta.object_name):
|
|
continue
|
|
appmodel_abstracts = self.get_appmodel_abstracts(appmodel)
|
|
abstract_fields = self.get_bases_abstract_fields(appmodel)
|
|
model = self.get_appmodel_context(appmodel, appmodel_abstracts)
|
|
attributes = self.get_appmodel_attributes(appmodel)
|
|
|
|
# find primary key and print it first
|
|
# ignoring implicit id if other pk exists
|
|
pk = appmodel._meta.pk
|
|
if pk and not appmodel._meta.abstract and pk in attributes:
|
|
model["fields"].append(self.add_attributes(pk, abstract_fields))
|
|
|
|
for field in attributes:
|
|
model = self.process_attributes(field, model, pk, abstract_fields)
|
|
|
|
if self.sort_fields:
|
|
model = self.sort_model_fields(model)
|
|
|
|
for field in appmodel._meta.local_fields:
|
|
model = self.process_local_fields(field, model, abstract_fields)
|
|
|
|
for field in appmodel._meta.local_many_to_many:
|
|
model = self.process_local_many_to_many(field, model)
|
|
|
|
if self.inheritance:
|
|
# add inheritance arrows
|
|
for parent in appmodel.__bases__:
|
|
model = self.process_parent(parent, appmodel, model)
|
|
|
|
app_graph["models"].append(model)
|
|
if app_graph["models"]:
|
|
self.graphs.append(app_graph)
|
|
|
|
def process_local_fields(self, field, model, abstract_fields):
|
|
newmodel = model.copy()
|
|
if (
|
|
field.attname.endswith("_ptr_id")
|
|
or field in abstract_fields
|
|
or self.skip_field(field)
|
|
):
|
|
# 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]", color
|
|
)
|
|
elif isinstance(field, ForeignKey):
|
|
relation = self.add_relation(
|
|
field,
|
|
newmodel,
|
|
"[arrowhead=none, arrowtail={}, dir=both]".format(self.arrow_shape),
|
|
color,
|
|
)
|
|
else:
|
|
relation = None
|
|
if relation is not None:
|
|
newmodel["relations"].append(relation)
|
|
return newmodel
|
|
|
|
def process_local_many_to_many(self, field, model):
|
|
newmodel = model.copy()
|
|
if self.skip_field(field):
|
|
return newmodel
|
|
relation = None
|
|
if isinstance(field, ManyToManyField):
|
|
if (
|
|
hasattr(field.remote_field.through, "_meta")
|
|
and field.remote_field.through._meta.auto_created
|
|
):
|
|
relation = self.add_relation(
|
|
field,
|
|
newmodel,
|
|
"[arrowhead={} arrowtail={}, dir=both]".format(
|
|
self.arrow_shape, self.arrow_shape
|
|
),
|
|
)
|
|
elif isinstance(field, GenericRelation):
|
|
relation = self.add_relation(
|
|
field,
|
|
newmodel,
|
|
mark_safe(
|
|
'[style="dotted", arrowhead=normal, arrowtail=normal, dir=both]'
|
|
),
|
|
)
|
|
if relation is not None:
|
|
newmodel["relations"].append(relation)
|
|
return newmodel
|
|
|
|
def process_parent(self, parent, appmodel, model):
|
|
newmodel = model.copy()
|
|
if hasattr(parent, "_meta"): # parent is a model
|
|
_rel = self.get_inheritance_context(appmodel, parent)
|
|
# TODO: seems as if abstract models aren't part of models.getModels,
|
|
# which is why they are printed by this without any attributes.
|
|
if _rel not in newmodel["relations"] and self.use_model(_rel["target"]):
|
|
newmodel["relations"].append(_rel)
|
|
return newmodel
|
|
|
|
def sort_model_fields(self, model):
|
|
newmodel = model.copy()
|
|
newmodel["fields"] = sorted(
|
|
newmodel["fields"],
|
|
key=lambda field: (
|
|
not field["primary_key"],
|
|
not field["relation"],
|
|
field["label"],
|
|
),
|
|
)
|
|
return newmodel
|
|
|
|
def use_model(self, model_name):
|
|
"""
|
|
Decide whether to use a model, based on the model name and the lists of
|
|
models to exclude and include.
|
|
"""
|
|
# Check against include list.
|
|
if self.include_models:
|
|
for model_pattern in self.include_models:
|
|
model_pattern = "^%s$" % model_pattern.replace("*", ".*")
|
|
if re.search(model_pattern, model_name):
|
|
return True
|
|
# Check against exclude list.
|
|
if self.exclude_models:
|
|
for model_pattern in self.exclude_models:
|
|
model_pattern = "^%s$" % model_pattern.replace("*", ".*")
|
|
if re.search(model_pattern, model_name):
|
|
return False
|
|
# Return `True` if `include_models` is falsey, otherwise return `False`.
|
|
return not self.include_models
|
|
|
|
def skip_field(self, field):
|
|
if self.exclude_columns:
|
|
if self.verbose_names and field.verbose_name:
|
|
if field.verbose_name in self.exclude_columns:
|
|
return True
|
|
if field.name in self.exclude_columns:
|
|
return True
|
|
if self.relation_fields_only:
|
|
if not isinstance(
|
|
field,
|
|
(
|
|
ForeignKey,
|
|
ManyToManyField,
|
|
OneToOneField,
|
|
RelatedField,
|
|
OneToOneRel,
|
|
ManyToOneRel,
|
|
),
|
|
):
|
|
return True
|
|
return False
|
|
|
|
|
|
def generate_dot(graph_data, template="django_extensions/graph_models/digraph.dot"):
|
|
if isinstance(template, str):
|
|
template = loader.get_template(template)
|
|
|
|
if not isinstance(template, Template) and not (
|
|
hasattr(template, "template") and isinstance(template.template, Template)
|
|
):
|
|
raise Exception(
|
|
"Default Django template loader isn't used. "
|
|
"This can lead to the incorrect template rendering. "
|
|
"Please, check the settings."
|
|
)
|
|
|
|
c = Context(graph_data).flatten()
|
|
dot = template.render(c)
|
|
|
|
return dot
|
|
|
|
|
|
def generate_graph_data(*args, **kwargs):
|
|
generator = ModelGraph(*args, **kwargs)
|
|
generator.generate_graph_data()
|
|
return generator.get_graph_data()
|
|
|
|
|
|
def use_model(model, include_models, exclude_models):
|
|
generator = ModelGraph(
|
|
[], include_models=include_models, exclude_models=exclude_models
|
|
)
|
|
return generator.use_model(model)
|