132 lines
4.3 KiB
Python
132 lines
4.3 KiB
Python
"""
|
|
Utility module for handling command contexts in Discord.py.
|
|
|
|
This module provides functions for setting up commands with different contexts,
|
|
allowing them to work in DMs, private channels, guilds, or any combination.
|
|
"""
|
|
|
|
import discord
|
|
from discord import app_commands
|
|
from enum import Enum, auto
|
|
from typing import Optional, List, Dict, Any, Union, Callable, Awaitable
|
|
|
|
class AppCommandContext(Enum):
|
|
"""
|
|
Enum representing the allowed contexts for application commands.
|
|
|
|
This enum defines where commands can be used:
|
|
- GUILD: Commands can only be used in guilds (servers)
|
|
- GUILD_INSTALL: Commands can only be used in guilds where the app is installed
|
|
- PRIVATE: Commands can only be used in private contexts (DMs and private channels)
|
|
- ALL: Commands can be used in both guilds and private contexts
|
|
"""
|
|
GUILD = auto()
|
|
GUILD_INSTALL = auto()
|
|
PRIVATE = auto()
|
|
ALL = auto()
|
|
|
|
def set_command_context(
|
|
command: app_commands.Command,
|
|
context: AppCommandContext = AppCommandContext.ALL
|
|
) -> app_commands.Command:
|
|
"""
|
|
Set the context for a command, determining where it can be used.
|
|
|
|
Args:
|
|
command: The command to modify
|
|
context: The context to set for the command
|
|
|
|
Returns:
|
|
The modified command
|
|
"""
|
|
# Set guild_only and dm_permission based on the context
|
|
if context == AppCommandContext.GUILD or context == AppCommandContext.GUILD_INSTALL:
|
|
# Guild-only commands
|
|
command.guild_only = True
|
|
# dm_permission is automatically set to False when guild_only is True
|
|
elif context == AppCommandContext.PRIVATE:
|
|
# Private-only commands (DMs and private channels)
|
|
command.guild_only = False
|
|
# Set dm_permission to True
|
|
if not hasattr(command, "extras"):
|
|
command.extras = {}
|
|
command.extras["dm_permission"] = True
|
|
# Disable in guilds by setting default_member_permissions to "0"
|
|
command.extras["default_member_permissions"] = "0"
|
|
elif context == AppCommandContext.ALL:
|
|
# Commands that work everywhere
|
|
command.guild_only = False
|
|
# Set dm_permission to True
|
|
if not hasattr(command, "extras"):
|
|
command.extras = {}
|
|
command.extras["dm_permission"] = True
|
|
# Allow default permissions in guilds
|
|
|
|
return command
|
|
|
|
def create_global_command(
|
|
bot_tree: app_commands.CommandTree,
|
|
name: str,
|
|
description: str,
|
|
callback: Callable[[discord.Interaction], Awaitable[None]],
|
|
context: AppCommandContext = AppCommandContext.ALL,
|
|
**kwargs
|
|
) -> app_commands.Command:
|
|
"""
|
|
Create a command that can be used globally with the specified context.
|
|
|
|
Args:
|
|
bot_tree: The command tree to add the command to
|
|
name: The name of the command
|
|
description: The description of the command
|
|
callback: The function to call when the command is used
|
|
context: The context to set for the command
|
|
**kwargs: Additional arguments to pass to the command constructor
|
|
|
|
Returns:
|
|
The created command
|
|
"""
|
|
# Create the command
|
|
command = app_commands.Command(
|
|
name=name,
|
|
description=description,
|
|
callback=callback,
|
|
**kwargs
|
|
)
|
|
|
|
# Set the context
|
|
set_command_context(command, context)
|
|
|
|
# Add the command to the tree
|
|
bot_tree.add_command(command)
|
|
|
|
return command
|
|
|
|
def sync_commands(
|
|
bot_tree: app_commands.CommandTree,
|
|
guild_specific: bool = True,
|
|
global_commands: bool = True
|
|
) -> Awaitable[List[app_commands.Command]]:
|
|
"""
|
|
Sync commands with Discord.
|
|
|
|
Args:
|
|
bot_tree: The command tree to sync
|
|
guild_specific: Whether to sync guild-specific commands
|
|
global_commands: Whether to sync global commands
|
|
|
|
Returns:
|
|
A coroutine that resolves to the list of synced commands
|
|
"""
|
|
if guild_specific and not global_commands:
|
|
# Only sync guild-specific commands
|
|
# This is handled by command_customization.register_all_guild_commands
|
|
return None
|
|
elif not guild_specific and global_commands:
|
|
# Only sync global commands
|
|
return bot_tree.sync()
|
|
else:
|
|
# Sync both guild-specific and global commands
|
|
# This is the default behavior of bot_tree.sync()
|
|
return bot_tree.sync()
|