botUtilities.py (5637B)
1 # Import dependencies 2 import discord 3 import re 4 5 6 COLORS = { 7 "error": 0xD52C42, 8 "embed": 0x248AD1, 9 } 10 11 12 13 # Error classes 14 class CommandError(Exception): 15 """Raised when a command encounters an anticipated error""" 16 pass 17 class ContributorNotFound(Exception): 18 """Raised when a matching countdown contributor cannot be found""" 19 pass 20 class CountdownNotFound(Exception): 21 """Raised when a matching countdown cannot be found""" 22 pass 23 24 25 26 # The rules for awarding leaderboard points 27 POINT_RULES = { 28 "r1": ("First Number", 0), 29 "r2": ("1000s", 1000), 30 "r3": ("1001s", 500), 31 "r4": ("200s", 200), 32 "r5": ("201s", 100), 33 "r6": ("100s", 100), 34 "r7": ("101s", 50), 35 # "r8": ("Prime Numbers", 15), 36 "r8": ("Odd Numbers", 12), 37 "r9": ("Even Numbers", 10), 38 } 39 40 41 42 async def getUsername(bot, id): 43 """ 44 Get a username from a user ID 45 46 Parameters 47 ---------- 48 bot : commands.Bot 49 The bot 50 id : int 51 The user ID 52 53 Returns 54 ------- 55 str 56 The username 57 """ 58 59 user = await bot.fetch_user(id) 60 return user.name 61 62 63 64 async def getContributor(bot, countdown, text): 65 """ 66 Get the ID of the countdown contributor refered to by a string 67 68 Parameters 69 ---------- 70 bot : commands.Bot 71 The bot 72 countdown : Countdown 73 The countdown 74 text : str 75 The string 76 77 Returns 78 ------- 79 int 80 The ID of the contributor 81 82 Raises 83 ------ 84 ContributorNotFound 85 If a matching contributor cannot be found 86 """ 87 88 if (re.match("^<@\d+>$", text)): 89 return int(text[2:-1]) 90 91 raise ContributorNotFound(text) 92 93 94 95 def isCountdown(cur, id): 96 """ 97 Determine whether a channel is a countdown 98 99 Parameters 100 ---------- 101 cur : psycopg.cursor 102 The database cursor 103 id : int 104 The countdown ID 105 106 Returns 107 ------- 108 bool 109 A boolean indicating whether the channel is a countdown 110 """ 111 112 cur.execute("CALL isCountdown(%s, null);", 113 (id,)) 114 return cur.fetchone()["result"] 115 116 117 118 def getContextCountdown(cur, ctx): 119 """ 120 Get the most relevant countdown to a certain context 121 122 Parameters 123 ---------- 124 cur : psycopg.cursor 125 The database cursor 126 ctx : discord.ext.commands.Context 127 The context 128 129 Returns 130 ------- 131 countdownID 132 The countdown ID 133 """ 134 135 if (isinstance(ctx.channel, discord.channel.TextChannel)): 136 # Channel inside a server 137 cur.execute("CALL getServerContextCountdown(%s, %s, %s, null);", 138 (ctx.channel.guild.id, ctx.channel.id, ctx.prefix)) 139 return cur.fetchone()["countdownid"] 140 141 if (isinstance(ctx.channel, discord.channel.DMChannel)): 142 # DM with a user 143 cur.execute("CALL getUserContextCountdown(%s, null);", 144 (ctx.author.id,)) 145 return cur.fetchone()["countdownid"] 146 147 return None 148 149 150 151 def getPrefix(conn, ctx, default): 152 """ 153 Get the bot prefix for a certain context 154 155 Parameters 156 ---------- 157 conn : psycopg.Connection 158 The database connection 159 ctx : discord.ext.commands.Context 160 The context 161 default : string 162 The default prefix 163 """ 164 165 with conn.cursor() as cur: 166 cur.execute("SELECT * FROM getServerPrefixes(%s, %s);", 167 (ctx.channel.guild.id if ctx.channel.guild else None, ctx.channel.id)) 168 prefixes = cur.fetchall() 169 return [x["prefix"] for x in prefixes] if prefixes else [default] 170 171 172 173 async def addMessage(cur, message): 174 """ 175 Parse a message and add it to a countdown 176 177 Notes 178 ----- 179 If the message is invalid or incorrect, a reaction will be added accordingly 180 181 Parameters 182 ---------- 183 cur : psycopg.cursor 184 The database cursor 185 message : discord.Message 186 The Discord message object 187 188 Returns 189 ------- 190 bool 191 Whether the message was valid and added to the countdown 192 """ 193 194 # Parse message number 195 match = re.search("^[0-9,]+", message.content) 196 if not match: return False 197 number = int(match[0].replace(",", "")) 198 199 # Attempt to add result 200 cur.execute("CALL addMessage(%s,%s,%s,%s,%s,null,null,null);", ( 201 message.id, message.channel.id, message.author.id, number, 202 message.created_at 203 )) 204 result = cur.fetchone() 205 206 # Process result 207 if result["result"] == 'badNumber': 208 await message.add_reaction("❌") 209 if result["result"] == 'badUser': 210 await message.add_reaction("⛔") 211 if result["pin"]: 212 await message.pin() 213 if result["reactions"]: 214 cur.execute("SELECT * FROM getReactions(%s, %s);", 215 (message.channel.id, number)) 216 for reaction in cur.fetchall(): 217 await message.add_reaction(reaction["value"]) 218 219 return result["result"] == 'good' 220 221 222 223 async def loadCountdown(bot, countdown): 224 """ 225 Loads countdown messages from a Discord countdown 226 227 Parameters 228 ---------- 229 bot : commands.Bot 230 The bot to load messages with 231 cur : psycopg.cursor 232 The database cursor 233 countdown : int 234 The ID of the countdown to load messages for 235 """ 236 237 with bot.db_connection.cursor() as cur: 238 # Clear countdown 239 cur.execute("CALL clearCountdown(%s);", (countdown,)) 240 241 # Get Discord messages 242 messages = [message async for message in 243 bot.get_channel(countdown).history(limit=10100)] 244 messages.reverse() 245 246 # Add messages to countdown 247 for message in messages: 248 await addMessage(cur, message) 249 250 # Commit changes 251 bot.db_connection.commit()