CODEWITHSUNDEEP
pythonpython-3.xpython-importpython-asynciodiscord.py
MajorTanya
MajorTanya

Reputation: 141

Interact with background task with commands in submodule [discord.py]

I have a Discord bot written with the non-rewrite version of discord.py that sends a heartbeat-like message (among other things). I don't know if I understood it correctly but from tests, I found out that I need to have the async def heartbeat() function in the main.py file.

Excerpt from main.py (heartbeat works as intended):

[...]
import asyncio
import datetime
from configparser import ConfigParser

startup_time = datetime.datetime.utcnow()
[...]

async def heartbeat():
    await bot.wait_until_ready()

    heartbeat_config = ConfigParser()
    heartbeat_config.read('./config/config.ini')
    hb_freq = int(heartbeat_config.get('Heartbeat', 'hb_freq'))  # frequency of heartbeat message
    hb_channel = heartbeat_config.get('Heartbeat', 'hb_channel')  # target channel of heartbeat message
    hb_channel = bot.get_channel(hb_channel)  # get channel from bot's channels

    await bot.send_message(hb_channel, "Starting up at: `" + str(startup_time) + "`")
    await asyncio.sleep(hb_freq)  # sleep for hb_freq seconds before entering loop
    while not bot.is_closed:
        now = datetime.datetime.utcnow()  # time right now
        tdelta = now - startup_time  # time since startup
        tdelta = tdelta - datetime.timedelta(microseconds=tdelta.microseconds)  # remove microseconds from tdelta
        beat = await bot.send_message(hb_channel, "Still running\nSince: `" + str(startup_time) + "`.\nCurrent uptime: `" + str(tdelta))
        await asyncio.sleep(hb_freq)  # sleep for hb_freq seconds before initialising next beat
        await bot.delete_message(beat)  # delete old beat so it can be replaced

[...]
if __name__ == "__main__":
    global heartbeat_task
    heartbeat_task = bot.loop.create_task(heartbeat())  # creates heartbeat task in the background
    bot.run(token)  # run bot

I have some commands that are supposed to interact with the created heartbeat_task, but they are in a different module, called dev.py (residing in the same directory as main.py).

Excerpt from dev.py:

[...]
from main import heartbeat_task, heartbeat
[...]
@commands.group(pass_context=True)
async def heart(self, ctx):
    if ctx.invoked_subcommand is None:
        return

@heart.command(pass_context=True)
async def stop(self, ctx):
    # should cancel the task from main.py
    heartbeat_task.cancel()
    await self.bot.say('Heartbeat stopped by user {}'.format(ctx.message.author.name))

@heart.command(pass_context=True)
async def start(self, ctx):
    # start the heartbeat if it is not running
    global heartbeat_task
    if heartbeat_task.cancelled():
        heartbeat_task = self.bot.loop.create_task(heartbeat())
        await self.bot.say('Heartbeat started by user {}'.format(ctx.message.author.name))
    else:
        return
[...]

These commands work perfectly fine when they are part of main.py (with the necessary adjustments of course, like removing self, the import, etc), but since I want all developer related commands to go into a module of their own, so I tried moving them.

I get the following error when I try to load the module:

ImportError: cannot import name 'heartbeat_task'.

Removing that import from dev.py leads to a successful loading of the module, but upon using either of the commands, the console throws an error:

NameError: name 'heartbeat_task' is not defined

Which traces back to the line heartbeat_task.cancel() (in the case of heart stop // if heartbeat_task.cancelled(): (in the case of heart start).

Now my question. How can I have the async heartbeat() in main.py but influence the task with the commands in the dev.py module?

And if I can't, what are feasible alternatives that keep the commands in dev.py (the function itself doesn't need to stay in main.py but is preferred to stay there)?

(I have searched for quite some time and couldn't find a problem like mine or a solution that happened to work for me too)

Upvotes: 2

Views: 2155

Answers (1)

Patrick Haugh
Patrick Haugh

Reputation: 60974

The easiest way to have a background task in a cog is to add an on_ready coroutine to the cog that will kick off the background task, instead of starting it manually:

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

    async def heartbeat(self):
        ...

    async def on_ready(self):
        self.heartbeat_task = self.bot.loop.create_task(heartbeat())

    @commands.command(pass_context=True)
    async def stop(self, ctx):
        self.heartbeat_task.cancel()
        await self.bot.say('Heartbeat stopped by user {}'.format(ctx.message.author.name))


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

Note that you don't need to decorate on_ready with anything in the cog, the add_cog machinery will pick it up based on its name.

Upvotes: 2

Related Questions