forked from Discord-TTS/Bot
-
Notifications
You must be signed in to change notification settings - Fork 0
/
main.py
143 lines (117 loc) · 4.73 KB
/
main.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
import asyncio
import traceback
from concurrent.futures import ProcessPoolExecutor
from configparser import ConfigParser
from os import listdir
from time import monotonic
from typing import Union
import aiohttp
import asyncgTTS
import asyncpg
import discord
from discord.ext import commands
from utils.basic import remove_chars
from utils.decos import wrap_with
print("Starting TTS Bot!")
start_time = monotonic()
# Read config file
config = ConfigParser()
config.read("config.ini")
# Setup activity and intents for logging in
activity = discord.Activity(name=config["Activity"]["name"], type=getattr(discord.ActivityType, config["Activity"]["type"]))
intents = discord.Intents(voice_states=True, messages=True, guilds=True, members=True)
status = getattr(discord.Status, config["Activity"]["status"])
# Custom prefix support
async def prefix(bot: commands.AutoShardedBot, message: discord.Message) -> str:
"Gets the prefix for a guild based on the passed message object"
return await bot.settings.get(message.guild, "prefix") if message.guild else "-"
class TTSBot(commands.AutoShardedBot):
def __init__(self, config, session, executor, *args, **kwargs):
self.channels = {}
self.config = config
self.session = session
self.executor = executor
self.sent_fallback = False
self.trusted = remove_chars(config["Main"]["trusted_ids"], "[", "]", "'").split(", ")
super().__init__(*args, **kwargs)
@property
def support_server(self):
return self.get_guild(int(self.config["Main"]["main_server"]))
async def check_gtts(self) -> Union[bool, Exception]:
try:
await self.gtts.get(text="RL Test", lang="en")
return True
except asyncgTTS.RatelimitException:
return False
except Exception as e:
return e
def load_extensions(self, folder):
filered_exts = filter(lambda e: e.endswith(".py"), listdir(folder))
for ext in filered_exts:
self.load_extension(f"{folder}.{ext[:-3]}")
async def start(self, *args, token, **kwargs):
# Get everything ready in async env
db_info = self.config["PostgreSQL Info"]
self.gtts, self.pool = await asyncio.gather(
asyncgTTS.setup(
premium=False,
session=self.session
),
asyncpg.create_pool(
host=db_info["ip"],
user=db_info["name"],
database=db_info["db"],
password=db_info["pass"]
)
)
# Fill up bot.channels, as a load of webhooks
for channel_name, webhook_url in self.config["Channels"].items():
self.channels[channel_name] = discord.Webhook.from_url(
url=webhook_url,
adapter=discord.AsyncWebhookAdapter(session=self.session)
)
# Load all of /cogs and /extensions
self.load_extensions("cogs")
self.load_extensions("extensions")
# Send starting message and actually start the bot
await self.channels["logs"].send("Starting TTS Bot!")
await super().start(token, *args, **kwargs)
def get_error_string(e: BaseException) -> str:
return f"{type(e).__name__}: {e}"
@wrap_with(ProcessPoolExecutor, aenter=False)
@wrap_with(aiohttp.ClientSession, aenter=True)
async def main(session, executor):
bot = TTSBot(
config=config,
status=status,
intents=intents,
session=session,
executor=executor,
help_command=None, # Replaced by FancyHelpCommand by FancyHelpCommandCog
activity=activity,
command_prefix=prefix,
case_insensitive=True,
chunk_guilds_at_startup=False,
allowed_mentions=discord.AllowedMentions(everyone=False, roles=False)
)
try:
print("\nLogging into Discord...")
ready_task = asyncio.create_task(bot.wait_until_ready())
bot_task = asyncio.create_task(bot.start(token=config["Main"]["Token"]))
done, pending = await asyncio.wait((bot_task, ready_task), return_when=asyncio.FIRST_COMPLETED)
if bot_task in done:
raise RuntimeError(f"Bot Shutdown before ready: {get_error_string(bot_task.exception())}")
print(f"Logged in as {bot.user} and ready!")
await bot.channels["logs"].send(f"Started and ready! Took `{monotonic() - start_time:.2f} seconds`")
await bot_task
except Exception as e:
print(get_error_string(e))
finally:
if not bot.user:
return
await bot.channels["logs"].send(f"{bot.user.mention} is shutting down.")
await bot.close()
try:
asyncio.run(main())
except (KeyboardInterrupt, RuntimeError) as e:
print(f"Shutdown forcefully: {get_error_string(e)}")