1# Imports 

2 

3from django.core.exceptions import ImproperlyConfigured 

4from .compat import ANYTREE_INSTALLED, Node, RenderTreeGraph 

5 

6# Exports 

7 

8__all__ = ( 

9 "get_ancestors", 

10 "get_descendants", 

11 "Diagram", 

12) 

13 

14# Functions 

15 

16 

17def get_ancestors(record, ancestors=None): 

18 """Get the ancestors of a given model instance. 

19 

20 :param record: The model instance. 

21 

22 :param ancestors: The existing list of ancestors. Used when overloading. 

23 :type ancestors: list 

24 

25 :rtype: list 

26 

27 """ 

28 if ancestors is None: 

29 ancestors = list() 

30 

31 if record.has_parent(): 

32 ancestors.append(record.parent) 

33 

34 ancestors = get_ancestors(record.parent, ancestors) 

35 

36 return ancestors 

37 

38 

39def get_descendants(parent, descendants=None): 

40 """Get the descendants of a given model instance. 

41 

42 :param parent: The model instance. 

43 

44 :param descendants: The existing list of ancestors. Used when overloading. 

45 :type descendants: list 

46 

47 :rtype: list 

48 

49 """ 

50 if descendants is None: 

51 descendants = list() 

52 

53 for c in parent.children.all(): 

54 descendants.append(c) 

55 if c.has_children(): 

56 descendants = get_descendants(c, descendants) 

57 

58 return descendants 

59 

60 

61# Classes 

62 

63 

64class Diagram(object): 

65 """Utility for graphing a parent-tree model instance.""" 

66 

67 def __init__(self, record): 

68 """Initialize the graph. 

69 

70 :param record: The model instance. 

71 

72 """ 

73 if not ANYTREE_INSTALLED: 

74 raise ImproperlyConfigured("Parent-tree graphing requires anytree to be installed.") 

75 

76 self.record = record 

77 

78 @staticmethod 

79 def get_color(instance): 

80 """Get the color to use for a record. 

81 

82 :param instance: The model instance. 

83 

84 :rtype: str 

85 

86 """ 

87 try: 

88 return instance.color 

89 except AttributeError: 

90 pass 

91 

92 try: 

93 return instance.get_color() 

94 except AttributeError: 

95 pass 

96 

97 return "#ffffff" 

98 

99 @staticmethod 

100 def get_display_name(instance): 

101 """Get the human-friendly name of a record. 

102 

103 :param instance: The model instance. 

104 

105 :rtype: str 

106 

107 """ 

108 try: 

109 return instance.get_display_name() 

110 except AttributeError: 

111 return str(instance) 

112 

113 @staticmethod 

114 def get_node_options(node): 

115 """Get the options for graphing the given node. See ``to_graph()``. 

116 

117 :param node: The node for which options are to be provided. 

118 :type node: Node 

119 

120 :rtype: str 

121 

122 """ 

123 options = list() 

124 options.append("shape=box") 

125 options.append("style=filled") 

126 

127 if node.color: 

128 options.append('fillcolor="%s"' % node.color) 

129 else: 

130 options.append('fillcolor=aliceblue') 

131 

132 return ";".join(options) 

133 

134 def get_structure(self): 

135 """Get the structure of the record and its children. 

136 

137 :rtype: Node 

138 

139 """ 

140 color = self.get_color(self.record) 

141 display_name = self.get_display_name(self.record) 

142 

143 # noinspection PyCallingNonCallable 

144 root_node = Node(display_name, color=color, level=self.record.level, pk=self.record.pk) 

145 

146 self._get_sub_structure(self.record, root_node) 

147 

148 return root_node 

149 

150 def to_graph(self): 

151 """Get the record as a graph ready for rendering. 

152 

153 :rtype: RenderTreeGraph 

154 

155 See also: ``get_structure()``. 

156 

157 """ 

158 root = self.get_structure() 

159 

160 # noinspection PyCallingNonCallable 

161 return RenderTreeGraph(root, nodeattrfunc=self.get_node_options) 

162 

163 def _get_sub_structure(self, instance, parent_node): 

164 """Get the sub-structure of a given record. 

165 

166 :param instance: The model instance. 

167 

168 :param parent_node: The node to which sub-nodes are added. 

169 :type parent_node: Node 

170 

171 """ 

172 for child in instance.children.all(): 

173 

174 color = self.get_color(child) 

175 display_name = self.get_display_name(instance) 

176 

177 # noinspection PyCallingNonCallable 

178 child_node = Node(display_name, parent=parent_node, color=color, level=child.level, pk=child.pk) 

179 

180 if child.children.exists(): 

181 self._get_sub_structure(child, child_node)