Building a Discord Chatbot with Python (7) - Refactoring Code Using Discord.py's Extension

In this article, we’ll leverage the discord.ext.commands extension in discord.py
to refactor our somewhat complex on_message function.

Recap

The code up to our last session is as follows:
We added !omikuji and !quiz commands, which made the on_message function somewhat bloated.

import asyncio
import json
import random

import discord


class MyClient(discord.Client):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        quiz_file = 'python_problems.en.json'
        with open(quiz_file) as f:
            self.quiz_list = json.load(f)

    async def on_ready(self):
        """Triggered when connecting to Discord."""
        print(f'Connected as {self.user}.')

    async def on_message(self, message: discord.Message):
        """Triggered when receiving a message."""
        # Ignore messages sent by the bot itself
        if message.author == self.user:
            return

        print(f'Received a message from {message.author}: {message.content}')

        if message.content == '!omikuji':
            choice = random.choice([
                "Great Blessing", "Blessing", "Small Blessing",
                "Misfortune", "Great Misfortune"
            ])
            await message.channel.send(
                f"Your fortune for today is: **{choice}**!"
            )

        if message.content == '!quiz':
            def quiz_reply_check(m: discord.Message):
                """Function to validate quiz replies."""
                return m.author == message.author and m.content.isdigit()

            quiz = random.choice(self.quiz_list)
            quiz_str = (
                f'Q: {quiz["question"]}\n\n' +
                '\n'.join([
                    f'[{i+1}] {c}' for (i, c) in enumerate(quiz["choices"])
                ])
            )
            await message.channel.send(quiz_str)
            try:
                reply_message = await self.wait_for(
                    'message', check=quiz_reply_check, timeout=30
                )
                if reply_message.content == str(quiz['answer_num']):
                    result = "😀 Correct!"
                else:
                    result = \
                        f"😢 Sorry, the correct answer is {quiz['answer_num']}."
            except asyncio.TimeoutError:
                result = \
                    f"😢 Time's up! The correct answer is {quiz['answer_num']}."

            await message.channel.send(result)


with open('.discord_token') as f:
    token = f.read().strip()

intents = discord.Intents.default()
intents.message_content = True
client = MyClient(intents=intents)
client.run(token)

Using the discord.ext.commands extension

As we add more commands like before, repetitive logic increases, and our on_message becomes bulky.

That’s where discord.py provides an extension library named discord.ext.commands
to make command additions more concise and simpler.
(Reference: API Reference, API Reference 2)

Now, let’s go ahead and refactor our code using this command extension.

1. Inherit from commands.Bot

Initially, we inherited from discord.Client to create our custom client MyClient.
This will now be changed to inherit from commands.Bot.

...
from discord.ext import commands


class MyClient(commands.Bot):
    ...

2. Modifying the second argument for the __init__ function

With command.Bot, the __init__ function has an additional argument named command_prefix.
We’ll adjust for this. By default, we’ll set it to !.

class MyClient(commands.Bot):
    def __init__(self, command_prefix='!', *args, **kwargs):
        super().__init__(command_prefix, *args, **kwargs)
        ...

3. Separate command definitions

Next, we’ll separate the command definitions from on_message and define them as individual functions.
The separated functions will have a first argument of type commands.Context.

For example, for the !omikuji command, it would look like this:

@commands.command(name='omikuji')
async def omikuji_command(ctx: commands.Context):
    """Pick a fortune slip."""
    choice = random.choice([
        "Great Blessing", "Blessing", "Small Blessing",
        "Misfortune", "Great Misfortune"
    ])
    await ctx.channel.send(
        f"Your fortune for today is: **{choice}**!"
    )

• Attach the @commands.command() decorator to the function. Specify the command name with name.
• The function’s first argument will be ctx: commands.Context. ctx represents the context and can send messages to the current channel with ctx.channel.send. For other attributes, refer to the API Reference.
• The docstring content is displayed in the command help.

Similarly, when separating the !quiz command, it looks like this:

