Update __init__.py, cached_messaging.py, and 5 more files...

Author: WardPearce

SourceQueryBot/__init__.py

@@ -1,20 +1,34 @@
 import discord
+from discord import Embed, Color, Status, Game
 from discord.ext import tasks
 
+import asyncio
+import sys
+import logging
+
+from aioquery.exceptions import InvalidServer, DidNotReceive, UnableToConnect
+
 from .translations import TranslationBase, English
 from .messages import Messages
+from .cached_messaging import CachedMessaging
 
 
 __version__ = "0.0.1"
 
 
+logging.basicConfig(level=logging.INFO)
+
+
 class SourceQueryBot(discord.Client):
-    refresh_rate = 0.0
+    _refresh_rate = 0.0
+    _presence_completed = True
+    _total_servers = 0
 
     def __init__(self, catagories: list, refresh_rate: float = 30.0,
                  lanague: TranslationBase = English,
                  messages: Messages = Messages(),
-                 *args, **kwargs):
+                 smart_presence: bool = False,
+                 *args, **kwargs) -> None:
         """
         Pass any discord bot client args / paramters.
 
@@ -26,24 +40,157 @@ def __init__(self, catagories: list, refresh_rate: float = 30.0,
             Defaults to english.
         messages: Messages
             Handles message caching.
+        smart_presence: bool
+            Defaults to false, if enabled presence will rotate between maps.
         """
 
+        super().__init__(*args, **kwargs)
+
         self.lanague = lanague
-        self.refresh_rate = refresh_rate
+        self._refresh_rate = refresh_rate
+        self.catagories = catagories
+        self.messages = messages
 
-        super().__init__(*args, **kwargs)
+        if not sys.version_info[1] >= 7 and smart_presence:
+            self.smart_presence = False
+            logging.log(
+                logging.INFO,
+                lanague.invalid_py_version.format(
+                    sys.version
+                )
+            )
+        else:
+            self.smart_presence = smart_presence
 
-    async def on_ready(self):
+        self.query_task.start()
+
+    async def on_ready(self) -> None:
         print(self.lanague.on_ready.format(self.user))
 
-    async def close(self):
+    async def close(self) -> None:
         print(self.lanague.shutdown)
+
+        await super().close()
         self.query_task.cancel()
 
-    @tasks.loop(seconds=refresh_rate)
-    async def query_task(self):
-        pass
+    async def _smart_presence(self, players, max_players,
+                              map_name, loop_index) -> None:
+
+        if loop_index == self._total_servers:
+            self._presence_completed = False
+
+        await asyncio.sleep(loop_index * 10)
+
+        await self.change_presence(
+            status=Status.online,
+            activity=Game(
+                self.lanague.smart_presence.format(
+                    players,
+                    max_players,
+                    map_name
+                )
+            )
+        )
+
+        if loop_index == self._total_servers:
+            self._presence_completed = True
+
+    @tasks.loop(seconds=30.0)
+    async def query_task(self) -> None:
+        total_players = 0
+        total_max_players = 0
+
+        loop_index = -1
+
+        for catgory in self.catagories:
+            embed = Embed(
+                title=catgory.name,
+                color=Color(catgory.color)
+            )
+
+            for server in catgory.servers:
+                try:
+                    info = await server.interact.info()
+
+                    embed.add_field(
+                        name=server.alt_name if server.alt_name else
+                        info.hostname[:catgory.server_name_limit] + "...",
+                        value=self.lanague.server.format(
+                            info.map,
+                            info.players,
+                            info.max_players,
+                            "{}:{}".format(server.ip, server.port)
+                        ),
+                        inline=catgory.inline
+                    )
+
+                    total_players += info.players
+                    total_max_players += info.max_players
+
+                    if self.smart_presence and self._presence_completed:
+                        loop_index += 1
+
+                        asyncio.create_task(
+                            self._smart_presence(
+                                info.players,
+                                info.max_players,
+                                info.map,
+                                loop_index
+                            )
+                        )
+
+                except (InvalidServer, DidNotReceive, UnableToConnect):
+                    embed.add_field(
+                        name=server.alt_name if server.alt_name else
+                        self.lanague.offline_title,
+                        value=self.lanague.offline_msg
+                    )
+
+                await asyncio.sleep(0.0001)
+
+            sent_message = await CachedMessaging(
+                catgory, self.get_channel(catgory.channel), embed=embed
+            ).send_edit()
+
+            await self.messages.save("{}:{}".format(
+                catgory.channel, sent_message.id
+            ))
+
+            await asyncio.sleep(0.0001)
+
+        if not self.smart_presence:
+            await self.change_presence(
+                status=Status.online,
+                activity=Game(
+                    self.lanague.normal_presence.format(
+                        total_players,
+                        total_max_players
+                    )
+                )
+            )
 
     @query_task.before_loop
-    async def before_query_task(self):
-        pass
+    async def before_query_task(self) -> None:
+        """
+        Handles deleting old messages & caching config details.
+        """
+
+        await self.wait_until_ready()
+
+        async for message in self.messages.get():
+            try:
+                channel_id, message_id = message.split(":")
+            except ValueError:
+                pass
+            else:
+                channel = self.get_channel(int(channel_id))
+                if channel:
+                    message_obj = await channel.fetch_message(int(message_id))
+                    if message_obj:
+                        await message_obj.delete()
+
+        for catgory in self.catagories:
+            for server in catgory.servers:
+                self._total_servers += 1
+
+            await asyncio.sleep(0.00001)

