Modernized codebase with NamedTuples, StrEnum, override decorators, slots, and other idiomatic improvements.

crabstero/bot.py

@@ -20,6 +20,7 @@ cog loading, ingestion worker pool, and graceful shutdown.
 
 import asyncio
 import logging
+from typing import override
 
 import discord
 from discord import app_commands
@@ -41,6 +42,8 @@ from crabstero.tasks.ingestion import ingest_channel
 
 logger = logging.getLogger(__name__)
 
+type IngestableChannel = discord.TextChannel | discord.VoiceChannel
+
 
 class Crabstero(commands.Bot):
     """Central bot subclass that owns all lifecycle state.
@@ -80,9 +83,7 @@ class Crabstero(commands.Bot):
         self._database_path = database_path
         self._ingestion_worker_count = ingestion_workers
         self._ingest_only = ingest_only
-        self._ingestion_queue: asyncio.Queue[
-            discord.TextChannel | discord.VoiceChannel
-        ] = asyncio.Queue()
+        self._ingestion_queue: asyncio.Queue[IngestableChannel] = asyncio.Queue()
         self._ingestion_workers: list[asyncio.Task[None]] = []
         self._db: Database | None = None
         self.ingest_cache = IngestCache()
@@ -92,7 +93,7 @@ class Crabstero(commands.Bot):
 
         DISCORD_LATENCY.set_function(lambda: self.latency)
         GUILD_COUNT.set_function(lambda: len(self.guilds))
-        INGESTION_BACKLOG.set_function(lambda: self._ingestion_queue.qsize())
+        INGESTION_BACKLOG.set_function(self._ingestion_queue.qsize)
 
         repo_url = "https://git.logal.dev/LogalDeveloper/Crabstero"
         self.http.user_agent = f"DiscordBot ({repo_url}, {crabstero_version})"
@@ -112,6 +113,7 @@ class Crabstero(commands.Bot):
         """Whether the bot is running in ingest-only mode."""
         return self._ingest_only
 
+    @override
     async def setup_hook(self) -> None:
         """Open the database, start ingestion workers, and load all cogs."""
         self._db = await Database.connect(self._database_path)
@@ -149,6 +151,7 @@ class Crabstero(commands.Bot):
                 logger.info("Slash command tree has changed, syncing with Discord.")
                 await self.tree.sync()
 
+    @override
     def dispatch(self, event: str, /, *args: object, **kwargs: object) -> None:
         """Dispatch an event, incrementing the events counter.
 
@@ -161,10 +164,12 @@ class Crabstero(commands.Bot):
         """Log that the bot has started successfully."""
         logger.info("Crabstero started!")
 
+    @override
     async def start(self, token: str = "", *, reconnect: bool = True) -> None:
         """Start the bot using the token provided at initialization."""
         await super().start(self._token, reconnect=reconnect)
 
+    @override
     async def close(self) -> None:
         """Cancel ingestion workers, close the database, and then the bot connection."""
         if self.is_closed():
@@ -181,9 +186,7 @@ class Crabstero(commands.Bot):
             await self._db.close()
         await super().close()
 
-    def queue_channel_for_ingestion(
-        self, channel: discord.TextChannel | discord.VoiceChannel
-    ) -> None:
+    def queue_channel_for_ingestion(self, channel: IngestableChannel) -> None:
         """Enqueue a single channel for background message history ingestion.
 
         :param channel: The channel to enqueue.

crabstero/cache.py

@@ -47,6 +47,8 @@ class IngestCache:
     :param cleanup_interval_seconds: How often the background task runs.
     """
 
