Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Migrated censor collection to ODM #6

Open
wants to merge 1 commit into
base: feat/beanie-odm/master
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions bot.py
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@
from discord.ext import commands
from rich.logging import RichHandler

import src.mongo.models
from env import env
from src.discord.globals import (
CHANNEL_BOTSPAM,
Expand Down Expand Up @@ -177,6 +178,7 @@ async def setup_hook(self) -> None:
await init_beanie(
database=self.mongo_database.client["data"],
document_models=[
src.mongo.models.Censor,
# TODO
],
)
Expand Down
10 changes: 5 additions & 5 deletions src/discord/censor.py
Original file line number Diff line number Diff line change
Expand Up @@ -92,10 +92,10 @@ async def on_message(self, message: discord.Message) -> None:

def word_present(self, content: str) -> bool:
with contextlib.suppress(asyncio.CancelledError):
for word in src.discord.globals.CENSOR["words"]:
for word in src.discord.globals.CENSOR.words:
if re.findall(rf"\b({word})\b", content, re.I):
return True
for emoji in src.discord.globals.CENSOR["emojis"]:
for emoji in src.discord.globals.CENSOR.emojis:
if len(re.findall(emoji, content)):
return True
return False
Expand Down Expand Up @@ -141,14 +141,14 @@ async def __censor(self, message: discord.Message):
author = message.author.nick or message.author.name

# Actually replace content found on the censored words/emojis list
for word in src.discord.globals.CENSOR["words"]:
for word in src.discord.globals.CENSOR.words:
content = re.sub(
rf"\b({word})\b",
"<censored>",
content,
flags=re.IGNORECASE,
)
for emoji in src.discord.globals.CENSOR["emojis"]:
for emoji in src.discord.globals.CENSOR.emojis:
content = re.sub(emoji, "<censored>", content, flags=re.I)

reply = (
Expand Down Expand Up @@ -301,7 +301,7 @@ async def on_raw_reaction_add(self, payload: discord.RawReactionActionEvent):
Args:
payload (discord.RawReactionActionEvent): The payload to use for the message.
"""
if str(payload.emoji) in src.discord.globals.CENSOR["emojis"]:
if str(payload.emoji) in src.discord.globals.CENSOR.emojis:
channel = self.bot.get_channel(payload.channel_id)
assert isinstance(channel, discord.TextChannel)

Expand Down
5 changes: 4 additions & 1 deletion src/discord/globals.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@
Holds global variables shared between cogs and variables that are initialized when
the bot is first setup.
"""
from src.mongo.models import Censor

##############
# CONSTANTS
Expand Down Expand Up @@ -120,7 +121,9 @@
# VARIABLES
##############
fish_now = 0
CENSOR = {}
CENSOR: Censor = {}
# FIXME: CENSOR for now has to be a dummy value since Beanie does
# not get initialized at the global scope before importing globals.py
EVENT_INFO = []
PING_INFO = []
INVITATIONAL_INFO = []
Expand Down
42 changes: 10 additions & 32 deletions src/discord/staff/censor.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@
from typing import TYPE_CHECKING, Literal

import discord
from beanie.odm.operators.update.array import Pull, Push
from discord import app_commands
from discord.ext import commands

Expand All @@ -14,6 +15,7 @@
ROLE_STAFF,
ROLE_VIP,
)
from src.mongo.models import Censor

if TYPE_CHECKING:
from bot import PiBot
Expand Down Expand Up @@ -54,36 +56,24 @@ async def censor_add(
)

if censor_type == "word":
if phrase in src.discord.globals.CENSOR["words"]:
if phrase in src.discord.globals.CENSOR.words:
await interaction.edit_original_response(
content=f"`{phrase}` is already in the censored words list. Operation cancelled.",
)
else:
src.discord.globals.CENSOR["words"].append(phrase)
await self.bot.mongo_database.update(
"data",
"censor",
src.discord.globals.CENSOR["_id"],
{"$push": {"words": phrase}},
)
await src.discord.globals.CENSOR.update(Push({Censor.words: phrase}))
first_letter = phrase[0]
last_letter = phrase[-1]
await interaction.edit_original_response(
content=f"Added `{first_letter}...{last_letter}` to the censor list.",
)
elif censor_type == "emoji":
if phrase in src.discord.globals.CENSOR["emojis"]:
if phrase in src.discord.globals.CENSOR.emojis:
await interaction.edit_original_response(
content="Emoji is already in the censored emoijs list. Operation cancelled.",
)
else:
src.discord.globals.CENSOR["emojis"].append(phrase)
await self.bot.mongo_database.update(
"data",
"censor",
src.discord.globals.CENSOR["_id"],
{"$push": {"emojis": phrase}},
)
await src.discord.globals.CENSOR.update(Push({Censor.emojis: phrase}))
await interaction.edit_original_response(
content="Added emoji to the censor list.",
)
Expand Down Expand Up @@ -112,34 +102,22 @@ async def censor_remove(
)

if censor_type == "word":
if phrase not in src.discord.globals.CENSOR["words"]:
if phrase not in src.discord.globals.CENSOR.words:
await interaction.response.send_message(
content=f"`{phrase}` is not in the list of censored words.",
)
else:
src.discord.globals.CENSOR["words"].remove(phrase)
await self.bot.mongo_database.update(
"data",
"censor",
src.discord.globals.CENSOR["_id"],
{"$pull": {"words": phrase}},
)
await src.discord.globals.CENSOR.update(Pull({Censor.words: phrase}))
await interaction.edit_original_response(
content=f"Removed `{phrase}` from the censor list.",
)
elif censor_type == "emoji":
if phrase not in src.discord.globals.CENSOR["emojis"]:
if phrase not in src.discord.globals.CENSOR.emojis:
await interaction.response.send_message(
content=f"{phrase} is not in the list of censored emojis.",
)
else:
src.discord.globals.CENSOR["emojis"].remove(phrase)
await self.bot.mongo_database.update(
"data",
"censor",
src.discord.globals.CENSOR["_id"],
{"$pull": {"emojis": phrase}},
)
await src.discord.globals.CENSOR.update(Pull({Censor.emojis: phrase}))
await interaction.edit_original_response(
content=f"Removed {phrase} from the emojis list.",
)
Expand Down
7 changes: 6 additions & 1 deletion src/discord/tasks.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@
from env import env
from src.discord.invitationals import update_invitational_list
from src.discord.views import UnselfmuteView
from src.mongo.models import Censor

if TYPE_CHECKING:
from bot import PiBot
Expand Down Expand Up @@ -81,7 +82,11 @@ async def pull_prev_info(self):
self.bot.settings = await self.bot.mongo_database.get_settings()
assert isinstance(self.bot.settings, dict)

src.discord.globals.CENSOR = await self.bot.mongo_database.get_censor()
src.discord.globals.CENSOR = await Censor.find_one({})

if not src.discord.globals.CENSOR:
src.discord.globals.CENSOR = Censor(words=[], emojis=[])
await src.discord.globals.CENSOR.save()
logger.info("Fetched previous variables.")

async def add_to_cron(self, item_dict: dict) -> None:
Expand Down
11 changes: 11 additions & 0 deletions src/mongo/models.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,14 @@
"""
Contains all database models
"""

from beanie import Document


class Censor(Document):
words: list[str]
emojis: list[str]

class Settings:
name = "censor"
use_cache = True
6 changes: 0 additions & 6 deletions src/mongo/mongo.py
Original file line number Diff line number Diff line change
Expand Up @@ -95,12 +95,6 @@ async def get_cron(self):
"""
return await self.get_entire_collection("data", "cron")

async def get_censor(self):
"""
Gets the document containing censor information from the censor collection.
"""
return await self.get_entire_collection("data", "censor", return_one=True)

async def get_pings(self):
"""
Gets all documents in the pings collection.
Expand Down
Loading