SourceQueryBot/cached_messaging.py

@@ -0,0 +1,29 @@
+from .settings import Category
+from .exceptions import InvalidChannel
+
+from discord import channel
+
+
+CACHED_CATAGORIES = {}
+
+
+class CachedMessaging:
+    def __init__(self, catgory: Category, channel: channel,
+                 *args, **kwargs) -> None:
+        self.catgory = catgory
+        self.args = args
+        self.kwargs = kwargs
+        self.channel = channel
+
+    async def send_edit(self) -> channel:
+        if self.catgory in CACHED_CATAGORIES:
+            message = await CACHED_CATAGORIES[self.catgory].edit(**self.kwargs)
+        else:
+            if self.channel:
+                message = await self.channel.send(**self.kwargs)
+
+                CACHED_CATAGORIES[self.catgory] = message
+            else:
+                raise InvalidChannel()
+
+        return CACHED_CATAGORIES[self.catgory]

SourceQueryBot/exceptions.py

@@ -4,3 +4,11 @@ class NoSuchConfigFile(Exception):
     """
 
     pass
+
+
+class InvalidChannel(Exception):
+    """
+    Raised when a channel ID is invalid or can't be accessed.
+    """
+
+    pass

SourceQueryBot/messages.py

@@ -1,16 +1,16 @@
 import aiofiles
 import typing
 
-from os import path
+from os import path, getcwd
 
 from .exceptions import NoSuchConfigFile
 
 
 class Messages:
+    _saved_messages = []
+
     def __init__(self,
-                 location: str = path.join(
-                     path.dirname(path.realpath(__file__)),
-                     "messages.txt")) -> None:
+                 location: str = path.join(getcwd(), "messages.txt")) -> None:
 
         self.location = location
 
@@ -34,12 +34,11 @@ async def get(self) -> list:
         """
 
         if path.exists(self.location):
-            async with aiofiles.open(self.location) as file:
-                data = await file.read()
+            async with aiofiles.open(self.location, "r+") as file:
+                async for line in file:
+                    yield line
 
                 await file.truncate(0)
-
-                return data.split("\n")
         else:
             raise NoSuchConfigFile()
 
@@ -53,5 +52,8 @@ async def save(self, line: typing.Any) -> None:
             line to save.
         """
 
-        async with aiofiles.open(self.location, mode="a") as file:
-            await file.write(str(line))
+        if line not in self._saved_messages:
+            self._saved_messages.append(line)
+
+            async with aiofiles.open(self.location, mode="a") as file:
+                await file.write("{}\n".format(line))

SourceQueryBot/settings.py

@@ -1,3 +1,6 @@
+import aioquery
+
+
 class Server:
     def __init__(self, ip: str, port: int = 27015,
                  alt_name: str = None) -> None:
@@ -18,10 +21,13 @@ def __init__(self, ip: str, port: int = 27015,
         self.port = port
         self.alt_name = alt_name
 
+        self.interact = aioquery.client(ip, port)
+
 
 class Category:
-    def __init__(self, name: str, channel: int,
-                 servers: list, server_name_limit: int = 20) -> None:
+    def __init__(self, name: str, channel: int, color: hex,
+                 servers: list, server_name_limit: int = 20,
+                 inline: bool = False) -> None:
         """
         Handles server commands.
 
@@ -33,11 +39,17 @@ def __init__(self, name: str, channel: int,
             Limits how many characters the server name can be.
         channel: int
             Discord channel ID for this category.
+        color: hex
+            Hex color code.
         servers: list
             List of Server objects.
+        inline: bool
+            Defaults to false, if servers should be inline.
         """
 
         self.name = name
         self.server_name_limit = server_name_limit
         self.channel = channel
         self.servers = servers
+        self.color = color
+        self.inline = inline

SourceQueryBot/translations.py

@@ -1,8 +1,22 @@
 class TranslationBase:
     shutdown: str
     on_ready: str
+    server: str
+    offline_title: str
+    offline_msg: str
+    normal_presence: str
+    smart_presence: str
+    invalid_py_version: str
 
 
 class English(TranslationBase):
     shutdown = "The bot is shutting down, please DO NOT press ctrl + C"
     on_ready = "Logged on as {}"
+    server = """**Map:** {}\n**Players:** {}/{}
+**Connect:**\nsteam://connect/{}"""
+    offline_title = "Unknown"
+    offline_msg = "**This server is offline.**"
+    normal_presence = "{}/{} on our servers."
+    smart_presence = "{}/{} on {}"
+    invalid_py_version = "Python 3.7 or above is required for" + \
+        " smart presence, you're running {}"

run.py

@@ -7,11 +7,23 @@
         Category(
             "Retake",
             705852327550124062,
+            0xffffff,
             [
-                Server("216.52.148.47")
+                Server("216.52.148.47"),
+                Server("54.37.111.216"),
+                Server("92.119.148.18"),
+            ]
+        ),
+        Category(
+            "Bhop",
+            705852327550124062,
+            0x444444,
+            [
+                Server("92.119.148.31")
             ]
         )
-    ]
+    ],
+    smart_presence=True
 )
 
 if __name__ == "__main__":