@commands.command(name='quiz')
async def quiz_command(ctx: commands.Context):
    """Pose a three-choice quiz"""
    def quiz_reply_check(m: discord.Message):
        """Function to validate quiz replies."""
        return m.author == ctx.author and m.content.isdigit()
    quiz = random.choice(ctx.bot.quiz_list)
    quiz_str = (
        f'Q: {quiz["question"]}\n\n' +
        '\n'.join([
            f'[{i+1}] {c}' for (i, c) in enumerate(quiz["choices"])
        ])
    )
    await ctx.channel.send(quiz_str)
    try:
        reply_message = await ctx.bot.wait_for(
            'message', check=quiz_reply_check, timeout=30
        )
        if reply_message.content == str(quiz['answer_num']):
            result = "😀 Correct!"
        else:
            result = \
                f"😢 Sorry, the correct answer is {quiz['answer_num']}."
    except asyncio.TimeoutError:
        result = \
            f"😢 Time's up! The correct answer is {quiz['answer_num']}."
    await ctx.channel.send(result)

• ctx.author is the sender of the message.
• Other sections previously marked with self have been changed to ctx.bot.

4. Adding commands with add_command

Add the separated commands to the bot as shown below:

client = MyClient(intents=intents)
bot_commands = [omikuji_command, quiz_command]
for command in bot_commands:
    client.add_command(command)
client.run(token)

5. Execute the inherited on_message within on_message

Finally, modify the on_message function to execute the inherited commands.Bot’s on_message as shown:

async def on_message(self, message: discord.Message):
    """Triggered when receiving a message."""
    # Ignore messages sent by the bot itself
    if message.author == self.user:
        return

    print(f'Received a message from {message.author}: {message.content}')
    await super().on_message(message)

6. Code after adjustments

The code after making the above adjustments looks like this:

import asyncio
import json
import random

import discord
from discord.ext import commands


class MyClient(commands.Bot):
    def __init__(self, command_prefix='!', *args, **kwargs):
        super().__init__(command_prefix, *args, **kwargs)
        quiz_file = 'python_problems.en.json'
        with open(quiz_file) as f:
            self.quiz_list = json.load(f)

    async def on_ready(self):
        """Triggered when connecting to Discord."""
        print(f'Connected as {self.user}.')

    async def on_message(self, message: discord.Message):
        """Triggered when receiving a message."""
        # Ignore messages sent by the bot itself
        if message.author == self.user:
            return

        print(f'Received a message from {message.author}: {message.content}')
        await super().on_message(message)


@commands.command(name='omikuji')
async def omikuji_command(ctx: commands.Context):
    """Pick a fortune slip."""
    choice = random.choice([
        "Great Blessing", "Blessing", "Small Blessing",
        "Misfortune", "Great Misfortune"
    ])
    await ctx.channel.send(
        f"Your fortune for today is: **{choice}**!"
    )


@commands.command(name='quiz')
async def quiz_command(ctx: commands.Context):
    """Pose a three-choice quiz"""
    def quiz_reply_check(m: discord.Message):
        """Function to validate quiz replies."""
        return m.author == ctx.author and m.content.isdigit()
    quiz = random.choice(ctx.bot.quiz_list)
    quiz_str = (
        f'Q: {quiz["question"]}\n\n' +
        '\n'.join([
            f'[{i+1}] {c}' for (i, c) in enumerate(quiz["choices"])
        ])
    )
    await ctx.channel.send(quiz_str)
    try:
        reply_message = await ctx.bot.wait_for(
            'message', check=quiz_reply_check, timeout=30
        )
        if reply_message.content == str(quiz['answer_num']):
            result = "😀 Correct!"
        else:
            result = \
                f"😢 Sorry, the correct answer is {quiz['answer_num']}."
    except asyncio.TimeoutError:
        result = \
            f"😢 Time's up! The correct answer is {quiz['answer_num']}."
    await ctx.channel.send(result)



with open('.discord_token') as f:
    token = f.read().strip()

