Merge pull request #298 from KoalaBotUK/feature/shedule-activity

Feat: Schedule Activity

Author: JayDwee

CHANGELOG.md

@@ -3,6 +3,9 @@ All notable changes to KoalaBot will be documented in this file. A lot of these
 administrators
 
 ## [Unreleased]
+### Base
+- Add scheduled activities for bot owners
+
 ## [0.5.3] - 12-04-2022
 ### TwitchAlert
 - Fix streams delete themselves when live (again)

koala/cogs/base/cog.py

@@ -10,20 +10,41 @@
 # Built-in/Generic Imports
 
 # Libs
-from discord.ext import commands
+import datetime
+
+import discord
+from discord.ext import commands, tasks
 
 # Own modules
+from discord.ext.commands import BadArgument
+
 import koalabot
 from koala.db import get_all_available_guild_extensions, give_guild_extension, \
     get_enabled_guild_extensions, remove_guild_extension
-from .utils import new_discord_activity, list_ext_embed
+from . import core
+from .utils import list_ext_embed
 from .log import logger
 
 # Constants
 
 # Variables
 
 
+def convert_activity_type(argument):
+    try:
+        return discord.ActivityType[argument]
+    except KeyError:
+        raise BadArgument('Unknown activity type %s' % argument)
+
+
+def convert_iso_datetime(argument):
+    try:
+        return datetime.datetime.fromisoformat(argument)
+    except ValueError:
+        raise BadArgument('Invalid ISO format "%s", instead use the format "2020-01-01 00:00:00"' % argument)
+
+
+
 class BaseCog(commands.Cog, name='KoalaBot'):
     """
         A discord.py cog with general commands useful to managers of the bot and servers
@@ -38,30 +59,97 @@ def __init__(self, bot):
         self._last_member = None
         self.started = False
         self.COGS_DIR = koalabot.COGS_DIR
+        self.current_activity = None
 
     @commands.Cog.listener()
     async def on_ready(self):
         """
         Ran after all cogs have been started and bot is ready
         """
-        await self.bot.change_presence(activity=new_discord_activity("playing", f"{koalabot.COMMAND_PREFIX}help"))
+        core.activity_clear_current()
+        await self.update_activity()
+        self.update_activity.start()
         self.started = True
         logger.info("Bot is ready.")
 
-    @commands.command(name="activity", aliases=["change_activity"])
+    @commands.group(name="activity")
+    @commands.check(koalabot.is_owner)
+    async def activity_group(self, ctx: commands.Context):
+        """
+        Group of commands for activity functionality.
+        :param ctx: Context of the command
+        :return:
+        """
+
+    @activity_group.command(name="set")
     @commands.check(koalabot.is_owner)
-    async def change_activity(self, ctx, new_activity, name):
+    async def activity_set(self, ctx, new_activity: convert_activity_type, name: str, url: str = None):
         """
         Change the activity of the bot
         :param ctx: Context of the command
         :param new_activity: The new activity of the bot
         :param name: The name of the activity
