# Imports
from django.core.exceptions import ImproperlyConfigured
from .compat import ANYTREE_INSTALLED, Node, RenderTreeGraph
# Exports
__all__ = (
"get_ancestors",
"get_descendants",
"Diagram",
)
# Functions
[docs]def get_ancestors(record, ancestors=None):
"""Get the ancestors of a given model instance.
:param record: The model instance.
:param ancestors: The existing list of ancestors. Used when overloading.
:type ancestors: list
:rtype: list
"""
if ancestors is None:
ancestors = list()
if record.has_parent():
ancestors.append(record.parent)
ancestors = get_ancestors(record.parent, ancestors)
return ancestors
[docs]def get_descendants(parent, descendants=None):
"""Get the descendants of a given model instance.
:param parent: The model instance.
:param descendants: The existing list of ancestors. Used when overloading.
:type descendants: list
:rtype: list
"""
if descendants is None:
descendants = list()
for c in parent.children.all():
descendants.append(c)
if c.has_children():
descendants = get_descendants(c, descendants)
return descendants
# Classes
[docs]class Diagram(object):
"""Utility for graphing a parent-tree model instance."""
[docs] def __init__(self, record):
"""Initialize the graph.
:param record: The model instance.
"""
if not ANYTREE_INSTALLED:
raise ImproperlyConfigured("Parent-tree graphing requires anytree to be installed.")
self.record = record
[docs] @staticmethod
def get_color(instance):
"""Get the color to use for a record.
:param instance: The model instance.
:rtype: str
"""
try:
return instance.color
except AttributeError:
pass
try:
return instance.get_color()
except AttributeError:
pass
return "#ffffff"
[docs] @staticmethod
def get_display_name(instance):
"""Get the human-friendly name of a record.
:param instance: The model instance.
:rtype: str
"""
try:
return instance.get_display_name()
except AttributeError:
return str(instance)
[docs] @staticmethod
def get_node_options(node):
"""Get the options for graphing the given node. See ``to_graph()``.
:param node: The node for which options are to be provided.
:type node: Node
:rtype: str
"""
options = list()
options.append("shape=box")
options.append("style=filled")
if node.color:
options.append('fillcolor="%s"' % node.color)
else:
options.append('fillcolor=aliceblue')
return ";".join(options)
[docs] def get_structure(self):
"""Get the structure of the record and its children.
:rtype: Node
"""
color = self.get_color(self.record)
display_name = self.get_display_name(self.record)
# noinspection PyCallingNonCallable
root_node = Node(display_name, color=color, level=self.record.level, pk=self.record.pk)
self._get_sub_structure(self.record, root_node)
return root_node
[docs] def to_graph(self):
"""Get the record as a graph ready for rendering.
:rtype: RenderTreeGraph
See also: ``get_structure()``.
"""
root = self.get_structure()
# noinspection PyCallingNonCallable
return RenderTreeGraph(root, nodeattrfunc=self.get_node_options)
def _get_sub_structure(self, instance, parent_node):
"""Get the sub-structure of a given record.
:param instance: The model instance.
:param parent_node: The node to which sub-nodes are added.
:type parent_node: Node
"""
for child in instance.children.all():
color = self.get_color(child)
display_name = self.get_display_name(instance)
# noinspection PyCallingNonCallable
child_node = Node(display_name, parent=parent_node, color=color, level=child.level, pk=child.pk)
if child.children.exists():
self._get_sub_structure(child, child_node)