intents = discord.Intents.default()
intents.message_content = True
client = MyClient(intents=intents)
bot_commands = [omikuji_command, quiz_command]
for command in bot_commands:
    client.add_command(command)
client.run(token)

Splitting commands into separate files

Next, for easier management, we’ll split each of the separated command functions into individual files.
We’ll create a commands directory and separate the commands into files within it.

• commands/omikuji.py

  import random
  
  from discord.ext import commands
  
  
  @commands.command(name='omikuji')
  async def omikuji_command(ctx: commands.Context):
      """Pick a fortune slip."""
      choice = random.choice([
          "Great Blessing", "Blessing", "Small Blessing",
          "Misfortune", "Great Misfortune"
      ])
      await ctx.channel.send(
          f"Your fortune for today is: **{choice}**!"
      )
  

• commands/quiz.py

  import asyncio
  import random
  
  import discord
  from discord.ext import commands
  
  
  @commands.command(name='quiz')
  async def quiz_command(ctx: commands.Context):
      """Pose a three-choice quiz"""
      def quiz_reply_check(m: discord.Message):
          """Function to validate quiz replies."""
          return m.author == ctx.author and m.content.isdigit()
      quiz = random.choice(ctx.bot.quiz_list)
      quiz_str = (
          f'Q: {quiz["question"]}\n\n' +
          '\n'.join([
              f'[{i+1}] {c}' for (i, c) in enumerate(quiz["choices"])
          ])
      )
      await ctx.channel.send(quiz_str)
      try:
          reply_message = await ctx.bot.wait_for(
              'message', check=quiz_reply_check, timeout=30
          )
          if reply_message.content == str(quiz['answer_num']):
              result = "😀 Correct!"
          else:
              result = \
                  f"😢 Sorry, the correct answer is {quiz['answer_num']}."
      except asyncio.TimeoutError:
          result = \
              f"😢 Time's up! The correct answer is {quiz['answer_num']}."
      await ctx.channel.send(result)
  

• chatbot.py (Moved bot_commands to the top for easier customization.)

  import json
  
  import discord
  from discord.ext import commands
  
  from commands.omikuji import omikuji_command
  from commands.quiz import quiz_command
  
  
  BOT_COMMANDS = [omikuji_command, quiz_command]
  
  
  class MyClient(commands.Bot):
      def __init__(self, command_prefix='!', *args, **kwargs):
          super().__init__(command_prefix, *args, **kwargs)
  
          quiz_file = 'python_problems.en.json'
          with open(quiz_file) as f:
              self.quiz_list = json.load(f)
  
      async def on_ready(self):
          """Triggered when connecting to Discord."""
          print(f'Connected as {self.user}.')
  
      async def on_message(self, message: discord.Message):
          """Triggered when receiving a message."""
          # Ignore messages sent by the bot itself
          if message.author == self.user:
              return
  
          print(f'Received a message from {message.author}: {message.content}')
          await super().on_message(message)
  
  
  with open('.discord_token') as f:
      token = f.read().strip()
  
  intents = discord.Intents.default()
  intents.message_content = True
  client = MyClient(intents=intents)
  
  for command in BOT_COMMANDS:
      client.add_command(command)
  
  client.run(token)
  

Remove the quiz data loading from __init__

Having separated the functionalities into different files for refactoring,
I noticed the MyClient’s __init__ still reads the quiz data.

There are several ways to fix this, but this time, let’s remove the quiz file reading part from the __init__ and instead, load it when quiz.py is loaded. Here’s how:

import asyncio
import json
import random

import discord
from discord.ext import commands


quiz_file = 'python_problems.en.json'
with open(quiz_file) as f:
    quiz_list = json.load(f)


