优化cmd2结构
This commit is contained in:
fengfeng
parent
6f78ceac53
commit
2e0e9425fc
326
main.py
326
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()
|
||||
# 3. 设置console
|
||||
self.saved_readline_settings = None
|
||||
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()
|
||||
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())
|
||||
|
||||
18
pikpakFs.py
18
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)
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user