13. Cog 사용해보기1

chaejm55·2021년 9월 15일
0

디스코드봇

목록 보기
17/18
post-thumbnail

이제서야 다시 찾아 뵙게 된 점 죄송합니다. 다만 앞으로는 예전처럼 3일 주기가 아닌 자유 주기로 포스팅을 할 예정입니다. 양해부탁드립니다.

0. 들어가기 전에

굉장히 오랜만에 쓰는 디스코드봇 포스팅이다. 무슨 주제로 포스팅을 할까 고민하다가 현재까지 작성한 명령어들을 카테고리별로 나누도록 해주는 Cog에 대해 다뤄보려한다. Cog를 통해 코드가 더욱 깔끔해지도록 리팩터링해보자.

1. Cog

현재까지는 그냥 bot.py에 모든 명령어를 작성했는데, cog를 사용하여 이를 각각 클래스를 통해 카테고리별로 나눌 수 있다.
먼저 기본 작성법을 알아보자. 공식 문서 예시

새로운 파일 Game.py를 만들어서 작성해보자.
Game.py

from discord.ext import commands  # Cog 사용을 위한 라이브러리 import

import random


class Game(commands.Cog):  # Cog를 상속하는 클래스를 선언
    
    def __init__(self, bot):  # 생성자 작성
        self.bot = bot

    @commands.command()  # 이전의 @bot.command() 대신 사용
    async def dice(self, ctx):  # 클래스 내부의 함수이므로 self 파라미터가 필수
        randnum = random.randint(1, 6)
        await ctx.send(f'주사위 결과는 {randnum} 입니다.')


def setup(bot):  # setup 함수로 cog를 추가한다. 
    bot.add_cog(Game(bot))

Game이라는 카테고리(클래스)를 만들어 이전에 작성하였던 커맨드 하나를 옮겨보았다. 이렇게 새로운 커맨드 카테고리를 만들어 코드를 더욱 깔끔하게 만들 수 있다.
이와 비슷한 형태로 나머지 커맨드들도 카테고리별로 나누어보자.

2. Cog로 bot.py 나눠 보기

Cog를 통해 현재까지 작성한 커맨드들을 카테코리에 따라 나눠보자. 해당 코드들은 cogs 디렉토리에 저장된다. 코드 전체를 작성하니 간단히 보고 싶으면 아래의 요약만 봐도 된다.

1) Game.py

이 카테고리에 속할 커맨드들은 디스코드 봇과 간단한 게임을 하는 커맨드들이다. 목록은 다음과 같다.

  • dice, mining, game, game_board

Game.py

import random
import os
from discord.ext import commands
from discord.ext.commands import MissingRequiredArgument


# game 커맨드가 바로 호출할 수 있도록 클래스 밖에 선언

async def make_dir(directory_name):
    try:
        if not os.path.exists(directory_name):
            os.makedirs(directory_name)
    except OSError:
        print('Error: makedirs()')


async def add_result(directory_name, user_name, result):
    file_path = directory_name + '/' + user_name + '.txt'
    if os.path.exists(file_path):
        with open(file_path, 'a', encoding='UTF-8') as f:
            f.write(result)
    else:
        with open(file_path, 'w', encoding='UTF-8') as f:
            f.write(result)


