def compute_nodes_assets_amount(self, nodes_with_assets):
self.stack = Stack()
nodes_items = []
for key, values in nodes_with_assets.items():
nodes_items.append({
"key": key,
"assets": values["assets"],
"all_assets": values["all_assets"],
"assets_amount": 0
})
nodes_items.sort(key=self.sorted_by)
nodes_items.append({
"key": "",
"assets": set(),
"all_assets": set(),
"assets_amount": 0
})
self.stack = Stack()
for item in nodes_items:
self.debug("准备: {} 栈顶: {}".format(
item['key'],
self.stack.top["key"] if self.stack.top else None))
# 入栈之前检查，该节点是不是栈顶节点的子节点
# 如果不是，则栈顶出栈
while self.stack.top and not self.is_children(
self.stack.top, item):
self.pop_from_stack_nodes_amount()
self.push_to_stack_nodes_amount(item)
# 出栈最后一个
self.debug("剩余: {}".format(', '.join([n["key"] for n in self.stack])))
return self._nodes
def get_nodes_family_and_system_users(self, nodes_with_system_users):
"""
返回所有nodes_with_system_users中的node的家族节点的信息，
并子会继承祖先的系统用户和actions信息
:param nodes_with_system_users:
{node.key: {system_user.id: actions,}, }
:return:
{node.key: {system_user.id: actions,}, }
"""
node_util = NodeUtil()
_nodes_keys = nodes_with_system_users.keys()
family_keys = node_util.get_some_nodes_family_keys_by_keys(_nodes_keys)
nodes_items = []
for i in family_keys:
system_users = nodes_with_system_users.get(i, defaultdict(int))
item = {"key": i, "system_users": system_users}
nodes_items.append(item)
# 按照父子关系排序
nodes_items.sort(key=self.sorted_by)
nodes_items.append({"key": "", "system_users": defaultdict(int)})
self.stack = Stack()
for item in nodes_items:
self.debug("准备: {} 栈顶: {}".format(
item['key'],
self.stack.top["key"] if self.stack.top else None))
# 入栈之前检查，该节点是不是栈顶节点的子节点
# 如果不是，则栈顶出栈
while self.stack.top and not self.is_children(
self.stack.top, item):
# 出栈
self.pop_from_stack_system_users()
# 入栈
self.push_to_stack_system_users(item)
# 出栈最后一个
self.debug("剩余: {}".format(', '.join([n["key"] for n in self.stack])))
return self._nodes
def generate(self):
# 准备排序好的资产信息数据
infos = []
for node in self.nodes:
assets = self.nodekey_assetsid_mapper.get(node.key, set())
info = NodeAssetsInfo(key=node.key, assets_amount=0, assets=assets)
infos.append(info)
infos = sorted(infos, key=lambda i: [int(i) for i in i.key.split(':')])
# 这个守卫需要添加一下，避免最后一个无法出栈
guarder = NodeAssetsInfo(key='', assets_amount=0, assets=set())
infos.append(guarder)
stack = Stack()
for info in infos:
# 如果栈顶的不是这个节点的父祖节点，那么可以出栈了，可以计算资产数量了
while stack.top and not info.key.startswith(f'{stack.top.key}:'):
pop_info = stack.pop()
pop_info.assets_amount = len(pop_info.assets)
self.nodekey_assetsinfo_mapper[pop_info.key] = pop_info
if not stack.top:
continue
stack.top.assets.update(pop_info.assets)
stack.push(info)
def __init__(self, with_assets_amount=False, debug=False):
self.stack = Stack()
self._nodes = {}
self.with_assets_amount = with_assets_amount
self._debug = debug
self.init()
class NodeUtil:
def __init__(self, with_assets_amount=False, debug=False):
self.stack = Stack()
self._nodes = {}
self.with_assets_amount = with_assets_amount
self._debug = debug
self.init()
@staticmethod
def sorted_by(node):
return [int(i) for i in node.key.split(':')]
def get_queryset(self):
all_nodes = Node.objects.all()
if self.with_assets_amount:
all_nodes = all_nodes.prefetch_related(
Prefetch('assets', queryset=Asset.objects.all().only('id')))
all_nodes = list(all_nodes)
for node in all_nodes:
node._assets = set(node.assets.all())
return all_nodes
def get_all_nodes(self):
all_nodes = sorted(self.get_queryset(), key=self.sorted_by)
guarder = Node(key='', value='Guarder')
guarder._assets = []
all_nodes.append(guarder)
return all_nodes
def push_to_stack(self, node):
# 入栈之前检查
# 如果栈是空的，证明是一颗树的根部
if self.stack.is_empty():
node._full_value = node.value
node._parents = []
else:
# 如果不是根节点,
# 该节点的祖先应该是父节点的祖先加上父节点
# 该节点的名字是父节点的名字+自己的名字
node._parents = [self.stack.top] + self.stack.top._parents
node._full_value = ' / '.join(
[self.stack.top._full_value, node.value])
node._children = []
node._all_children = []
self.debug("入栈: {}".format(node.key))
self.stack.push(node)
# 出栈
def pop_from_stack(self):
_node = self.stack.pop()
self.debug("出栈: {} 栈顶: {}".format(
_node.key, self.stack.top.key if self.stack.top else None))
self._nodes[_node.key] = _node
if not self.stack.top:
return
if self.with_assets_amount:
self.stack.top._assets.update(_node._assets)
_node._assets_amount = len(_node._assets)
delattr(_node, '_assets')
self.stack.top._children.append(_node)
self.stack.top._all_children.extend([_node] + _node._all_children)
def init(self):
all_nodes = self.get_all_nodes()
for node in all_nodes:
self.debug("准备: {} 栈顶: {}".format(
node.key, self.stack.top.key if self.stack.top else None))
# 入栈之前检查，该节点是不是栈顶节点的子节点
# 如果不是，则栈顶出栈
while self.stack.top and not self.stack.top.is_children(node):
self.pop_from_stack()
self.push_to_stack(node)
# 出栈最后一个
self.debug("剩余: {}".format(', '.join([n.key for n in self.stack])))
def get_nodes_by_queryset(self, queryset):
nodes = []
for n in queryset:
node = self.get_node_by_key(n.key)
if not node:
continue
nodes.append(node)
return nodes
def get_node_by_key(self, key):
return self._nodes.get(key)
def debug(self, msg):
self._debug and logger.debug(msg)
def set_assets_amount(self):
for node in self._nodes.values():
node.assets_amount = node._assets_amount
def set_full_value(self):
for node in self._nodes.values():
node.full_value = node._full_value
@property
def nodes(self):
return list(self._nodes.values())
def get_family_by_key(self, key):
tree_nodes = set()
node = self.get_node_by_key(key)
if not node:
return []
tree_nodes.update(node._parents)
tree_nodes.add(node)
tree_nodes.update(node._all_children)
return list(tree_nodes)
# 使用给定节点生成一颗树
# 找到他们的祖先节点
# 可选找到他们的子孙节点
def get_family(self, node):
return self.get_family_by_key(node.key)
def get_family_keys_by_key(self, key):
nodes = self.get_family_by_key(key)
return [n.key for n in nodes]
def get_some_nodes_family_by_keys(self, keys):
family = set()
for key in keys:
family.update(self.get_family_by_key(key))
return family
def get_some_nodes_family_keys_by_keys(self, keys):
family = self.get_some_nodes_family_by_keys(keys)
return [n.key for n in family]
def get_nodes_parents_by_key(self, key, with_self=True):
parents = set()
node = self.get_node_by_key(key)
if not node:
return []
parents.update(set(node._parents))
if with_self:
parents.add(node)
return list(parents)
def get_node_parents(self, node, with_self=True):
return self.get_nodes_parents_by_key(node.key, with_self=with_self)
def get_nodes_parents_keys_by_key(self, key, with_self=True):
nodes = self.get_nodes_parents_by_key(key, with_self=with_self)
return [n.key for n in nodes]
def get_all_children_by_key(self, key, with_self=True):
children = set()
node = self.get_node_by_key(key)
if not node:
return []
children.update(set(node._all_children))
if with_self:
children.add(node)
return list(children)
def get_all_children(self, node, with_self=True):
return self.get_all_children_by_key(node.key, with_self=with_self)
def get_all_children_keys_by_key(self, key, with_self=True):
nodes = self.get_all_children_by_key(key, with_self=with_self)
return [n.key for n in nodes]
class PermSystemUserNodeUtil(PermStackUtilMixin):
"""
self._nodes: {node.key: {system_user.id: actions,}}
"""
@timeit
def get_nodes_family_and_system_users(self, nodes_with_system_users):
"""
返回所有nodes_with_system_users中的node的家族节点的信息，
并子会继承祖先的系统用户和actions信息
:param nodes_with_system_users:
{node.key: {system_user.id: actions,}, }
:return:
{node.key: {system_user.id: actions,}, }
"""
node_util = NodeUtil()
_nodes_keys = nodes_with_system_users.keys()
family_keys = node_util.get_some_nodes_family_keys_by_keys(_nodes_keys)
nodes_items = []
for i in family_keys:
system_users = nodes_with_system_users.get(i, defaultdict(int))
item = {"key": i, "system_users": system_users}
nodes_items.append(item)
# 按照父子关系排序
nodes_items.sort(key=self.sorted_by)
nodes_items.append({"key": "", "system_users": defaultdict(int)})
self.stack = Stack()
for item in nodes_items:
self.debug("准备: {} 栈顶: {}".format(
item['key'],
self.stack.top["key"] if self.stack.top else None))
# 入栈之前检查，该节点是不是栈顶节点的子节点
# 如果不是，则栈顶出栈
while self.stack.top and not self.is_children(
self.stack.top, item):
# 出栈
self.pop_from_stack_system_users()
# 入栈
self.push_to_stack_system_users(item)
# 出栈最后一个
self.debug("剩余: {}".format(', '.join([n["key"] for n in self.stack])))
return self._nodes
def push_to_stack_system_users(self, item):
"""
:param item:
{"key": node.key, "system_users": {system_user.id: actions,},}
"""
if not self.stack.is_empty():
item_system_users = item["system_users"]
for system_user, action in self.stack.top["system_users"].items():
# 更新栈顶的系统用户和action到将要入栈的item中
item_system_users[system_user] |= action
item["system_users"] = item_system_users
self.debug("入栈: {}".format(item['key']))
self.stack.push(item)
# 出栈
def pop_from_stack_system_users(self):
_node = self.stack.pop()
self._nodes[_node["key"]] = _node["system_users"]
self.debug("出栈: {} 栈顶: {}".format(
_node['key'], self.stack.top['key'] if self.stack.top else None))