实现中断指令

PikPakFs.py

@@ -7,7 +7,7 @@ import os
 import logging
 from enum import Enum
 import asyncio
-import uuid
+import shortuuid
 from utils import PathWalker
 from typing import Callable, Awaitable
 
@@ -20,8 +20,8 @@ class TaskStatus(Enum):
 
 class PikPakTaskStatus(Enum):
     PENDING = "pending"
-    REMOTE_DOWNLOADING = "remote downloading"
+    REMOTE_DOWNLOADING = "remote"
-    LOCAL_DOWNLOADING = "local downloading"
+    LOCAL_DOWNLOADING = "local"
 
 class FileDownloadTaskStatus(Enum):
     PENDING = "pending"
@@ -33,7 +33,7 @@ class UnRecoverableError(Exception):
 
 class TaskBase:
     def __init__(self, id : str, tag : str = "", maxConcurrentNumber = -1):
-        self.id : str = uuid.uuid4() if id is None else id 
+        self.id : str = shortuuid.uuid() if id is None else id 
         self.tag : str = tag
         self.maxConcurrentNumber : int = maxConcurrentNumber
 
@@ -59,11 +59,12 @@ class FileDownloadTask(TaskBase):
     TAG = "FileDownloadTask"
     MAX_CONCURRENT_NUMBER = 5
 
-    def __init__(self, nodeId : str, PikPakTaskId : str, id : str = None, status : FileDownloadTaskStatus = FileDownloadTaskStatus.PENDING):
+    def __init__(self, nodeId : str, PikPakTaskId : str, relativePath : str, status : FileDownloadTaskStatus = FileDownloadTaskStatus.PENDING):
-        super().__init__(id, FileDownloadTask.TAG, FileDownloadTask.MAX_CONCURRENT_NUMBER)
+        super().__init__(nodeId, FileDownloadTask.TAG, FileDownloadTask.MAX_CONCURRENT_NUMBER)
         self.status : FileDownloadTaskStatus = status
         self.PikPakTaskId : str = PikPakTaskId
         self.nodeId : str = nodeId
+        self.relativePath : str = relativePath
 
 async def TaskWorker(task : TaskBase):
     try:
@@ -158,11 +159,33 @@ class PikPakFs:
             self.nodes[task.nodeId] = DirNode(task.nodeId, task.name, task.toDirId)    
         else:
             self.nodes[task.nodeId] = FileNode(task.nodeId, task.name, task.toDirId)
+        
         father = self.GetNodeById(task.toDirId)
         if father.id is not None and task.nodeId not in father.childrenId:
             father.childrenId.append(task.nodeId)
         task.status = PikPakTaskStatus.LOCAL_DOWNLOADING
 
+    async def _pikpak_local_downloading(self, task : PikPakTask):
+        node = self.GetNodeById(task.nodeId)
+        if IsFile(node):
+            fileDownloadTask = FileDownloadTask(task.nodeId, task.id, self.NodeToPath(node, node))
+            fileDownloadTask.handler = self._file_download_task_handler
+            self._add_task(fileDownloadTask)
+        elif IsDir(node):
+            # 使用广度优先遍历
+            queue : list[DirNode] = [node]
+            while len(queue) > 0:
+                current = queue.pop(0)
+                await self.Refresh(current)
+                for childId in current.childrenId:
+                    child = self.GetNodeById(childId)
+                    if IsDir(child):
+                        queue.append(child)
+                    elif IsFile(child):
+                        fileDownloadTask = FileDownloadTask(childId, task.id, self.NodeToPath(child, node))
+                        fileDownloadTask.handler = self._file_download_task_handler
+                        self._add_task(fileDownloadTask)
+
     async def _pikpak_task_handler(self, task : PikPakTask):
         while True:
             if task.status == PikPakTaskStatus.PENDING:
@@ -170,7 +193,7 @@ class PikPakFs:
             elif task.status == PikPakTaskStatus.REMOTE_DOWNLOADING:
                 await self._pikpak_offline_downloading(task)
             elif task.status == PikPakTaskStatus.LOCAL_DOWNLOADING:
-                break
+                await self._pikpak_local_downloading(task)
             else:
                 break
 
@@ -180,7 +203,11 @@ class PikPakFs:
     def _add_task(self, task : TaskBase):
         if self.taskQueues.get(task.tag) is None:
             self.taskQueues[task.tag] = []
-        self.taskQueues[task.tag].append(task)
+        taskQueue = self.taskQueues[task.tag]
+        for t in taskQueue:
+            if t.id == task.id:
+                return
+        taskQueue.append(task)
     
     async def StopTask(self, task : TaskBase):
         pass
@@ -317,16 +344,16 @@ class PikPakFs:
         
         node.lastUpdate = datetime.now()
 
-    async def PathToNode(self, pathStr : str) -> FsNode:
+    async def PathToNode(self, path : str) -> FsNode:
-        father, sonName = await self.PathToFatherNodeAndNodeName(pathStr)
+        father, sonName = await self.PathToFatherNodeAndNodeName(path)
         if sonName == "":
             return father
         if not IsDir(father):
             return None
         return self.FindChildInDirByName(father, sonName)
       
-    async def PathToFatherNodeAndNodeName(self, pathStr : str) -> tuple[FsNode, str]:
+    async def PathToFatherNodeAndNodeName(self, path : str) -> tuple[FsNode, str]:
-        pathWalker = PathWalker(pathStr)
+        pathWalker = PathWalker(path)
         father : FsNode = None
         sonName : str = None
         current = self.root if pathWalker.IsAbsolute() else self.currentLocation
@@ -354,12 +381,14 @@ class PikPakFs:
         
         return father, sonName
 
-    def NodeToPath(self, node : FsNode) -> str:
+    def NodeToPath(self, node : FsNode, root : FsNode = None) -> str:
-        if node is self.root:
+        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 = self.GetFatherNode(current)
         spots.append("")
@@ -393,6 +422,12 @@ class PikPakFs:
         self._add_task(task)
         return task
     
+    async def Pull(self, node : FsNode) -> PikPakTask:
+        task = PikPakTask("", node.fatherId, node.id, PikPakTaskStatus.LOCAL_DOWNLOADING)
+        task.handler = self._pikpak_local_downloading
+        self._add_task(task)
+        return task
+
     async def QueryPikPakTasks(self, filterStatus : TaskStatus = None) -> list[PikPakTask]:
         if PikPakTask.TAG not in self.taskQueues:
             return []
@@ -401,6 +436,14 @@ class PikPakFs:
             return taskQueue
         return [task for task in taskQueue if task._status == filterStatus]
     
+    async def QueryFileDownloadTasks(self, filterStatus : TaskStatus = None) -> list[FileDownloadTask]:
+        if FileDownloadTask.TAG not in self.taskQueues:
+            return []
+        taskQueue = self.taskQueues[FileDownloadTask.TAG]
+        if filterStatus is None:
+            return taskQueue
+        return [task for task in taskQueue if task._status == filterStatus]
+
     async def Delete(self, nodes : list[FsNode]) -> None:
         nodeIds = [node.id for node in nodes]
         await self.client.delete_to_trash(nodeIds)

main.py

@@ -6,7 +6,7 @@ import threading
 import colorlog
 from PikPakFs import PikPakFs, IsDir, IsFile, TaskStatus
 import os