class Game(commands.Cog):

    def __init__(self, bot):
        self.bot = bot

    @commands.command()
    async def dice(self, ctx):
        randnum = random.randint(1, 6)
        await ctx.send(f'주사위 결과는 {randnum} 입니다.')

    @commands.command()
    async def mining(self, ctx):
        minerals = ['다이아몬드', '루비', '에메랄드', '자수정', '철', '석탄']
        weights = [1, 3, 6, 15, 25, 50]
        results = random.choices(minerals, weights=weights, k=5)
        await ctx.send(', '.join(results) + ' 광물들을 획득하였습니다.')

    @commands.command()
    async def game(self, ctx, user: str):
        rps_table = ['가위', '바위', '보']
        bot = random.choice(rps_table)
        result = rps_table.index(user) - rps_table.index(bot)
        if result == 0:
            result_text = f'{user} vs {bot} 비김'
            await ctx.send(f'{user} vs {bot}  비겼습니다.')
        elif result == 1 or result == -2:
            result_text = f'{user} vs {bot} 승리!'
            await ctx.send(f'{user} vs {bot}  유저가 이겼습니다.')
        else:
            result_text = f'{user} vs {bot} 패배...'
            await ctx.send(f'{user} vs {bot}  봇이 이겼습니다.')

        directory_name = "game_result"
        await make_dir(directory_name)
        await add_result(directory_name, str(ctx.author), result_text + '\n')

    @game.error  # @<명령어>.error의 형태로 된 데코레이터를 사용한다.
    async def game_error(self, ctx, error):  # 파라미터에 ctx, error를 필수로 한다.
        if isinstance(error, MissingRequiredArgument):  # isinstance로 에러에 따라 시킬 작업을 결정한다.
            await ctx.send("가위/바위/보 중 낼 것을 입력해주세요.")

    @commands.command(name="전적")
    async def game_board(self, ctx):
        user_name = str(ctx.author)
        file_path = "game_result/" + user_name + ".txt"
        if os.path.exists(file_path):
            with open(file_path, "r", encoding="UTF-8") as f:
                result = f.read()
            await ctx.send(f'{ctx.author}님의 가위바위보 게임 전적입니다.\n==============================\n' + result)
        else:
            await ctx.send(f'{ctx.author}님의 가위바위보 전적이 존재하지 않습니다.')


def setup(bot):
    bot.add_cog(Game(bot))

2) User.py

이 카테고리에 속할 커맨드들은 채널의 유저 관리와 관련된 커맨드들이다. 목록은 다음과 같다.

  • kick_user, ban_user, unban_user, role_user

User.py

from discord.ext import commands
import discord


class User(commands.Cog):

    def __init__(self, bot):
        self.bot = bot

    @commands.command(aliases=['추방'])
    async def kick_user(self, ctx, nickname: discord.Member):
        await nickname.kick()
        await ctx.send(f"{nickname} 님이 추방되었습니다.")

    @commands.command(aliases=['차단'])
    async def ban_user(self, ctx, nickname: discord.Member):
        await nickname.ban()
        await ctx.send(f"{nickname} 님이 차단되었습니다.")

    @commands.command(aliases=['해제'])
    async def unban_user(self, ctx, nickname: str):
        ban_entry = await ctx.guild.bans()
        for users in ban_entry:
            if nickname == users.user.name:
                forgive_user = users.user
                await ctx.guild.unban(forgive_user)
                return await ctx.send(f"{nickname} 님이 차단 해제되었습니다.")
        return await ctx.send(f"{nickname} 님은 차단 목록에 없습니다.")

    @commands.command(aliases=['역할부여'])
    async def role_user(self, ctx, nickname: discord.Member, role_name):
        roles = ctx.guild.roles
        for role in roles:
            if role_name == role.name:
                await nickname.add_roles(role)
                return await ctx.send(f"{nickname} 님에게 {role_name} 역할이 부여 되었습니다.")
        return await ctx.send(f"{role_name} 역할이 존재하지 않습니다.")


def setup(bot):
    bot.add_cog(User(bot))

3) Msg.py

이 카테고리에 속할 커맨드들은 채널의 메시지 관리와 관련된 커맨드, 이벤트 기능이다. 목록은 다음과 같다.
-remove_thombsdown(on_raw_reaction_add에서 이름 변경), num_echo, embed, delete_msg, edit_msg, reaction, wait

Msg.py

from discord.ext import commands
from discord.ext.commands import BadArgument
import discord
import asyncio


