1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
# ~*~ coding: utf-8 ~*~
#
from treelib import Tree
from treelib.exceptions import NodeIDAbsentError
from collections import defaultdict
from copy import deepcopy
from common.utils import get_object_or_none, get_logger, timeit
from .models import SystemUser, Asset
logger = get_logger(__file__)
def get_system_user_by_name(name):
system_user = get_object_or_none(SystemUser, name=name)
return system_user
def get_system_user_by_id(id):
system_user = get_object_or_none(SystemUser, id=id)
return system_user
class TreeService(Tree):
tag_sep = ' / '
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self.nodes_assets_map = defaultdict(set)
self.all_nodes_assets_map = {}
self._invalid_assets = frozenset()
@classmethod
@timeit
def new(cls):
from .models import Node
from orgs.utils import tmp_to_root_org
with tmp_to_root_org():
all_nodes = list(Node.objects.all().values("key", "value"))
all_nodes.sort(key=lambda x: len(x["key"].split(":")))
tree = cls()
tree.create_node(tag='', identifier='')
for node in all_nodes:
key = node["key"]
value = node["value"]
parent_key = ":".join(key.split(":")[:-1])
tree.safe_create_node(
tag=value, identifier=key,
parent=parent_key,
)
tree.init_assets()
return tree
@timeit
def init_assets(self):
from orgs.utils import tmp_to_root_org
self.all_nodes_assets_map = {}
self.nodes_assets_map = defaultdict(set)
with tmp_to_root_org():
queryset = Asset.objects.all().values_list('id', 'nodes__key')
invalid_assets = Asset.objects.filter(is_active=False)\
.values_list('id', flat=True)
self._invalid_assets = frozenset(invalid_assets)
for asset_id, key in queryset:
if not key:
continue
self.nodes_assets_map[key].add(asset_id)
def safe_create_node(self, **kwargs):
parent = kwargs.get("parent")
if not self.contains(parent):
kwargs['parent'] = self.root
self.create_node(**kwargs)
def all_children_ids(self, nid, with_self=True):
children_ids = self.expand_tree(nid)
if not with_self:
next(children_ids)
return list(children_ids)
def all_children(self, nid, with_self=True, deep=False):
children_ids = self.all_children_ids(nid, with_self=with_self)
return [self.get_node(i, deep=deep) for i in children_ids]
def ancestors(self, nid, with_self=False, deep=False):
ancestor_ids = list(self.rsearch(nid))
ancestor_ids.pop()
if not with_self:
ancestor_ids.pop(0)
return [self.get_node(i, deep=deep) for i in ancestor_ids]
def get_node_full_tag(self, nid):
ancestors = self.ancestors(nid, with_self=True)
ancestors.reverse()
return self.tag_sep.join([n.tag for n in ancestors])
def get_family(self, nid, deep=False):
ancestors = self.ancestors(nid, with_self=False, deep=deep)
children = self.all_children(nid, with_self=False)
return ancestors + [self[nid]] + children
@staticmethod
def is_parent(child, parent):
parent_id = child.bpointer
return parent_id == parent.identifier
def root_node(self):
return self.get_node(self.root)
def get_node(self, nid, deep=False):
node = super().get_node(nid)
if deep:
node = self.copy_node(node)
return node
def parent(self, nid, deep=False):
parent = super().parent(nid)
if deep:
parent = self.copy_node(parent)
return parent
def set_assets(self, nid, assets):
self.nodes_assets_map[nid] = set(assets)
def assets(self, nid):
assets = self.nodes_assets_map[nid]
return assets
def valid_assets(self, nid):
return set(self.assets(nid)) - set(self._invalid_assets)
def all_assets(self, nid):
assets = self.all_nodes_assets_map.get(nid)
if assets:
return assets
assets = set(self.assets(nid))
try:
children = self.children(nid)
except NodeIDAbsentError:
children = []
for child in children:
assets.update(self.all_assets(child.identifier))
self.all_nodes_assets_map[nid] = assets
return assets
def all_valid_assets(self, nid):
return set(self.all_assets(nid)) - set(self._invalid_assets)
def assets_amount(self, nid):
return len(self.all_assets(nid))
def valid_assets_amount(self, nid):
return len(self.all_valid_assets(nid))
@staticmethod
def copy_node(node):
new_node = deepcopy(node)
new_node.fpointer = None
return new_node
def safe_add_ancestors(self, node, ancestors):
# 如果没有祖先节点，那么添加该节点, 父节点是root node
if len(ancestors) == 0:
parent = self.root_node()
else:
parent = ancestors[0]
# 如果当前节点已再树中，则移动当前节点到父节点中
# 这个是由于 当前节点放到了二级节点中
if not self.contains(parent.identifier):
# logger.debug('Add parent: {}'.format(parent.identifier))
self.safe_add_ancestors(parent, ancestors[1:])
if self.contains(node.identifier):
# msg = 'Move node to parent: {} => {}'.format(
# node.identifier, parent.identifier
# )
# logger.debug(msg)
self.move_node(node.identifier, parent.identifier)
else:
# logger.debug('Add node: {}'.format(node.identifier))
self.add_node(node, parent)
#
# def __getstate__(self):
# self.mutex = None
# return self.__dict__
#
def __setstate__(self, state):
self.__dict__ = state
if '_invalid_assets' not in state:
self._invalid_assets = frozenset()
# self.mutex = threading.Lock()