Merge pull request #11 from Pagliacii/main

feat: add `--disable-command` option
This commit is contained in:
yihong 2023-12-16 16:30:04 +08:00 committed by GitHub
commit ea236d47a8
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 31 additions and 3 deletions

View File

@ -26,6 +26,9 @@ for yihong0618's channel: https://t.me/hyi0618
2. Get tg token, ask Google or ChatGPT, need get it from [BotFather](https://t.me/BotFather) 2. Get tg token, ask Google or ChatGPT, need get it from [BotFather](https://t.me/BotFather)
3. python tg.py ${tg_token} 3. python tg.py ${tg_token}
> [!Note]
> If you don't want to use one of these command, you can use `--disable-command <command>` option to disable it. This option can be used multiple times.
## Contribution ## Contribution

View File

@ -55,12 +55,14 @@ def wrap_handler(handler: T, bot: TeleBot) -> T:
return update_wrapper(wrapper, handler) return update_wrapper(wrapper, handler)
def load_handlers(bot: TeleBot) -> None: def load_handlers(bot: TeleBot, disable_commands: list[str]) -> None:
# import all submodules # import all submodules
this_path = Path(__file__).parent this_path = Path(__file__).parent
for child in this_path.iterdir(): for child in this_path.iterdir():
if child.name.startswith("_"): if child.name.startswith("_"):
continue continue
if child.stem in disable_commands:
continue
module = importlib.import_module(f".{child.stem}", __package__) module = importlib.import_module(f".{child.stem}", __package__)
if hasattr(module, "register"): if hasattr(module, "register"):
print(f"Loading {child.stem} handlers.") print(f"Loading {child.stem} handlers.")
@ -78,3 +80,13 @@ def load_handlers(bot: TeleBot) -> None:
if all_commands: if all_commands:
bot.set_my_commands(all_commands) bot.set_my_commands(all_commands)
print("Setting commands done.") print("Setting commands done.")
def list_available_commands() -> list[str]:
commands = []
this_path = Path(__file__).parent
for child in this_path.iterdir():
if child.name.startswith("_"):
continue
commands.append(child.stem)
return commands

17
tg.py
View File

@ -2,19 +2,32 @@ import argparse
from telebot import TeleBot from telebot import TeleBot
from handlers import load_handlers from handlers import list_available_commands, load_handlers
def main(): def main():
# Init args # Init args
parser = argparse.ArgumentParser() parser = argparse.ArgumentParser()
parser.add_argument("tg_token", help="tg token") parser.add_argument("tg_token", help="tg token")
# 'disable-command' option
# The action 'append' will allow multiple entries to be saved into a list
# The variable name is changed to 'disable_commands'
parser.add_argument(
"--disable-command",
action="append",
dest="disable_commands",
help="Specify a command to disable. Can be used multiple times.",
default=[],
choices=list_available_commands(),
)
options = parser.parse_args() options = parser.parse_args()
print("Arg parse done.") print("Arg parse done.")
# Init bot # Init bot
bot = TeleBot(options.tg_token) bot = TeleBot(options.tg_token)
load_handlers(bot) load_handlers(bot, options.disable_commands)
print("Bot init done.") print("Bot init done.")
# Start bot # Start bot