@commands.command(name='quiz')
async def quiz_command(ctx: commands.Context):
    """Pose a three-choice quiz"""

    def quiz_reply_check(m: discord.Message):
        """Function to validate quiz replies."""
        return m.author == ctx.author and m.content.isdigit()

    quiz = random.choice(quiz_list)
    quiz_str = (
        f'Q: {quiz["question"]}\n\n' +
        '\n'.join([
            f'[{i+1}] {c}' for (i, c) in enumerate(quiz["choices"])
        ])
    )
    await ctx.channel.send(quiz_str)
    try:
        reply_message = await ctx.bot.wait_for(
            'message', check=quiz_reply_check, timeout=30
        )
        if reply_message.content == str(quiz['answer_num']):
            result = "😀 Correct!"
        else:
            result = \
                f"😢 Sorry, the correct answer is {quiz['answer_num']}."
    except asyncio.TimeoutError:
        result = \
            f"😢 Time's up! The correct answer is {quiz['answer_num']}."
    await ctx.channel.send(result)

Let’s Test It Out

Finally, let’s verify if the commands work correctly.

Additionally, the !help command showcases the command list.

Conclusion

In this article, we refactored our code using discord.py’s extension, discord.ext.commands.
By splitting the functionalities into indivisual files, our code became more organized, and managing it became more convenient.

In the next session, we’ll dive deeper into the features of discord.ext.commands again and improve our quiz feature.

Code for this session

• chatbot.py

  import json
  
  import discord
  from discord.ext import commands
  
  from commands.omikuji import omikuji_command
  from commands.quiz import quiz_command
  
  
  BOT_COMMANDS = [omikuji_command, quiz_command]
  
  
  class MyClient(commands.Bot):
      def __init__(self, command_prefix='!', *args, **kwargs):
          super().__init__(command_prefix, *args, **kwargs)
  
      async def on_ready(self):
          """Triggered when connecting to Discord."""
          print(f'Connected as {self.user}.')
  
      async def on_message(self, message: discord.Message):
          """Triggered when receiving a message."""
          # Ignore messages sent by the bot itself
          if message.author == self.user:
              return
  
          print(f'Received a message from {message.author}: {message.content}')
          await super().on_message(message)
  
  
  with open('.discord_token') as f:
      token = f.read().strip()
  
  intents = discord.Intents.default()
  intents.message_content = True
  client = MyClient(intents=intents)
  
  for command in BOT_COMMANDS:
      client.add_command(command)
  
  client.run(token)
  

• commands/omikuji.py

  import random
  
  from discord.ext import commands
  
  
  @commands.command(name='omikuji')
  async def omikuji_command(ctx: commands.Context):
      """Pick a fortune slip."""
      choice = random.choice([
          "Great Blessing", "Blessing", "Small Blessing",
          "Misfortune", "Great Misfortune"
      ])
      await ctx.channel.send(
          f"Your fortune for today is: **{choice}**!"
      )
  

• commands/quiz.py

  import asyncio
  import json
  import random
  
  import discord
  from discord.ext import commands
  
  
  quiz_file = 'python_problems.en.json'
  with open(quiz_file) as f:
      quiz_list = json.load(f)
  
  
  @commands.command(name='quiz')
  async def quiz_command(ctx: commands.Context):
      """Pose a three-choice quiz"""
  
      def quiz_reply_check(m: discord.Message):
          """Function to validate quiz replies."""
          return m.author == ctx.author and m.content.isdigit()
  
      quiz = random.choice(quiz_list)
      quiz_str = (
          f'Q: {quiz["question"]}\n\n' +
          '\n'.join([
              f'[{i+1}] {c}' for (i, c) in enumerate(quiz["choices"])
          ])
      )
      await ctx.channel.send(quiz_str)
      try:
          reply_message = await ctx.bot.wait_for(
              'message', check=quiz_reply_check, timeout=30
          )
          if reply_message.content == str(quiz['answer_num']):
              result = "😀 Correct!"
          else:
              result = \
                  f"😢 Sorry, the correct answer is {quiz['answer_num']}."
      except asyncio.TimeoutError:
          result = \
              f"😢 Time's up! The correct answer is {quiz['answer_num']}."
      await ctx.channel.send(result)