优化cmd2结构

main.py

@@ -4,41 +4,10 @@ from functools import wraps
 import logging
 import threading
 import colorlog
-from pikpakFs import FsNode, DirNode, FileNode, PKFs, IsDir, IsFile
+from pikpakFs import PKFs, IsDir, IsFile
 import os
 
-def RunSyncInLoop(loop):
+def setup_logging():
-    def decorator(func):
-        @wraps(func)
-        def decorated(*args, **kwargs):
-            currentLoop = None
-            try:
-                currentLoop = asyncio.get_running_loop()
-            except RuntimeError:
-                pass
-            if currentLoop is loop:
-                return loop.run_until_complete(func(*args, **kwargs))
-            else:
-                return asyncio.run_coroutine_threadsafe(func(*args, **kwargs), loop).result()
-        return decorated
-    return decorator
-
-def ProvideDecoratorSelfArgs(decorator, argsProvider):
-    def wrapper(func):
-        @wraps(func)
-        def decorated(*args, **kwargs):
-            namespace = args[0]
-            return decorator(argsProvider(namespace))(func)(*args, **kwargs)
-        return decorated
-    return wrapper
-
-class PikpakConsole(cmd2.Cmd):
-    def LoopProvider(self):
-        return self.loop
-
-    RunSync = ProvideDecoratorSelfArgs(RunSyncInLoop, LoopProvider)
-    
-    def _setup_logging(self):
     formatter = colorlog.ColoredFormatter(
         "%(log_color)s%(asctime)s - %(levelname)s - %(name)s - %(message)s",
         datefmt='%Y-%m-%d %H:%M:%S',
@@ -64,123 +33,77 @@ class PikpakConsole(cmd2.Cmd):
     logger.addHandler(file_handler)
     logger.setLevel(logging.DEBUG)
 
-    def IOWorker(self, loop):
+setup_logging()
-        self.terminal_lock.acquire() # 我看cmdloop是这么做的，所以我也这么做
+MainLoop : asyncio.AbstractEventLoop = None
+Client = PKFs("token.json", proxy="http://127.0.0.1:7890")
+
+def RunSync(func):
+    @wraps(func)
+    def decorator(*args, **kwargs):
+        currentLoop = None
+        try:
+            currentLoop = asyncio.get_running_loop()
+        except RuntimeError:
+            pass
+        if currentLoop is MainLoop:
+            return MainLoop.run_until_complete(func(*args, **kwargs))
+        else:
+            return asyncio.run_coroutine_threadsafe(func(*args, **kwargs), MainLoop).result()
+    return decorator
+
+class Console(cmd2.Cmd):
+    def _io_worker(self, loop):
+        self.terminal_lock.acquire()
         asyncio.set_event_loop(loop)
         try:
             loop.run_forever()
         finally:
             self.terminal_lock.release()
 
-    async def ainput(self, prompt):
+    async def AsyncInput(self, prompt):
-        async def ReadInput(prompt):
+        async def _input(prompt):
             return self._read_command_line(prompt)
-        future = asyncio.run_coroutine_threadsafe(ReadInput(prompt), self.ioLoop)
+        future = asyncio.run_coroutine_threadsafe(_input(prompt), self.ioLoop)
         return await asyncio.wrap_future(future)
 
-    async def aoutput(self, output):
+    async def AsyncPrint(self, output):
-        async def PrintOuput(output):
+        async def _print(output):
             print(output)
-        future = asyncio.run_coroutine_threadsafe(PrintOuput(output), self.ioLoop)
+        future = asyncio.run_coroutine_threadsafe(_print(output), self.ioLoop)
         await asyncio.wrap_future(future)
 
-    def ParserProvider(self):
-        return cmd2.Cmd2ArgumentParser()
-
-    def AddPathParser(parserProvider, nargs = "?"):
-        def PathParserProvider(self):
-            parser = parserProvider(self)
-            parser.add_argument("path", help="path", default="", nargs=nargs, type=RunSyncInLoop(self.loop)(self.client.PathToNode))
-            return parser
-        return PathParserProvider
-    
-    def AddFatherAndSonParser(parserProvider):
-        def PathParserProvider(self):
-            parser = parserProvider(self)
-            parser.add_argument("path", help="path", default="", nargs="?", type=RunSyncInLoop(self.loop)(self.client.PathToFatherNodeAndNodeName))
-            return parser
-        return PathParserProvider
-
-    def AddUrlParser(parserProvider):
-        def PathParserProvider(self):
-            parser = parserProvider(self)
-            parser.add_argument("url", help="url")
-            return parser
-        return PathParserProvider
-    
-    def AddUsernamePasswordParser(parserProvider):
-        def PathParserProvider(self):
-            parser = parserProvider(self)
-            parser.add_argument("username", help="username", nargs="?")
-            parser.add_argument("password", help="password", nargs="?")
-            return parser
-        return PathParserProvider
-
-    async def PathCompleter(self, text, line, begidx, endidx, filterFiles):   
-        father, sonName = await self.client.PathToFatherNodeAndNodeName(text)
-        if not IsDir(father):
-            return []
-        matches = []
-        matchesNode = []
-        for childId in father.childrenId:
-            child = self.client.GetNodeById(childId)
-            if filterFiles and IsFile(child):
-                continue
-            if child.name.startswith(sonName):
-                self.display_matches.append(child.name)
-                if sonName == "":
-                    matches.append(text + child.name)
-                elif text.endswith(sonName):
-                    matches.append(text[:text.rfind(sonName)] + child.name)
-                matchesNode.append(child)
-        if len(matchesNode) == 1 and IsDir(matchesNode[0]):
-            matches[0] += "/"
-            self.allow_appended_space = False
-            self.allow_closing_quote = False
-        return matches
-
     def __init__(self):
         super().__init__()
-        self._setup_logging()
-        self.client = PKFs("token.json", proxy="http://127.0.0.1:7890")
 
-    async def Run(self):
+    def preloop(self):
         # 1. 设置忽略SIGINT
         import signal
         def signal_handler(sig, frame):
             pass
         signal.signal(signal.SIGINT, signal_handler)
 
-        # 2. 创建一个新的事件循环
+        # 2. 创建IO线程处理输入输出
-        self.loop = asyncio.get_running_loop()
         self.ioLoop = asyncio.new_event_loop()
-        thread = threading.Thread(target=self.IOWorker, args=(self.ioLoop,))
+        self.ioThread = threading.Thread(target=self._io_worker, args=(self.ioLoop,))
-        thread.start()
+        self.ioThread.start()
 
-        # 3. 启动cmd2
+        # 3. 设置console
-        saved_readline_settings = None
+        self.saved_readline_settings = None
-        try:
-            # Get sigint protection while we set up readline for cmd2
         with self.sigint_protection:
             saved_readline_settings = self._set_up_cmd2_readline()
     
-            stop = False
+    def postloop(self):
-            while not stop:
+        # 1. 还原console设置
-                # Get sigint protection while we read the command line
-                line = await self.ainput(self.prompt)
-                # Run the command along with all associated pre and post hooks
-                stop = self.onecmd_plus_hooks(line)
-        finally:
-            # Get sigint protection while we restore readline settings
         with self.sigint_protection:
-                if saved_readline_settings is not None:
+            if self.saved_readline_settings is not None:
-                    self._restore_readline(saved_readline_settings)
+                self._restore_readline(self.saved_readline_settings)
+        
+        # 2. 停止IO线程
         # https://stackoverflow.com/questions/51642267/asyncio-how-do-you-use-run-forever
         self.ioLoop.call_soon_threadsafe(self.ioLoop.stop)
-            thread.join()
+        self.ioThread.join()
     
     # commands #
-
     def do_debug(self, args):
         """
         Enable debug mode
@@ -195,59 +118,89 @@ class PikpakConsole(cmd2.Cmd):
         logging.getLogger().setLevel(logging.INFO)
         logging.info("Debug mode disabled")
 
+    login_parser = cmd2.Cmd2ArgumentParser()
+    login_parser.add_argument("username", help="username")
+    login_parser.add_argument("password", help="password")
     @RunSync
-    @ProvideDecoratorSelfArgs(cmd2.with_argparser, AddUsernamePasswordParser(ParserProvider))
+    @cmd2.with_argparser(login_parser)
     async def do_login(self, args):
         """
         Login to pikpak
         """
-        await self.client.Login(args.username, args.password)
+        await Client.Login(args.username, args.password)
-        await self.aoutput("Logged in successfully")
+        await self.AsyncPrint("Logged in successfully")
+
+    async def _path_completer(self, text, line, begidx, endidx, filterfiles):   
+        father, sonName = await Client.PathToFatherNodeAndNodeName(text)
+        if not IsDir(father):
+            return []
+        matches = []
+        matchesNode = []
+        for childId in father.childrenId:
+            child = Client.GetNodeById(childId)
+            if filterfiles and IsFile(child):
+                continue
+            if child.name.startswith(sonName):
+                self.display_matches.append(child.name)
+                if sonName == "":
+                    matches.append(text + child.name)
+                elif text.endswith(sonName):
+                    matches.append(text[:text.rfind(sonName)] + child.name)
+                matchesNode.append(child)
+        if len(matchesNode) == 1 and IsDir(matchesNode[0]):
+            matches[0] += "/"
+            self.allow_appended_space = False
+            self.allow_closing_quote = False
+        return matches
 
     @RunSync
     async def complete_ls(self, text, line, begidx, endidx):
-        return await self.PathCompleter(text, line, begidx, endidx, filterFiles = False)
+        return await self._path_completer(text, line, begidx, endidx, False)
 
+    ls_parser = cmd2.Cmd2ArgumentParser()
+    ls_parser.add_argument("path", help="path", default="", nargs="?", type=RunSync(Client.PathToNode))
     @RunSync
-    @ProvideDecoratorSelfArgs(cmd2.with_argparser, AddPathParser(ParserProvider))
+    @cmd2.with_argparser(ls_parser)
     async def do_ls(self, args):
         """
         List files in a directory
         """
         node = args.path
         if node is None:
-            await self.aoutput("Invalid path")
+            await self.AsyncPrint("Invalid path")
             return
-        await self.client.Refresh(node)
+        await Client.Refresh(node)
         if IsDir(node):
             for childId in node.childrenId:
-                child = self.client.GetNodeById(childId)
+                child = Client.GetNodeById(childId)
-                await self.aoutput(child.name)
+                await self.AsyncPrint(child.name)
         elif IsFile(node):
-            await self.aoutput(f"{node.name}: {node.url}")            
+            await self.AsyncPrint(f"{node.name}: {node.url}")            
     
     @RunSync
     async def complete_cd(self, text, line, begidx, endidx):
-        return await self.PathCompleter(text, line, begidx, endidx, filterFiles = True)
+        return await self._path_completer(text, line, begidx, endidx, True)
 
+    cd_parser = cmd2.Cmd2ArgumentParser()
+    cd_parser.add_argument("path", help="path", default="", nargs="?", type=RunSync(Client.PathToNode))
     @RunSync
-    @ProvideDecoratorSelfArgs(cmd2.with_argparser, AddPathParser(ParserProvider))
+    @cmd2.with_argparser(cd_parser)
     async def do_cd(self, args):
         """
         Change directory
         """
         node = args.path
         if not IsDir(node):
-            await self.aoutput("Invalid directory")
+            await self.AsyncPrint("Invalid directory")
             return
-        self.client.currentLocation = node
+        Client.currentLocation = node
     
     @RunSync
     async def do_cwd(self, args):
         """
         Print current working directory
         """
-        await self.aoutput(self.client.NodeToPath(self.client.currentLocation))
+        await self.AsyncPrint(Client.NodeToPath(Client.currentLocation))
 
     def do_clear(self, args):
         """
@@ -257,50 +210,69 @@ class PikpakConsole(cmd2.Cmd):
 
     @RunSync
     async def complete_rm(self, text, line, begidx, endidx):
-        return await self.PathCompleter(text, line, begidx, endidx, filterFiles = False)
+        return await self._path_completer(text, line, begidx, endidx, False)
 
+    rm_parser = cmd2.Cmd2ArgumentParser()
+    rm_parser.add_argument("paths", help="paths", default="", nargs="+", type=RunSync(Client.PathToNode))
     @RunSync
-    @ProvideDecoratorSelfArgs(cmd2.with_argparser, AddPathParser(ParserProvider, "+"))
+    @cmd2.with_argparser(rm_parser)
     async def do_rm(self, args):
         """
         Remove a file or directory
         """
-        await self.client.Delete(args.path)
+        await Client.Delete(args.paths)
 
     @RunSync
     async def complete_mkdir(self, text, line, begidx, endidx):
-        return await self.PathCompleter(text, line, begidx, endidx, filterFiles = True)
+        return await self._path_completer(text, line, begidx, endidx, True)
 
+    mkdir_parser = cmd2.Cmd2ArgumentParser()
+    mkdir_parser.add_argument("path_and_son", help="path and son", default="", nargs="?", type=RunSync(Client.PathToFatherNodeAndNodeName))
     @RunSync
-    @ProvideDecoratorSelfArgs(cmd2.with_argparser, AddFatherAndSonParser(ParserProvider))
+    @cmd2.with_argparser(mkdir_parser)
     async def do_mkdir(self, args):
         """
         Create a directory
         """
-        father, sonName = args.path
+        father, sonName = args.path_and_son
         if not IsDir(father) or sonName == "" or sonName == None:
-            await self.aoutput("Invalid path")
+            await self.AsyncPrint("Invalid path")
             return
-        child = self.client.FindChildInDirByName(father, sonName)
+        child = Client.FindChildInDirByName(father, sonName)
         if child is not None:
-            await self.aoutput("Path already exists")
+            await self.AsyncPrint("Path already exists")
             return
-        await self.client.MakeDir(father, sonName)
+        await Client.MakeDir(father, sonName)
 
+    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
-    @ProvideDecoratorSelfArgs(cmd2.with_argparser, AddPathParser(AddUrlParser(ParserProvider)))
+    @cmd2.with_argparser(download_parser)
     async def do_download(self, args):
         """
         Download a file
         """
         node = args.path
         if not IsDir(node):
-            await self.aoutput("Invalid directory")
+            await self.AsyncPrint("Invalid directory")
             return
-        await self.client.Download(args.url, node)
+        await Client.Download(args.url, node)
     
+async def mainLoop():
+    global MainLoop, Client
+    MainLoop = asyncio.get_running_loop()
+    
+    console = Console()
+    console.preloop()
+    try:
+        stop = False
+        while not stop:
+            line = await console.AsyncInput(console.prompt)
+            stop = console.onecmd_plus_hooks(line)
+    finally:
+        console.postloop()
 
 if __name__ == "__main__":  
     nest_asyncio.apply()
-    prog = PikpakConsole()
+    asyncio.run(mainLoop())
-    asyncio.run(prog.Run())

pikpakFs.py

@@ -1,13 +1,12 @@
 import httpx
-from hashlib import md5
 from pikpakapi import PikPakApi
 from typing import Dict
 from datetime import datetime
 import json
-import re
 import os
 import logging
 from enum import Enum
+import asyncio
 
 
 class PKTaskStatus(Enum):
@@ -24,6 +23,7 @@ class PkTask:
         self.taskId = id
         self.status = status
         
+        self.runningTask : asyncio.Task = None
         self.toDirId = toDirId
         self.torrent = torrent
         self.url = None
@@ -86,9 +86,14 @@ class PkToken:
         return cls(**data)
 
 class PKFs:
-    async def RunTasks(self):
+    async def _task_worker(self, task : PkTask):
         pass
 
+    def RunTask(self, task : PkTask):
+        self.tasks.append(task)
+        if task.runningTask is None or task.runningTask.done():
+            task.runningTask = asyncio.create_task(self._task_worker(task))
+
     def __init__(self, loginCachePath : str = None, proxy : str = None, rootId = None):
         self.nodes : Dict[str, FsNode] = {} 
         self.root = DirNode(rootId, "", None)
@@ -275,9 +280,14 @@ class PKFs:
 
     async def Download(self, url : str, dirNode : DirNode) -> PkTask :
         task = PkTask(url, dirNode.id)
-        self.tasks.append(task)
+        self.RunTask(task)
         return task
 
+    async def QueryTasks(self, filterByStatus : PKTaskStatus = None) -> list[PkTask]:
+        if filterByStatus is None:
+            return self.tasks
+        return [task for task in self.tasks if task.status == filterByStatus]
+
     async def Delete(self, nodes : list[FsNode]) -> None:
         nodeIds = [node.id for node in nodes]
         await self.client.delete_to_trash(nodeIds)