添加任务

PikPakFs.py

@@ -10,6 +10,7 @@ import asyncio
 import shortuuid
 from utils import PathWalker
 from typing import Callable, Awaitable
+import random
 
 class TaskStatus(Enum):
     PENDING = "pending"
@@ -22,11 +23,13 @@ class PikPakTaskStatus(Enum):
     PENDING = "pending"
     REMOTE_DOWNLOADING = "remote"
     LOCAL_DOWNLOADING = "local"
+    DONE = "done"
 
 class FileDownloadTaskStatus(Enum):
     PENDING = "pending"
     DOWNLOADING = "downloading"
-
+    DONE = "done"
+    
 class UnRecoverableError(Exception):
     def __init__(self, message):
         super().__init__(message)
@@ -36,6 +39,7 @@ class TaskBase:
         self.id : str = shortuuid.uuid() if id is None else id 
         self.tag : str = tag
         self.maxConcurrentNumber : int = maxConcurrentNumber
+        self.name : str = ""
 
         self._status : TaskStatus = TaskStatus.PENDING
         
@@ -46,25 +50,36 @@ class PikPakTask(TaskBase):
     TAG = "PikPakTask"
     MAX_CONCURRENT_NUMBER = 5
 
-    def __init__(self, torrent : str, toDirId : str, id : str = None, status : PikPakTaskStatus = PikPakTaskStatus.PENDING):
+    def __init__(self, torrent : str, toDirId : str, nodeId : str = None, status : PikPakTaskStatus = PikPakTaskStatus.PENDING, id : str = None):
         super().__init__(id, PikPakTask.TAG, PikPakTask.MAX_CONCURRENT_NUMBER)
         self.status : PikPakTaskStatus = status
         self.toDirId : str = toDirId
-        self.nodeId : str = None
+        self.nodeId : str = nodeId
-        self.name : str = ""
         self.torrent : str = torrent # todo: 将torrent的附加参数去掉再加入
         self.remoteTaskId : str = None
 
+        self.progress : str = ""
+
+    def __eq__(self, other):
+        if isinstance(other, PikPakTask):
+            return self is other or self.nodeId == other.nodeId
+        return False
+
 class FileDownloadTask(TaskBase):
     TAG = "FileDownloadTask"
     MAX_CONCURRENT_NUMBER = 5
 
-    def __init__(self, nodeId : str, PikPakTaskId : str, relativePath : str, status : FileDownloadTaskStatus = FileDownloadTaskStatus.PENDING):
+    def __init__(self, nodeId : str, PikPakTaskId : str, relativePath : str, status : FileDownloadTaskStatus = FileDownloadTaskStatus.PENDING, id : str = None):
-        super().__init__(nodeId, FileDownloadTask.TAG, FileDownloadTask.MAX_CONCURRENT_NUMBER)
+        super().__init__(id, FileDownloadTask.TAG, FileDownloadTask.MAX_CONCURRENT_NUMBER)
         self.status : FileDownloadTaskStatus = status
         self.PikPakTaskId : str = PikPakTaskId
         self.nodeId : str = nodeId
         self.relativePath : str = relativePath
+    
+    def __eq__(self, other):
+        if isinstance(other, FileDownloadTask):
+            return self is other or (self.nodeId == other.nodeId and self.PikPakTaskId == other.PikPakTaskId)
+        return False
 
 async def TaskWorker(task : TaskBase):
     try:
@@ -82,7 +97,7 @@ async def TaskWorker(task : TaskBase):
 async def TaskManager(taskQueues : Dict[str, list[TaskBase]]):
     # todo: 处理取消的情况
     while True:
-        await asyncio.sleep(1)
+        await asyncio.sleep(0.5)
         for taskQueue in taskQueues.values():
             notRunningTasks = [task for task in taskQueue if task.worker is None or task.worker.done()]
             runningTasksNumber = len(taskQueue) - len(notRunningTasks)
@@ -170,6 +185,7 @@ class PikPakFs:
         if IsFile(node):
             fileDownloadTask = FileDownloadTask(task.nodeId, task.id, self.NodeToPath(node, node))
             fileDownloadTask.handler = self._file_download_task_handler
+            fileDownloadTask.name = task.name
             self._add_task(fileDownloadTask)
         elif IsDir(node):
             # 使用广度优先遍历
@@ -184,7 +200,39 @@ class PikPakFs:
                     elif IsFile(child):
                         fileDownloadTask = FileDownloadTask(childId, task.id, self.NodeToPath(child, node))
                         fileDownloadTask.handler = self._file_download_task_handler
+                        fileDownloadTask.name = task.name
                         self._add_task(fileDownloadTask)
+        
+        # 开始等待下载任务完成
+        while True:
+            fileDownloadTasks = self.taskQueues.get(FileDownloadTask.TAG)
+            myTasks = [myTask for myTask in fileDownloadTasks if myTask.PikPakTaskId == task.id]
+            allNumber = len(myTasks)
+            notCompletedNumber = 0
+            pausedNumber = 0
+            errorNumber = 0
+            for myTask in myTasks:
+                if myTask._status == TaskStatus.PAUSED:
+                    pausedNumber += 1
+                if myTask._status == TaskStatus.ERROR:
+                    errorNumber += 1
+                if myTask._status in {TaskStatus.PENDING, TaskStatus.RUNNING}:
+                    notCompletedNumber += 1
+            
+            runningNumber = allNumber - notCompletedNumber - pausedNumber - errorNumber
+            task.progress = f"{runningNumber}/{allNumber} ({pausedNumber}|{errorNumber})"
+            
+            if notCompletedNumber > 0:
+                await asyncio.sleep(0.5)
+                continue
+            if errorNumber > 0:
+                raise Exception("file download failed")
+            if pausedNumber > 0:
+                raise asyncio.CancelledError()
+            break
+            
+        task.status = PikPakTaskStatus.DONE
+            
 
     async def _pikpak_task_handler(self, task : PikPakTask):
         while True:
@@ -198,32 +246,39 @@ class PikPakFs:
                 break
 
     async def _file_download_task_handler(self, task : FileDownloadTask):
+        if random.randint(1, 5) == 2:
+            raise asyncio.CancelledError()
         pass
 
     def _add_task(self, task : TaskBase):
         if self.taskQueues.get(task.tag) is None:
             self.taskQueues[task.tag] = []
         taskQueue = self.taskQueues[task.tag]
-        for t in taskQueue:
+        for other in taskQueue:
-            if t.id == task.id:
+            if other == task:
+                if other._status != TaskStatus.DONE:
+                    other._status = TaskStatus.PENDING
                 return
         taskQueue.append(task)
     
-    async def StopTask(self, task : TaskBase):
+    async def GetTaskById(self, taskId : str) -> TaskBase:
-        pass
+        for taskQueue in self.taskQueues.values():
+            for task in taskQueue:
+                if task.id == taskId:
+                    return task
+        return None
 
-    async def ResumeTask(self, task : TaskBase):
+    async def StopTask(self, taskId : str):
-        pass
+        task = await self.GetTaskById(taskId)
+        if task is not None and task._status in {TaskStatus.PENDING, TaskStatus.RUNNING}:
+            if task.worker is not None:
+                task.worker.cancel()
+            task.worker = None
 
-    async def RetryTask(self, taskId : str):
+    async def ResumeTask(self, taskId : str):
-        if PikPakTask.TAG not in self.taskQueues:
+        task = await self.GetTaskById(taskId)
-            return
+        if task is not None and task._status in {TaskStatus.PAUSED, TaskStatus.ERROR}:
-        for task in self.taskQueues[PikPakTask.TAG]:
+            task._status = TaskStatus.PENDING
-            if task.id == taskId and task._status == TaskStatus.ERROR:
-                task._status = TaskStatus.PENDING
-                break
-            elif task.id == taskId:
-                break
 
     def Start(self):
         return asyncio.create_task(TaskManager(self.taskQueues))
@@ -424,7 +479,8 @@ class PikPakFs:
     
     async def Pull(self, node : FsNode) -> PikPakTask:
         task = PikPakTask("", node.fatherId, node.id, PikPakTaskStatus.LOCAL_DOWNLOADING)
-        task.handler = self._pikpak_local_downloading
+        task.name = node.name
+        task.handler = self._pikpak_task_handler
         self._add_task(task)
         return task
 

main.py

@@ -7,6 +7,8 @@ import colorlog
 from PikPakFs import PikPakFs, IsDir, IsFile, TaskStatus
 import os
 from tabulate import tabulate
+import wcwidth
+import types
 
 LogFormatter = colorlog.ColoredFormatter(
         "%(log_color)s%(asctime)s - %(levelname)s - %(name)s - %(message)s",
@@ -61,6 +63,12 @@ class RunSync:
         else:
             return asyncio.run_coroutine_threadsafe(func(*args, **kwargs), MainLoop).result()
 
+    def __get__(self, instance, cls):
+        if instance is None:
+            return self
+        else:
+            return types.MethodType(self, instance)
+
 class Console(cmd2.Cmd):
     def _io_worker(self, loop):
         asyncio.set_event_loop(loop)
@@ -165,7 +173,8 @@ class Console(cmd2.Cmd):
                     matches.append(text[:text.rfind(sonName)] + child.name)
                 matchesNode.append(child)
         if len(matchesNode) == 1 and IsDir(matchesNode[0]):
-            matches[0] += "/"
+            if matches[0] == sonName:
+                matches[0] += "/"
             self.allow_appended_space = False
             self.allow_closing_quote = False
         return matches
@@ -304,8 +313,8 @@ class Console(cmd2.Cmd):
         if args.type == "pikpak":
             tasks = await Client.QueryPikPakTasks(TaskStatus(args.filter) if args.filter is not None else None)
             # 格式化输出所有task信息id，status，lastStatus的信息，输出表格
-            table = [[task.id, task._status.value, task.status.value] for task in tasks]
+            table = [[task.id, task._status.value, task.status.value, task.progress] for task in tasks]
-            headers = ["id", "status", "details"]
+            headers = ["id", "status", "details", "progress"]
             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)
@@ -314,15 +323,24 @@ class Console(cmd2.Cmd):
             headers = ["id", "status", "details", "path"]
             await self.Print(tabulate(table, headers, tablefmt="grid"))
 
-    retry_parser = cmd2.Cmd2ArgumentParser()
+    taskid_parser = cmd2.Cmd2ArgumentParser()
-    retry_parser.add_argument("taskId", help="taskId")
+    taskid_parser.add_argument("taskId", help="taskId")
-    @cmd2.with_argparser(retry_parser)
+    
+    @cmd2.with_argparser(taskid_parser)
     @RunSync
-    async def do_retry(self, args):
+    async def do_pause(self, args):
         """
-        Retry a task
+        Stop a task
         """
-        await Client.RetryTask(args.taskId)
+        await Client.StopTask(args.taskId)
+
+    @cmd2.with_argparser(taskid_parser)
+    @RunSync
+    async def do_resume(self, args):
+        """
+        Resume a task
+        """
+        await Client.ResumeTask(args.taskId)
 
 async def mainLoop():
     global MainLoop, Client