diff --git a/main.py b/main.py index ecde4a4..c426cf2 100644 --- a/main.py +++ b/main.py @@ -4,183 +4,106 @@ 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 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 +def setup_logging(): + formatter = colorlog.ColoredFormatter( + "%(log_color)s%(asctime)s - %(levelname)s - %(name)s - %(message)s", + datefmt='%Y-%m-%d %H:%M:%S', + reset=True, + log_colors={ + 'DEBUG': 'cyan', + 'INFO': 'green', + 'WARNING': 'yellow', + 'ERROR': 'red', + 'CRITICAL': 'red,bg_white', + } + ) + handler = logging.StreamHandler() + handler.setFormatter(formatter) + handler.setLevel(logging.INFO) + + file_handler = logging.FileHandler('app.log') + file_handler.setFormatter(formatter) + file_handler.setLevel(logging.DEBUG) + + logger = logging.getLogger() + logger.addHandler(handler) + logger.addHandler(file_handler) + logger.setLevel(logging.DEBUG) + +setup_logging() +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 -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', - reset=True, - log_colors={ - 'DEBUG': 'cyan', - 'INFO': 'green', - 'WARNING': 'yellow', - 'ERROR': 'red', - 'CRITICAL': 'red,bg_white', - } - ) - handler = logging.StreamHandler() - handler.setFormatter(formatter) - handler.setLevel(logging.INFO) - - file_handler = logging.FileHandler('app.log') - file_handler.setFormatter(formatter) - file_handler.setLevel(logging.DEBUG) - - logger = logging.getLogger() - logger.addHandler(handler) - logger.addHandler(file_handler) - logger.setLevel(logging.DEBUG) - - def IOWorker(self, loop): - self.terminal_lock.acquire() # 我看cmdloop是这么做的，所以我也这么做 +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 ReadInput(prompt): + async def AsyncInput(self, 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 PrintOuput(output): + async def AsyncPrint(self, 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. 创建一个新的事件循环 - self.loop = asyncio.get_running_loop() + # 2. 创建IO线程处理输入输出 self.ioLoop = asyncio.new_event_loop() - thread = threading.Thread(target=self.IOWorker, args=(self.ioLoop,)) - thread.start() + self.ioThread = threading.Thread(target=self._io_worker, args=(self.ioLoop,)) + self.ioThread.start() - # 3. 启动cmd2 - 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 - while not stop: - # 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: - self._restore_readline(saved_readline_settings) - # https://stackoverflow.com/questions/51642267/asyncio-how-do-you-use-run-forever - self.ioLoop.call_soon_threadsafe(self.ioLoop.stop) - thread.join() + # 3. 设置console + self.saved_readline_settings = None + with self.sigint_protection: + saved_readline_settings = self._set_up_cmd2_readline() + + def postloop(self): + # 1. 还原console设置 + with self.sigint_protection: + if self.saved_readline_settings is not None: + 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) + 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 self.aoutput("Logged in successfully") + await Client.Login(args.username, args.password) + 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) - await self.aoutput(child.name) + child = Client.GetNodeById(childId) + 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(prog.Run()) + asyncio.run(mainLoop()) diff --git a/pikpakFs.py b/pikpakFs.py index 5fedcf7..d57a87b 100644 --- a/pikpakFs.py +++ b/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)