class Msg(commands.Cog):

    def __init__(self, bot):
        self.bot = bot
	
    @commands.Cog.listener('on_raw_reaction_add')  # @bot.event()가 Cog.listener로 변경됨
    async def remove_thumbsdown(self, payload):  # 그에 따라 함수 이름도 알맞게 변경
        banned_emoji = "👎"
        author = payload.user_id
        channel = await self.bot.fetch_channel(payload.channel_id)  # bot이 self.bot으로 변경됨
        message = await channel.fetch_message(payload.message_id)
        if payload.emoji.name == banned_emoji and author != self.bot.user.id: # bot이 self.bot으로 
            await message.clear_reaction(banned_emoji)

    @commands.command(name="숫자")
    async def num_echo(self, ctx, user: int):
        await ctx.send(f"입력한 숫자는 {user}입니다.")

    @num_echo.error
    async def num_echo_error(self, ctx, error):
        if isinstance(error, BadArgument):
            await ctx.send("정수를 입력 해주세요")

    @commands.command()
    async def embed(self, ctx):
        embed = discord.Embed(title="Embed title", description="Embed description", color=0x36ccf2)
        embed.set_thumbnail(url="https://cdn.discordapp.com/attachments/721307978455580742/762760779409129513/img.png")
        embed.set_image(url="https://cdn.discordapp.com/attachments/721307978455580742/762760779409129513/img.png")
        embed.add_field(name="field_name1", value="field value1", inline=False)
        embed.add_field(name="field_name2", value="field value2", inline=False)
        embed.add_field(name="field_name3", value="field value3", inline=False)
        embed.add_field(name="field_name4", value="field value4", inline=False)
        embed.set_footer(text="footer_text",
                         icon_url="https://cdn.discordapp.com/attachments/721307978455580742/762760779409129513/img.png")
        embed.set_author(name="author_name", url="https://velog.io/@chaejm55",
                         icon_url="https://cdn.discordapp.com/attachments/721307978455580742/762760779409129513/img.png")

        await ctx.send(embed=embed)

    @commands.command(aliases=['삭제'])
    async def delete_msg(self, ctx):
        msg = await ctx.send("3초 뒤에 삭제 됩니다!")
        await msg.delete(delay=3)

    @commands.command(aliases=['수정'])
    async def edit_msg(self, ctx):
        msg = await ctx.send("곧 수정 됩니다!")
        await msg.edit(content="수정 되었습니다!")

    @commands.command(name="따봉")
    async def reaction(self, ctx):
        await ctx.message.add_reaction('👍')

    @commands.command(name="기다리기")
    async def wait(self, ctx):
        timeout = 5
        send_message = await ctx.send(f'{timeout}초간 기다립니다!')

        def check(m):
            return m.author == ctx.message.author and m.channel == ctx.message.channel

        try:
            msg = await self.bot.wait_for('message', check=check, timeout=timeout)  # bot이 self.bot으로 변경됨
        except asyncio.TimeoutError:
            await ctx.send(f'시간초과 입니다...({timeout}초)')
        else:
            await ctx.send(f'{msg.content}메시지를 {timeout}초 안에 입력하셨습니다!')

            
def setup(bot):
    bot.add_cog(Msg(bot))

4) Util.py

이 카테고리에 속할 커맨드들은 기타 다용도 커맨드들이다. 목록은 다음과 같다.

  • crawl, write_excel, read_excel
from discord.ext import commands
from bs4 import BeautifulSoup
import requests
import openpyxl
import discord


class Util(commands.Cog):

    def __init__(self, bot):
        self.bot = bot

    @commands.command(aliases=['코로나'])  # !코로나 입력 시에도 실행 가능
    async def crawl(self, ctx):
        url = "http://ncov.mohw.go.kr/"
        response = requests.get(url)
        response_code = int(response.status_code)  # 응답 코드 받기

        if response_code == 200:  # 정상 작동(코드 200 반환) 시
            soup = BeautifulSoup(response.content, 'lxml')
        else:  # 오류 발생
            await ctx.send("웹 페이지 오류입니다.")

        # element 찾기

        # soup.find ()로 <div class="liveNum_today_new"> 에서 확진자 수 데이터가 들어 있는 <span class="data"> 리스트 가져오기
        today = soup.find("div", {"class": "liveNum_today_new"}).findAll("span", {"class": "data"})
        today_domestic = int(today[0].text)  # 리스트 첫 번째 요소 (국내발생)
        # today_domestic = int(soup.select_one("body > div > div.mainlive_container > div.container > div > div.liveboard_layout > div.liveNumOuter > div.liveNum_today_new > div > ul > li:nth-child(1) > span.data").text)
        today_overseas = int(today[1].text)  # 리스트 두 번째 요소 (해외유입)
        accumulate_confirmed = soup.find("div", {"class": "liveNum"}).find("span", {"class": "num"}).text[
                               4:]  # 앞에 (누적) 글자 자르기
        embed = discord.Embed(title="국내 코로나 확진자 수 현황", description="http://ncov.mohw.go.kr/ 의 정보를 가져옵니다.",
                              color=0x005666)
        embed.add_field(name="일일 확진자",
                        value=f"총: {today_domestic + today_overseas}, 국내: {today_domestic}, 해외유입: {today_overseas}",
                        inline=False)
        embed.add_field(name="누적 확진자", value=f"{accumulate_confirmed}명", inline=False)
        await ctx.send(embed=embed)

    @commands.command(name="엑셀쓰기")
    async def write_excel(self, ctx, write_str: str):
        filename = "discord_bot.xlsx"
        book = openpyxl.load_workbook(filename)
        ws = book["discord_bot"]
        ws.append([str(ctx.author), write_str])
        book.save(filename)
        await ctx.send("엑셀 입력 완료!")

    @commands.command(name="엑셀읽기")
    async def read_excel(self, ctx):
        filename = "discord_bot.xlsx"
        book = openpyxl.load_workbook(filename)
        ws = book["discord_bot"]
        result = []
        for row in ws.rows:
            if row[0].value == str(ctx.author):
                result.append(row[1].value)
        book.close()
        if result:
            await ctx.send(f'{ctx.author}님이 엑셀 파일에 입력한 내용들입니다.')
            await ctx.send('\n'.join(result))
        else:
            await ctx.send(f'{ctx.author}님이 엑셀 파일에 입력한 내용이 없습니다.')