-import keyboard
+from tabulate import tabulate
 
 LogFormatter = colorlog.ColoredFormatter(
         "%(log_color)s%(asctime)s - %(levelname)s - %(name)s - %(message)s",
@@ -34,19 +34,32 @@ setup_logging()
 MainLoop : asyncio.AbstractEventLoop = None
 Client = PikPakFs("token.json", proxy="http://127.0.0.1:7897")
 
-def RunSync(func):
+class RunSync:
-    @wraps(func)
+    _current_task : asyncio.Task = None
-    def decorator(*args, **kwargs):
+
+    def StopCurrentRunningCoroutine():
+        if RunSync._current_task is not None:
+            RunSync._current_task.cancel()
+
+    def __init__(self, func):
+        wraps(func)(self)
+
+    def __call__(self, *args, **kwargs):
         currentLoop = None
         try:
             currentLoop = asyncio.get_running_loop()
         except RuntimeError:
+            logging.error("Not in an event loop")
             pass
+        func = self.__wrapped__
         if currentLoop is MainLoop:
-            return MainLoop.run_until_complete(func(*args, **kwargs))
+            task = asyncio.Task(func(*args, **kwargs))
+            RunSync._current_task = task
+            result = MainLoop.run_until_complete(task)
+            RunSync._current_task = None
+            return result
         else:
             return asyncio.run_coroutine_threadsafe(func(*args, **kwargs), MainLoop).result()
-    return decorator
 
 class Console(cmd2.Cmd):
     def _io_worker(self, loop):
@@ -76,7 +89,7 @@ class Console(cmd2.Cmd):
         # 1. 设置忽略SIGINT
         import signal
         def signal_handler(sig, frame):
-            pass
+            RunSync.StopCurrentRunningCoroutine()
         signal.signal(signal.SIGINT, signal_handler)
 
         # 2. 创建IO线程处理输入输出
@@ -125,8 +138,8 @@ class Console(cmd2.Cmd):
     login_parser = cmd2.Cmd2ArgumentParser()
     login_parser.add_argument("username", help="username", nargs="?")
     login_parser.add_argument("password", help="password", nargs="?")
-    @RunSync
     @cmd2.with_argparser(login_parser)
+    @RunSync
     async def do_login(self, args):
         """
         Login to pikpak
@@ -134,7 +147,7 @@ class Console(cmd2.Cmd):
         await Client.Login(args.username, args.password)
         await self.Print("Logged in successfully")
 
-    async def _path_completer(self, text, line, begidx, endidx, filterfiles):   
+    async def _path_completer(self, text, line, begidx, endidx, ignoreFiles):   
         father, sonName = await Client.PathToFatherNodeAndNodeName(text)
         if not IsDir(father):
             return []
@@ -142,7 +155,7 @@ class Console(cmd2.Cmd):
         matchesNode = []
         for childId in father.childrenId:
             child = Client.GetNodeById(childId)
-            if filterfiles and IsFile(child):
+            if ignoreFiles and IsFile(child):
                 continue
             if child.name.startswith(sonName):
                 self.display_matches.append(child.name)
@@ -163,8 +176,8 @@ class Console(cmd2.Cmd):
 
     ls_parser = cmd2.Cmd2ArgumentParser()
     ls_parser.add_argument("path", help="path", default="", nargs="?", type=RunSync(Client.PathToNode))
-    @RunSync
     @cmd2.with_argparser(ls_parser)
+    @RunSync
     async def do_ls(self, args):
         """
         List files in a directory
@@ -187,8 +200,8 @@ class Console(cmd2.Cmd):
 
     cd_parser = cmd2.Cmd2ArgumentParser()
     cd_parser.add_argument("path", help="path", default="", nargs="?", type=RunSync(Client.PathToNode))
-    @RunSync
     @cmd2.with_argparser(cd_parser)
+    @RunSync
     async def do_cd(self, args):
         """
         Change directory
@@ -218,8 +231,8 @@ class Console(cmd2.Cmd):
 
     rm_parser = cmd2.Cmd2ArgumentParser()
     rm_parser.add_argument("paths", help="paths", default="", nargs="+", type=RunSync(Client.PathToNode))
-    @RunSync
     @cmd2.with_argparser(rm_parser)
+    @RunSync
     async def do_rm(self, args):
         """
         Remove a file or directory
@@ -232,8 +245,8 @@ class Console(cmd2.Cmd):
 
     mkdir_parser = cmd2.Cmd2ArgumentParser()
     mkdir_parser.add_argument("path_and_son", help="path and son", default="", nargs="?", type=RunSync(Client.PathToFatherNodeAndNodeName))
-    @RunSync
     @cmd2.with_argparser(mkdir_parser)
+    @RunSync
     async def do_mkdir(self, args):
         """
         Create a directory
@@ -251,11 +264,11 @@ class Console(cmd2.Cmd):
     download_parser = cmd2.Cmd2ArgumentParser()
     download_parser.add_argument("url", help="url")
     download_parser.add_argument("path", help="path", default="", nargs="?", type=RunSync(Client.PathToNode))
-    @RunSync
     @cmd2.with_argparser(download_parser)
+    @RunSync
     async def do_download(self, args):
         """
-        Download a file
+        Download a file or directory
         """
         node = args.path
         if not IsDir(node):
@@ -264,24 +277,47 @@ class Console(cmd2.Cmd):
         task = await Client.Download(args.url, node)
         await self.Print(f"Task {task.id} created")
 
-    query_parser = cmd2.Cmd2ArgumentParser()
-    query_parser.add_argument("-f", "--filter", help="filter", nargs="?", choices=[member.value for member in TaskStatus])
     @RunSync
+    async def complete_pull(self, text, line, begidx, endidx):
+        return await self._path_completer(text, line, begidx, endidx, False)
+
+    pull_parser = cmd2.Cmd2ArgumentParser()
+    pull_parser.add_argument("target", help="pull target", type=RunSync(Client.PathToNode))
+    @cmd2.with_argparser(pull_parser)
+    @RunSync
+    async def do_pull(self, args):
+        """
+        Pull a file or directory
+        """
+        await Client.Pull(args.target)
+        
+
+    query_parser = cmd2.Cmd2ArgumentParser()
+    query_parser.add_argument("-t", "--type", help="type", nargs="?", choices=["pikpak", "filedownload"], default="pikpak")
+    query_parser.add_argument("-f", "--filter", help="filter", nargs="?", choices=[member.value for member in TaskStatus])
     @cmd2.with_argparser(query_parser)
+    @RunSync
     async def do_query(self, args):
         """
         Query All Tasks
         """
-        tasks = await Client.QueryPikPakTasks(TaskStatus(args.filter) if args.filter is not None else None)
+        if args.type == "pikpak":
-        # 格式化输出所有task信息id，status，lastStatus的信息，输出表格
+            tasks = await Client.QueryPikPakTasks(TaskStatus(args.filter) if args.filter is not None else None)
-        await self.Print("tstatus\tdetails\tid")
+            # 格式化输出所有task信息id，status，lastStatus的信息，输出表格
-        for task in tasks:
+            table = [[task.id, task._status.value, task.status.value] for task in tasks]
-            await self.Print(f"{task._status.value}\t{task.status.value}\t{task.id}")
+            headers = ["id", "status", "details"]
+            await self.Print(tabulate(table, headers, tablefmt="grid"))
+        elif args.type == "filedownload":
+            tasks = await Client.QueryFileDownloadTasks(TaskStatus(args.filter) if args.filter is not None else None)
+            # 格式化输出所有task信息id，status，lastStatus的信息，输出表格
+            table = [[task.id, task._status.value, task.status.value, task.relativePath] for task in tasks]
+            headers = ["id", "status", "details", "path"]
+            await self.Print(tabulate(table, headers, tablefmt="grid"))
 
     retry_parser = cmd2.Cmd2ArgumentParser()
     retry_parser.add_argument("taskId", help="taskId")
-    @RunSync
     @cmd2.with_argparser(retry_parser)
+    @RunSync
     async def do_retry(self, args):
         """
         Retry a task
@@ -299,7 +335,10 @@ async def mainLoop():
         stop = False
         while not stop:
             line = await console.Input(console.prompt)
-            stop = console.onecmd_plus_hooks(line)
+            try:
+                stop = console.onecmd_plus_hooks(line)
+            except asyncio.CancelledError:
+                await console.Print("^C: Task cancelled")
     finally:
         console.postloop()
         clientWorker.cancel()