diff --git a/PikPakFileSystem.py b/PikPakFileSystem.py index 859eff8..cb9035b 100644 --- a/PikPakFileSystem.py +++ b/PikPakFileSystem.py @@ -221,12 +221,14 @@ class PikPakFileSystem: return father, sonName - async def _node_to_path(self, node : NodeBase) -> str: - if node is self._root: + async def _node_to_path(self, node : NodeBase, root : NodeBase = None) -> str: + if root is None: + root = self._root + if node is root: return "/" spots : list[str] = [] current = node - while current is not self._root: + while current is not root: spots.append(current.name) current = await self._get_father_node(current) spots.append("") @@ -262,7 +264,7 @@ class PikPakFileSystem: async def IsDir(self, path : str) -> bool: node = await self._path_to_node(path) return isinstance(node, DirNode) - + async def SplitPath(self, path : str) -> tuple[str, str]: father, son_name = await self._path_to_father_node_and_son_name(path) return await self._node_to_path(father), son_name @@ -307,19 +309,21 @@ class PikPakFileSystem: async def GetCwd(self) -> str: return await self._node_to_path(self._cwd) - async def GetNodeIdByPath(self, path : str) -> str: + async def GetChildren(self, node : NodeBase) -> list[NodeBase]: + if not isinstance(node, DirNode): + return [] + await self._refresh(node) + return [await self._get_node_by_id(child_id) for child_id in node.children_id] + + async def PathToNode(self, path : str) -> NodeBase: node = await self._path_to_node(path) if node is None: return None - return node.id + return node - async def IfExists(self, path : str) -> bool: - return await self._path_to_node(path) is not None + async def NodeToPath(self, from_node : NodeBase, to_node : NodeBase) -> str: + return await self._node_to_path(to_node, from_node) - async def ToFullPath(self, path : str) -> str: - node = await self._path_to_node(path) - return await self._node_to_path(node) - async def RemoteDownload(self, torrent : str, remote_base_path : str) -> tuple[str, str]: node = await self._path_to_node(remote_base_path) info = await self._pikpak_client.offline_download(torrent, node.id) @@ -328,25 +332,19 @@ class PikPakFileSystem: async def QueryTaskStatus(self, task_id : str, node_id : str) -> DownloadStatus: return await self._pikpak_client.get_task_status(task_id, node_id) - async def UpdateDirectory(self, path : str, son_id : str) -> str: - await self._path_to_node(path) - son_info = await self._pikpak_client.offline_file_info(son_id) - kind = son_info["kind"] - parent_id = son_info["parent_id"] - name = son_info["name"] - son : NodeBase = None - if kind.endswith("folder"): - son = DirNode(son_id, name, parent_id) - else: - son = FileNode(son_id, name, parent_id) - await self._add_node(son) - return await self._node_to_path(son) - - async def JoinPath(self, father : str, son : str) -> str: - father_node = await self._path_to_node(father) - son_node = await self._find_child_in_dir_by_name(father_node, son) - if son_node is None: - raise Exception("Son not found") - return await self._node_to_path(son_node) + async def UpdateNode(self, node_id : str) -> NodeBase: + node : NodeBase = await self._get_node_by_id(node_id) + if node is None: + info = await self._pikpak_client.offline_file_info(node_id) + kind = info["kind"] + parent_id = info["parent_id"] + name = info["name"] + if kind.endswith("folder"): + node = DirNode(node_id, name, parent_id) + else: + node = FileNode(node_id, name, parent_id) + await self._add_node(node) + node.lastUpdate = None + return node #endregion \ No newline at end of file diff --git a/TaskManager.py b/TaskManager.py index fc6fc2a..15a8fb0 100644 --- a/TaskManager.py +++ b/TaskManager.py @@ -3,7 +3,7 @@ from typing import Awaitable, Callable, Dict import asyncio import logging import shortuuid -from PikPakFileSystem import PikPakFileSystem +from PikPakFileSystem import PikPakFileSystem, FileNode, DirNode from pikpakapi import DownloadStatus import random @@ -49,11 +49,11 @@ class TorrentTask(TaskBase): super().__init__(self) self.torrent_status : TorrentTaskStatus = TorrentTaskStatus.PENDING self.torrent : str = torrent - self.remote_full_path : str = None - self.remote_base_path : str = None self.info : str = "" + self.name : str = "" # 和PikPak交互需要的信息 + self.remote_base_path : str = None self.node_id : str = None self.task_id : str = None @@ -62,9 +62,10 @@ class FileDownloadTask(TaskBase): TAG = "FileDownloadTask" MAX_CONCURRENT_NUMBER = 5 - def __init__(self, remote_path : str, owner_id : str): + def __init__(self, node_id : str, remote_path : str, owner_id : str): super().__init__(self) self.file_download_status : FileDownloadTaskStatus = FileDownloadTaskStatus.PENDING + self.node_id : str = node_id self.remote_path : str = remote_path self.owner_id : str = owner_id @@ -141,25 +142,28 @@ class TaskManager: await asyncio.sleep(wait_seconds) wait_seconds = wait_seconds * 1.5 - task.remote_full_path = await self.client.UpdateDirectory(task.remote_base_path, task.node_id) task.torrent_status = TorrentTaskStatus.LOCAL_DOWNLOADING async def _on_torrent_local_downloading(self, task : TorrentTask): - path = task.remote_full_path + node = await self.client.UpdateNode(task.node_id) + task.name = node.name + task.node_id = node.id - if not await self.client.IsDir(path): - await self._init_file_download_task(path, task.id) - else: + if isinstance(node, FileNode): + await self._init_file_download_task(task.node_id, task.name, task.id) + elif isinstance(node, DirNode): # 使用广度优先遍历 - queue : list[str] = [path] + queue : list[str] = [node] while len(queue) > 0: - current_path = queue.pop(0) - for child_name in await self.client.GetChildrenNames(current_path, False): - child_path = await self.client.JoinPath(current_path, child_name) - if await self.client.IsDir(child_path): - queue.append(child_path) - else: - await self._init_file_download_task(child_path, task.id) + current = queue.pop(0) + for child in await self.client.GetChildren(current): + if isinstance(child, DirNode): + queue.append(child) + if isinstance(child, FileNode): + child_path = task.name + await self.client.NodeToPath(node, child) + await self._init_file_download_task(child.id, child_path, task.id) + else: + raise Exception("unknown node type") # 开始等待下载任务完成 while True: @@ -214,22 +218,21 @@ class TaskManager: #region 文件下载部分 - async def _init_file_download_task(self, remote_path : str, owner_id : str) -> str: + async def _init_file_download_task(self, node_id : str, remote_path : str, owner_id : str) -> str: queue = await self._get_file_download_queue(owner_id) for task in queue: if not isinstance(task, FileDownloadTask): continue - if task.remote_path == remote_path: + if task.node_id == node_id: if task.status in {TaskStatus.PAUSED, TaskStatus.ERROR}: task.status = TaskStatus.PENDING return task.id - task = FileDownloadTask(remote_path, owner_id) + task = FileDownloadTask(node_id, remote_path, owner_id) task.handler = self._file_download_task_handler await self._append_task(task) return task.id async def _file_download_task_handler(self, task : FileDownloadTask): - await asyncio.sleep(30) if random.randint(1, 5) == 2: raise asyncio.CancelledError() if random.randint(1, 5) == 3: @@ -260,17 +263,19 @@ class TaskManager: await self._append_task(task) return task.id - async def PullRemote(self, remote_full_path : str) -> str: - if not await self.client.IfExists(remote_full_path): + async def PullRemote(self, path : str) -> str: + target = await self.client.PathToNode(path) + if target is None: raise Exception("target not found") - queue = await self._get_torrent_queue() + queue = await self._get_torrent_queue() for task in queue: if not isinstance(task, TorrentTask): continue - if task.remote_full_path == remote_full_path: + if task.node_id == target.id: return task.id task = TorrentTask(None) - task.remote_full_path = remote_full_path + task.name = target.name + task.node_id = target.id task.handler = self._torrent_task_handler task.torrent_status = TorrentTaskStatus.LOCAL_DOWNLOADING await self._append_task(task) diff --git a/main.py b/main.py index 62350a0..f2ba45e 100644 --- a/main.py +++ b/main.py @@ -166,10 +166,10 @@ class App(cmd2.Cmd): await Client.Login(args.username, args.password) await self.print("Logged in successfully") - async def _path_completer(self, text, line, begidx, endidx, ignoreFiles): + async def _path_completer(self, text, line, begidx, endidx, ignore_files): father_path, son_name = await Client.SplitPath(text) - children_names = await Client.GetChildrenNames(father_path, ignoreFiles) - matches = [] + children_names = await Client.GetChildrenNames(father_path, ignore_files) + matches : list[str] = [] for child_name in children_names: if child_name.startswith(son_name): self.display_matches.append(child_name) @@ -177,8 +177,8 @@ class App(cmd2.Cmd): matches.append(text + child_name) elif text.endswith(son_name): matches.append(text[:text.rfind(son_name)] + child_name) - if len(matches) == 1 and await Client.IsDir(father_path + matches[0]): - if matches[0] == son_name: + if len(matches) == 1 and await Client.IsDir(matches[0]): + if matches[0].endswith(son_name): matches[0] += "/" self.allow_appended_space = False self.allow_closing_quote = False @@ -278,7 +278,7 @@ class App(cmd2.Cmd): """ Pull a file or directory """ - task_id = await self.task_manager.PullRemote(await Client.ToFullPath(args.target)) + task_id = await self.task_manager.PullRemote(args.target) await self.print(f"Task {task_id} created")