def setup(bot):
    bot.add_cog(Util(bot))

코드에서 변경된 점들을 다시 살펴보겠다.

  • @bot.commands() 대신 @commands.command()를 사용하였다.
  • 커맨드들이 class 내의 함수가 되었으므로 첫 인자로 self가 추가 되었다.
  • 커맨드 내에서 bot으로 바로 사용하던게 self.bot으로 변경되었다.
  • @bot.event() + async def 이벤트이름()@commands.Cog.listener('이벤트 이름')으로 변경되었다.

3. 사용 예시

사용 시에는 bot.py에 load를 해서 사용해야한다. bot.py에 코드를 추가해보자.

bot.py

import discord
import os
from dotenv import load_dotenv
from discord.ext import commands

load_dotenv()

token = os.getenv('TOKEN')
bot = commands.Bot(command_prefix='!')  # 봇의 접두사 설정

# cogs 폴더의 절대 경로 얻기
# Pycharm에서 바로 상대 경로를 사용하면 오류가 발생하기 때문에 따로 절대경로를 얻어야한다.
cogs_path = 'Cogs'
abs_cogs_path = os.path.join(os.path.abspath(os.path.dirname(__file__)), cogs_path)

# cogs 폴더에 존재하는 cogs(.py파일) 로드
for ext in os.listdir(abs_cogs_path):
    if ext.endswith(".py"):
        bot.load_extension(f"Cogs.{ext.split('.')[0]}")  # .py 부분을 떼고 cog의 이름만 추출


@bot.event
async def on_ready():  # 봇 준비 시 1회 동작하는 부분
    # 봇 이름 하단에 나오는 상태 메시지 설정
    await bot.change_presence(status=discord.Status.online, activity=discord.Game("반갑습니다 :D"))
    print("Bot is ready")


@bot.command()  # 봇 명령어
async def hello(ctx):  # !hello라고 사용자가 입력하면
    await ctx.send("Hello world")  # 봇이 Hello world!라고 대답함


@bot.command(usage='test usage', description='test description', help='test help')
async def test(ctx):
    await ctx.send("test")


bot.run(token)

Cog 로드 후 사용 시 파일 나눴지만 커맨드들은 정상 작동하는 것을 볼 수 있다.

4. 발생할 법할 에러

  • ExtensionNotFound: cogs의 경로가 제대로 설정되지 않을 때 발생한다. 경로를 다시 살펴보자.
  • 기타 에러: 커맨드들을 옮길 때 오타 등의 문제일 확률이 높다. @commands.command(), 커맨드의 첫 인자로 self가 있는지, 각 파일별로 import가 적절히 되었는지 등을 살펴보자.

5. 마무리

Cog를 통해 커맨드들을 카테코리에 따라 나누어 코드를 좀더 깔끔하게 리팩토링 해보았다. 다음에는 Cog의 서브커맨드 등 Cog의 다른 기능을 알아보겠다.

6. Reference

discord.py 공식문서
상대경로와 PyCharm 그리고 명령행

github 전체 코드

profile
여러가지를 시도하는 학생입니다

0개의 댓글