+    __slots__ = ("_cleanup_interval", "_entries", "_task", "_ttl")
+
     def __init__(
         self,
         ttl_seconds: float = 120,

crabstero/database.py

@@ -14,10 +14,36 @@
 
 """Async SQLite database access for Crabstero's persistent storage."""
 
-from typing import Self
+from typing import NamedTuple, Self
 
 import aiosqlite
 
+
+class StartWord(NamedTuple):
+    """A Markov chain starting word row."""
+
+    channel_id: int
+    user_id: int
+    word: str
+
+
+class Transition(NamedTuple):
+    """A Markov chain word transition row."""
+
+    channel_id: int
+    user_id: int
+    word: str
+    next_word: str
+
+
+class ChannelImage(NamedTuple):
+    """An image URL associated with a channel."""
+
+    channel_id: int
+    user_id: int
+    url: str
+
+
 # SQL statements for creating the database schema.
 _SCHEMA = """
 -- Markov chain starting words.
@@ -107,15 +133,15 @@ class Database:
 
     async def add_markov_data(
         self,
-        start_words: list[tuple[int, int, str]],
-        transitions: list[tuple[int, int, str, str]],
+        start_words: list[StartWord],
+        transitions: list[Transition],
     ) -> None:
         """Insert Markov start words and transitions, then commit.
 
         Both inserts happen in a single transaction.
 
-        :param start_words: A list of (channel_id, user_id, word) tuples.
-        :param transitions: A list of (channel_id, user_id, word, next_word) tuples.
+        :param start_words: Starting word entries to insert.
+        :param transitions: Transition entries to insert.
         """
         if start_words:
             await self._connection.executemany(
@@ -136,16 +162,16 @@ class Database:
 
     async def remove_markov_data(
         self,
-        start_words: list[tuple[int, int, str]],
-        transitions: list[tuple[int, int, str, str]],
+        start_words: list[StartWord],
+        transitions: list[Transition],
     ) -> None:
         """Remove one matching row per entry from the Markov tables.
 
         Each entry removes at most one duplicate row, preserving remaining
         frequency weight.
 
-        :param start_words: A list of (channel_id, user_id, word) tuples.
-        :param transitions: A list of (channel_id, user_id, word, next_word) tuples.
+        :param start_words: Starting word entries to remove.
+        :param transitions: Transition entries to remove.
         """
         for channel_id, user_id, word in start_words:
             await self._connection.execute(
@@ -213,6 +239,8 @@ class Database:
         Weighted by occurrence frequency. A completing word is one whose last
         character is '.', '!', '?', or '§'.
 
+        The set of terminators must match ``markov._TERMINATORS``.
+
         :param channel_id: The Discord channel ID.
         :param word: The current word to find a completing transition for.
         :return: A random completing next word, or None if none exist.
@@ -228,10 +256,10 @@ class Database:
             row = await cursor.fetchone()
             return row[0] if row else None
 
-    async def add_images(self, images: list[tuple[int, int, str]]) -> None:
+    async def add_images(self, images: list[ChannelImage]) -> None:
         """Store image URLs for a given channel.
 
-        :param images: A list of (channel_id, user_id, url) tuples.
+        :param images: Image entries to insert.
         """
         if images:
             await self._connection.executemany(
@@ -242,13 +270,13 @@ class Database:
             )
             await self._connection.commit()
 
-    async def remove_images(self, images: list[tuple[int, int, str]]) -> None:
+    async def remove_images(self, images: list[ChannelImage]) -> None:
         """Remove one matching row per entry from the images table.
 
         Each entry removes at most one duplicate row, preserving remaining
         frequency weight.
 
-        :param images: A list of (channel_id, user_id, url) tuples.
+        :param images: Image entries to remove.
         """
         for channel_id, user_id, url in images:
             await self._connection.execute(

crabstero/flags.py

@@ -27,7 +27,7 @@ if TYPE_CHECKING:
     from crabstero.database import Database
 
 
-class Flag(enum.Enum):
+class Flag(enum.StrEnum):
     """Enum of all supported flag names."""
 
     NO_REPLY = "noReply"
@@ -35,7 +35,7 @@ class Flag(enum.Enum):
     ALLOW_PINGS = "allowPings"
 
 
-class EntityType(enum.Enum):
+class EntityType(enum.StrEnum):
     """Enum of entity types that can have flags."""
 
     CHANNEL = "channel"
@@ -65,7 +65,7 @@ async def set_flag(
     :param entity_type: The type of the entity.
     :param flag: The flag to set.
     """
-    await db.set_flag(entity_type.value, _entity_id(entity), flag.value)
+    await db.set_flag(entity_type, _entity_id(entity), flag)
 
 
 async def clear_flag(
@@ -81,7 +81,7 @@ async def clear_flag(
     :param entity_type: The type of the entity.
     :param flag: The flag to clear.
     """
-    await db.clear_flag(entity_type.value, _entity_id(entity), flag.value)
+    await db.clear_flag(entity_type, _entity_id(entity), flag)
 
 
 async def is_flag_set(
@@ -98,4 +98,4 @@ async def is_flag_set(
     :param flag: The flag to check for.
     :return: True if the flag is set, False otherwise.
     """
-    return await db.is_flag_set(entity_type.value, _entity_id(entity), flag.value)
+    return await db.is_flag_set(entity_type, _entity_id(entity), flag)

crabstero/markov.py

@@ -24,6 +24,7 @@ import re
 from typing import TYPE_CHECKING
 
 from crabstero import metrics
+from crabstero.database import StartWord, Transition
 
 if TYPE_CHECKING:
     from crabstero.database import Database
@@ -116,8 +117,8 @@ async def _ingest_sentence(
     :param sentence: The sentence to ingest.
     """
     raw_starts, raw_transitions = _tokenize_sentence(sentence)
-    start_words = [(channel_id, user_id, w) for w in raw_starts]
-    transitions = [(channel_id, user_id, w, nw) for w, nw in raw_transitions]
+    start_words = [StartWord(channel_id, user_id, w) for w in raw_starts]
+    transitions = [Transition(channel_id, user_id, w, nw) for w, nw in raw_transitions]
     await db.add_markov_data(start_words, transitions)
 
 
@@ -147,8 +148,8 @@ async def _uningest_sentence(
     :param sentence: The sentence to uningest.
     """
     raw_starts, raw_transitions = _tokenize_sentence(sentence)
-    start_words = [(channel_id, user_id, w) for w in raw_starts]
-    transitions = [(channel_id, user_id, w, nw) for w, nw in raw_transitions]
+    start_words = [StartWord(channel_id, user_id, w) for w in raw_starts]
+    transitions = [Transition(channel_id, user_id, w, nw) for w, nw in raw_transitions]
     await db.remove_markov_data(start_words, transitions)
 
 
@@ -201,8 +202,7 @@ async def generate(
             result = result[:hard_limit]
 
         # Strip the internal sentence-end marker so it never appears in output.
-        if result.endswith(DEFAULT_SENTENCE_END):
-            result = result[:-1]
+        result = result.removesuffix(DEFAULT_SENTENCE_END)
 
         metrics.GENERATED_MESSAGE_LENGTH.observe(len(result))
         return result

crabstero/messages.py

@@ -27,14 +27,15 @@ import discord
 
 from crabstero import flags, markov, metrics
 from crabstero.cache import CachedMessage, IngestCache
+from crabstero.database import ChannelImage
 from crabstero.flags import EntityType, Flag
 
 if TYPE_CHECKING:
     from crabstero.database import Database
 
 # Copied from discordjs/discord-api-types:
-# https://github.com/discordjs/discord-api-types/blob/7fe434114e91c80ed79f0204ae6c73047672d55d/globals.ts#L30
-MENTION_PATTERN = re.compile(r"<@!?(?P<id>\d{17,20})>")
+# https://github.com/discordjs/discord-api-types/blob/662cb0cb0ac9c6f9ad93e180849476714bfceb0c/globals.ts#L39
+_MENTION_PATTERN = re.compile(r"<@!?(?P<id>\d{17,20})>")
 
 _EMBED_CHANCE_THRESHOLD = 95  # Out of 100; sends an embed ~5% of the time.
 
@@ -89,7 +90,7 @@ async def reply_to_message(db: Database, message: discord.Message) -> None:
     # Suppress all mentions by default; only ping users who opted in.
     allowed_user_ids: set[int] = set()
 
-    for match in MENTION_PATTERN.finditer(body):
+    for match in _MENTION_PATTERN.finditer(body):
         user_id = int(match.group("id"))
 
         if await flags.is_flag_set(db, user_id, EntityType.USER, Flag.ALLOW_PINGS):
@@ -169,7 +170,9 @@ async def ingest_message(
                 image_urls.append(embed.image.url)
 
         if image_urls:
-            await db.add_images([(channel_id, user_id, url) for url in image_urls])
+            await db.add_images(
+                [ChannelImage(channel_id, user_id, url) for url in image_urls]
+            )
 
     metrics.MESSAGES_INGESTED.inc()
 
@@ -208,7 +211,10 @@ async def uningest_message(db: Database, cache: IngestCache, message_id: int) ->
 
     if entry.image_urls:
         await db.remove_images(
-            [(entry.channel_id, entry.user_id, url) for url in entry.image_urls]
+            [
+                ChannelImage(entry.channel_id, entry.user_id, url)
+                for url in entry.image_urls
+            ]
         )
 
     metrics.MESSAGES_UNINGESTED.inc()

crabstero/metrics.py

@@ -131,6 +131,8 @@ class MetricsServer:
     which handles compression and content negotiation automatically.
     """
 
+    __slots__ = ("_host", "_port", "_runner")
+
     def __init__(self, host: str, port: int) -> None:
         """Store the listen address.
 

crabstero/tasks/ingestion.py

@@ -26,7 +26,7 @@ from crabstero import metrics
 from crabstero.messages import ingest_message
 
 if TYPE_CHECKING:
-    from crabstero.bot import Crabstero
+    from crabstero.bot import Crabstero, IngestableChannel
     from crabstero.database import Database
 
 MAXIMUM_MESSAGES_PER_CHANNEL = (
@@ -48,7 +48,7 @@ def queue_channels_for_ingestion(guild: discord.Guild, bot: Crabstero) -> None:
 
 
 async def ingest_channel(
-    channel: discord.TextChannel | discord.VoiceChannel,
+    channel: IngestableChannel,
     db: Database,
 ) -> None:
     """Bulk-ingest the message history of a given channel.