from anytree import NodeMixin class TorNode(NodeMixin): def __init__( self, name, is_folder=False, is_file=False, parent=None, size=None, priority=None, file_id=None, progress=None, ): super().__init__() self.name = name self.is_folder = is_folder self.is_file = is_file if parent is not None: self.parent = parent if size is not None: self.fsize = size if priority is not None: self.priority = priority if file_id is not None: self.file_id = file_id if progress is not None: self.progress = progress def qb_get_folders(path): return path.split("/") def get_folders(path, root_path): fs = path.split(root_path)[-1] return fs.split("/") def make_tree(res, tool, root_path=""): if tool == "qbittorrent": parent = TorNode("QBITTORRENT") folder_id = 0 for i in res: folders = qb_get_folders(i.name) if len(folders) > 1: previous_node = parent for j in range(len(folders) - 1): current_node = next( (k for k in previous_node.children if k.name == folders[j]), None, ) if current_node is None: previous_node = TorNode( folders[j], is_folder=True, parent=previous_node, file_id=folder_id, ) folder_id += 1 else: previous_node = current_node TorNode( folders[-1], is_file=True, parent=previous_node, size=i.size, priority=i.priority, file_id=i.index, progress=round(i.progress * 100, 5), ) else: TorNode( folders[-1], is_file=True, parent=parent, size=i.size, priority=i.priority, file_id=i.index, progress=round(i.progress * 100, 5), ) elif tool == "aria2": parent = TorNode("ARIA2") folder_id = 0 for i in res: folders = get_folders(i["path"], root_path) priority = 1 if i["selected"] == "false": priority = 0 if len(folders) > 1: previous_node = parent for j in range(len(folders) - 1): current_node = next( (k for k in previous_node.children if k.name == folders[j]), None, ) if current_node is None: previous_node = TorNode( folders[j], is_folder=True, parent=previous_node, file_id=folder_id, ) folder_id += 1 else: previous_node = current_node try: progress = round( (int(i["completedLength"]) / int(i["length"])) * 100, 5 ) except ZeroDivisionError: progress = 0 TorNode( folders[-1], is_file=True, parent=previous_node, size=int(i["length"]), priority=priority, file_id=i["index"], progress=progress, ) else: try: progress = round( (int(i["completedLength"]) / int(i["length"])) * 100, 5 ) except ZeroDivisionError: progress = 0 TorNode( folders[-1], is_file=True, parent=parent, size=int(i["length"]), priority=priority, file_id=i["index"], progress=progress, ) else: parent = TorNode("SABNZBD+") priority = 1 for i in res["files"]: TorNode( i["filename"], is_file=True, parent=parent, size=float(i["mb"]) * 1048576, priority=priority, file_id=i["nzf_id"], progress=round( ((float(i["mb"]) - float(i["mbleft"])) / float(i["mb"])) * 100, 5, ), ) result = create_list(parent) return {"files": result, "engine": tool} """ def print_tree(parent): for pre, _, node in RenderTree(parent): treestr = u"%s%s" % (pre, node.name) print(treestr.ljust(8), node.is_folder, node.is_file) """ def create_list(parent, contents=None): if contents is None: contents = [] for i in parent.children: if i.is_folder: children = [] create_list(i, children) contents.append( { "id": f"folderNode_{i.file_id}", "name": i.name, "type": "folder", "children": children, } ) else: contents.append( { "id": i.file_id, "name": i.name, "size": i.fsize, "type": "file", "selected": bool(i.priority), "progress": i.progress, } ) return contents def extract_file_ids(data): selected_files = [] unselected_files = [] for item in data: if item.get("type") == "file": if item.get("selected"): selected_files.append(str(item["id"])) else: unselected_files.append(str(item["id"])) if item.get("children"): child_selected, child_unselected = extract_file_ids(item["children"]) selected_files.extend(child_selected) unselected_files.extend(child_unselected) return selected_files, unselected_files