+        :param url: url for streaming
         """
-        if str.lower(new_activity) in ["playing", "watching", "listening", "streaming"]:
-            await self.bot.change_presence(activity=new_discord_activity(new_activity, name))
-            await ctx.send(f"I am now {new_activity} {name}")
-        else:
-            await ctx.send("That is not a valid activity, sorry!\nTry 'playing' or 'watching'")
+        await core.activity_set(new_activity, name, url, bot=self.bot)
+        await ctx.send(f"I am now {new_activity.name} {name}")
+
+    @activity_group.command(name="schedule")
+    @commands.check(koalabot.is_owner)
+    async def activity_schedule(self, ctx, new_activity: convert_activity_type, message: str,
+                                start_time: convert_iso_datetime, end_time: convert_iso_datetime, url: str = None):
+        """
+        Schedule an activity
+        :param ctx: Context of the command
+        :param new_activity: activity type (watching, playing etc.)
+        :param message: message
+        :param start_time: iso format start time
+        :param end_time: iso format end time
+        :param url: url
+        """
+        core.activity_schedule(new_activity, message, url, start_time, end_time)
+        await ctx.send("Activity saved")
+
+    @activity_group.command(name="list")
+    @commands.check(koalabot.is_owner)
+    async def activity_list(self, ctx, show_all: bool = False):
+        """
+        List scheduled activities
+        :param ctx: Context of the command
+        :param show_all: false=future activities, true=all activities
+        """
+        activities = core.activity_list(show_all)
+        result = "Activities:"
+        for activity in activities:
+            result += "\n%s, %s, %s, %s, %s, %s" % (activity.activity_id, activity.activity_type.name,
+                                                    activity.stream_url, activity.message, activity.time_start,
+                                                    activity.time_end)
+        await ctx.send(result)
+
+    @activity_group.command(name="remove")
+    @commands.check(koalabot.is_owner)
+    async def activity_remove(self, ctx, activity_id: int):
+        """
+        Remove an existing activity
+        :param ctx: Context of the command
+        :param activity_id: Activity ID
+        """
+        activity = core.activity_remove(activity_id)
+        result = "Removed:"
+        result += "\n%s, %s, %s, %s, %s, %s" % (activity.activity_id, activity.activity_type.name,
+                                                activity.stream_url, activity.message, activity.time_start,
+                                                activity.time_end)
+        await ctx.send(result)
+
+    @tasks.loop(seconds=1.0)
+    async def update_activity(self):
+        """
+        Loop for updating the activity of the bot according to scheduled activities
+        """
+        try:
+            await core.activity_set_current_scheduled(self.bot)
+        except Exception as err:
+            logger.error("Error in update_activity loop %s" % err, exc_info=err)
 
     @commands.command()
     async def ping(self, ctx):

koala/cogs/base/core.py

@@ -0,0 +1,98 @@
+import datetime
+from typing import List, Optional
+
+import discord
+from discord.ext.commands import Bot
+
+from . import db
+from .log import logger
+from .models import ScheduledActivities
+from koala.db import assign_session
+from .utils import DEFAULT_ACTIVITY, activity_eq
+
+# Constants
+
+# Variables
+current_activity = None
+
+
+def activity_clear_current():
+    global current_activity
+    current_activity = None
+
+
+async def activity_set(activity_type: discord.ActivityType, name: str, url: Optional[str], bot: Bot):
+    """
+    Set presence for this bot to the given presence
+    :param activity_type:
+    :param name:
+    :param url:
+    :param bot: 
+    :return:
+    """
+    new_activity = discord.Activity(type=activity_type, name=name, url=url)
+    await bot.change_presence(activity=new_activity)
+
+
+@assign_session
+def activity_schedule(activity_type: discord.ActivityType, message: str, url: Optional[str], start_time: datetime.datetime,
+                      end_time: datetime.datetime, **kwargs):
+    """
+    Schedule an activity to be used for a timed presence
+    :param activity_type:
+    :param message:
+    :param url:
+    :param start_time:
+    :param end_time:
+    :param kwargs:
+    :return:
+    """
+    db.add_scheduled_activity(activity_type, message, url, start_time, end_time, **kwargs)
+
+
+@assign_session
+def activity_list(show_all: bool, **kwargs) -> List[ScheduledActivities]:
+    """
+    Get a list of all scheduled activity
+    :param show_all:
+    :param kwargs:
+    :return:
+    """
+    return db.get_scheduled_activities(False, not show_all, **kwargs)
+
+
+@assign_session
+def activity_remove(activity_id: int, **kwargs) -> ScheduledActivities:
+    """
+    Remove a scheduled activity
+    :param activity_id:
+    :param kwargs:
+    :return:
+    """
+    return db.remove_scheduled_activities(activity_id, **kwargs)
+
+
+@assign_session
+async def activity_set_current_scheduled(bot: Bot, **kwargs):
+    """
+    Set the current scheduled activity as the bot presence
+    :param bot:
+    :param kwargs:
+    :return:
+    """
+    activities = db.get_scheduled_activities(True, True, **kwargs)
+    if len(activities) > 1:
+        logger.warn("Multiple activities found for this timeslot, %s" % activities)
+
+    if len(activities) != 0:
+        activity = activities[0]
+        new_activity = discord.Activity(
+            type=activity.activity_type, name=activity.message, url=activity.stream_url)
+    else:
+        new_activity = DEFAULT_ACTIVITY
+
+    global current_activity
+    if not activity_eq(new_activity, current_activity):
+        await bot.change_presence(activity=new_activity)
+        logger.info("Auto changing bot presence: %s" % new_activity)
+        current_activity = new_activity

koala/cogs/base/db.py

@@ -0,0 +1,64 @@
+import datetime
+from typing import List, Optional
+
+import discord
+import sqlalchemy.orm
+from sqlalchemy import select
+
+from koala.cogs.base.models import ScheduledActivities
+from koala.db import assign_session
+
+
+@assign_session
+def add_scheduled_activity(activity_type: discord.ActivityType, message: str, url: Optional[str],
+                           time_start: datetime.datetime, time_end: datetime.datetime,
+                           session: sqlalchemy.orm.Session) -> ScheduledActivities:
+    """
+    Add scheduled activity to database
+    :param activity_type:
+    :param message:
+    :param url:
+    :param time_start:
+    :param time_end:
+    :param session:
+    :return:
+    """
+    activity = ScheduledActivities(activity_type=activity_type, message=message, stream_url=url, time_start=time_start,
+                                   time_end=time_end)
+    session.add(activity)
+    session.commit()
+    return activity
+
+
+@assign_session
+def get_scheduled_activities(start_time_restricted: bool, end_time_restricted: bool,
+                             session: sqlalchemy.orm.Session) -> List[ScheduledActivities]:
+    """
+    Get all scheduled activities
+    :param start_time_restricted:
+    :param end_time_restricted:
+    :param session:
+    :return:
+    """
+    current_time = datetime.datetime.now()
+    query = select(ScheduledActivities)
+    if start_time_restricted:
+        query = query.where(ScheduledActivities.time_start < current_time)
+    if end_time_restricted:
+        query = query.where(ScheduledActivities.time_end > current_time)
+
+    return session.execute(query).scalars().all()
+
+
+@assign_session
+def remove_scheduled_activities(activity_id: int, session: sqlalchemy.orm.Session) -> ScheduledActivities:
+    """
+    Delete a specified scheduled activity from a database
+    :param activity_id:
+    :param session:
+    :return:
+    """
+    activity = session.execute(select(ScheduledActivities).filter_by(activity_id=activity_id)).scalar()
+    session.delete(activity)
+    session.commit()
+    return activity

koala/cogs/base/models.py

@@ -0,0 +1,32 @@
+# Futures
+# Built-in/Generic Imports
+# Libs
+from discord import ActivityType
+from sqlalchemy import Column, Integer, String, TIMESTAMP, Enum
+
+# Own modules
+from koala.models import mapper_registry
+from koala.db import setup
+
+# Constants
+
+# Variables
+
+
+@mapper_registry.mapped
+class ScheduledActivities:
+    __tablename__ = 'ScheduledActivities'
+    activity_id = Column(Integer, primary_key=True, autoincrement=True)
+    activity_type = Column(Enum(ActivityType), comment="0: Playing, 1: Streaming, 2: Listening, 3: Watching, 4: Custom,"
+                                                       " 5: Competing")
+    stream_url = Column(String, nullable=True)
+    message = Column(String)
+    time_start = Column(TIMESTAMP)
+    time_end = Column(TIMESTAMP)
+
+    def __repr__(self):
+        return "<ScheduledActivities(%s, %s, %s)>" % \
+               (self.activity_id, self.activity_type, self.message)
+
+
+setup()