Improved code quality with more idiomatic Python patterns and safer initialization.
CI / Formatting (push) Successful in 5s
CI / Linting (push) Successful in 4s
CI / Tests (push) Successful in 21s
CI / Type Checking (push) Failing after 10s
CI / Spelling (push) Successful in 5s

This commit is contained in:
2026-03-22 14:41:38 -04:00
parent 371f7d2ca3
commit 2cf00fde4f
13 changed files with 148 additions and 79 deletions
+11 -1
View File
@@ -80,10 +80,20 @@ def _parse_args(argv: list[str] | None = None) -> argparse.Namespace:
' variable or "crabstero.db").'
),
)
workers_env = os.environ.get("INGESTION_WORKERS", "4")
try:
workers_default = int(workers_env)
except ValueError:
logger.warning(
"Invalid INGESTION_WORKERS value '%s', defaulting to 4.",
workers_env,
)
workers_default = 4
parser.add_argument(
"--ingestion-workers",
type=int,
default=int(os.environ.get("INGESTION_WORKERS", "4")),
default=workers_default,
help=(
"Number of concurrent ingestion workers"
" (default: INGESTION_WORKERS environment"
+31 -14
View File
@@ -30,6 +30,7 @@ from crabstero.cache import IngestCache
from crabstero.database import Database
from crabstero.listeners import interaction, message, server_events
from crabstero.metrics import (
CHANNEL_INGESTION_ERRORS,
DISCORD_EVENTS,
DISCORD_LATENCY,
GUILD_COUNT,
@@ -83,10 +84,11 @@ class Crabstero(commands.Bot):
discord.TextChannel | discord.VoiceChannel
] = asyncio.Queue()
self._ingestion_workers: list[asyncio.Task[None]] = []
self.db: Database
self._db: Database | None = None
self.ingest_cache = IngestCache()
self._metrics_address = metrics_address
self._metrics_server: MetricsServer | None = None
DISCORD_LATENCY.set_function(lambda: self.latency)
GUILD_COUNT.set_function(lambda: len(self.guilds))
@@ -95,15 +97,31 @@ class Crabstero(commands.Bot):
repo_url = "https://git.logal.dev/LogalDeveloper/Crabstero"
self.http.user_agent = f"DiscordBot ({repo_url}, {crabstero_version})"
@property
def db(self) -> Database:
"""The active database connection.
:raises RuntimeError: If accessed before :meth:`setup_hook` has run.
"""
if self._db is None:
raise RuntimeError("Database is not initialized")
return self._db
@property
def ingest_only(self) -> bool:
"""Whether the bot is running in ingest-only mode."""
return self._ingest_only
async def setup_hook(self) -> None:
"""Open the database, start ingestion workers, and load all cogs."""
self.db = await Database.connect(self._database_path)
self._db = await Database.connect(self._database_path)
self.ingest_cache.start()
self._start_ingestion_workers()
if self._metrics_address is not None:
self._metrics_server = MetricsServer(*self._metrics_address)
await self._metrics_server.start()
server = MetricsServer(*self._metrics_address)
await server.start()
self._metrics_server = server
if not self._ingest_only:
await interaction.setup(self)
@@ -143,13 +161,9 @@ class Crabstero(commands.Bot):
"""Log that the bot has started successfully."""
logger.info("Crabstero started!")
async def start(self, token: str = "", *, reconnect: bool = True) -> None:
"""Start the bot using the stored token by default.
:param token: Optional token override. Falls back to the stored token if empty.
:param reconnect: Whether to automatically reconnect on disconnect.
"""
await super().start(token or self._token, reconnect=reconnect)
async def start(self, **kwargs: object) -> None:
"""Start the bot using the token provided at initialization."""
await super().start(self._token, **kwargs)
async def close(self) -> None:
"""Cancel ingestion workers, close the database, and then the bot connection."""
@@ -161,10 +175,10 @@ class Crabstero(commands.Bot):
await asyncio.gather(*self._ingestion_workers, return_exceptions=True)
self._ingestion_workers.clear()
self.ingest_cache.stop()
if hasattr(self, "_metrics_server"):
if self._metrics_server is not None:
await self._metrics_server.stop()
if hasattr(self, "db"):
await self.db.close()
if self._db is not None:
await self._db.close()
await super().close()
def queue_channel_for_ingestion(
@@ -188,5 +202,8 @@ class Crabstero(commands.Bot):
channel = await self._ingestion_queue.get()
try:
await ingest_channel(channel, self.db)
except Exception:
CHANNEL_INGESTION_ERRORS.inc()
logger.exception("Ingestion failed for channel %s", channel.id)
finally:
self._ingestion_queue.task_done()
+14
View File
@@ -1 +1,15 @@
# Copyright 2026 Logan Fick
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Listener cogs for Discord events."""
+1 -1
View File
@@ -47,7 +47,7 @@ class InteractionCog(commands.Cog):
name="pingme",
description=(
"Opts into (or back out of) receiving pings"
" for generated message which mention you."
" for generated messages which mention you."
),
)
async def pingme(self, interaction: discord.Interaction) -> None:
+1 -1
View File
@@ -58,7 +58,7 @@ class MessageCog(commands.Cog):
metrics.MESSAGES_PROCESSED.inc()
if self.bot._ingest_only:
if self.bot.ingest_only:
# In ingest-only mode, never reply — only ingest eligible messages.
if message.type == discord.MessageType.default and not isinstance(
channel, discord.Thread
+2 -2
View File
@@ -62,13 +62,13 @@ class ServerEventsCog(commands.Cog):
value=f"{guild.member_count} members",
)
if guild.icon:
if guild.icon is not None:
embed.set_image(url=guild.icon.url)
embed.set_footer(text=f"{len(self.bot.guilds)} total servers")
app_info = await self.bot.application_info()
if app_info.owner:
if app_info.owner is not None:
with contextlib.suppress(discord.HTTPException):
await app_info.owner.send(embed=embed)
+25 -19
View File
@@ -19,6 +19,7 @@ to represent frequency weight. Random selection via ORDER BY RANDOM()
LIMIT 1 naturally preserves this weighting.
"""
import itertools
import re
from typing import TYPE_CHECKING
@@ -31,6 +32,19 @@ if TYPE_CHECKING:
# as is common on Discord.
DEFAULT_SENTENCE_END = "\u00a7"
_TERMINATORS = frozenset({DEFAULT_SENTENCE_END, ".", "!", "?"})
_MULTI_SPACE = re.compile(r" +")
_SENTENCE_SPLIT = re.compile(r"(?<=[.!?]) ")
def _normalize_whitespace(text: str) -> str:
"""Collapse runs of spaces into single spaces and strip edges.
:param text: The raw text to normalize.
:return: The normalized text.
"""
return _MULTI_SPACE.sub(" ", text.strip())
def is_complete_sentence(sentence: str) -> bool:
"""Check whether a sentence ends with a valid terminator.
@@ -44,7 +58,7 @@ def is_complete_sentence(sentence: str) -> bool:
if not sentence:
return False
return sentence[-1] in (DEFAULT_SENTENCE_END, ".", "!", "?")
return sentence[-1] in _TERMINATORS
def _split_sentences(paragraph: str) -> list[str]:
@@ -55,8 +69,8 @@ def _split_sentences(paragraph: str) -> list[str]:
"""
if not is_complete_sentence(paragraph):
paragraph += DEFAULT_SENTENCE_END
normalized = re.sub(r" +", " ", paragraph.strip().replace("\n", " "))
return re.split(r"(?<=[.!?]) ", normalized)
normalized = _normalize_whitespace(paragraph.replace("\n", " "))
return _SENTENCE_SPLIT.split(normalized)
def _tokenize_sentence(
@@ -69,14 +83,10 @@ def _tokenize_sentence(
"""
if not is_complete_sentence(sentence):
sentence += DEFAULT_SENTENCE_END
words = re.sub(r" +", " ", sentence.strip()).split(" ")
words = _normalize_whitespace(sentence).split()
start_words: list[str] = []
transitions: list[tuple[str, str]] = []
for i in range(len(words) - 1):
if i == 0:
start_words.append(words[i])
transitions.append((words[i], words[i + 1]))
start_words: list[str] = [words[0]] if len(words) >= 2 else []
transitions: list[tuple[str, str]] = list(itertools.pairwise(words))
return start_words, transitions
@@ -162,8 +172,7 @@ async def generate(
" Chat a bit more so I can learn how this channel talks!"
)
parts: list[str] = []
parts.append(word)
parts: list[str] = [word]
current_length = len(word)
# The loop is skipped if the start word already ends a sentence (e.g. "Yes.").
@@ -185,17 +194,14 @@ async def generate(
current_length += 1 + len(word) # +1 for the joining space.
if current_length >= hard_limit:
result = " ".join(parts)[:hard_limit]
# Strip the internal sentence-end marker if it ended up at the boundary.
if result and result[-1] == DEFAULT_SENTENCE_END:
result = result[:-1]
metrics.GENERATED_MESSAGE_LENGTH.observe(len(result))
return result
break
result = " ".join(parts)
if current_length >= hard_limit:
result = result[:hard_limit]
# Strip the internal sentence-end marker so it never appears in output.
if result and result[-1] == DEFAULT_SENTENCE_END:
if result.endswith(DEFAULT_SENTENCE_END):
result = result[:-1]
metrics.GENERATED_MESSAGE_LENGTH.observe(len(result))
+5 -3
View File
@@ -66,7 +66,7 @@ async def reply_to_message(db: Database, message: discord.Message) -> None:
else:
channel_id = channel.id
body = await markov.generate(db, channel_id, 750, 1000)
body = await markov.generate(db, channel_id)
embed = None
@@ -76,8 +76,10 @@ async def reply_to_message(db: Database, message: discord.Message) -> None:
and channel.permissions_for(guild.me).embed_links
):
embed = discord.Embed(
title=await markov.generate(db, channel_id, 200, 300),
description=await markov.generate(db, channel_id, 300, 500),
title=await markov.generate(db, channel_id, soft_limit=200, hard_limit=300),
description=await markov.generate(
db, channel_id, soft_limit=300, hard_limit=500
),
)
random_image = await db.get_random_image(channel_id)
+14
View File
@@ -1 +1,15 @@
# Copyright 2026 Logan Fick
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Background tasks for Crabstero."""
+27 -35
View File
@@ -59,40 +59,32 @@ async def ingest_channel(
:param channel: The channel to ingest.
:param db: The database instance.
"""
try:
if not channel.permissions_for(channel.guild.me).read_message_history:
logger.warning(
"[%s] Unable to ingest channel history"
" due to lacking permissions. Ignoring.",
channel.id,
)
return
if await db.is_channel_ingested(channel.id):
return
await db.mark_channel_ingested(channel.id)
logger.info("[%s] Starting ingestion of textable channel history.", channel.id)
with metrics.CHANNEL_INGESTION_DURATION.time():
count = 0
async for message in channel.history(limit=MAXIMUM_MESSAGES_PER_CHANNEL):
count += 1
await ingest_message(db, message)
metrics.CHANNEL_INGESTION_MESSAGES.observe(count)
logger.info(
"[%s] Ingestion of channel history complete. %d messages ingested.",
channel.id,
count,
)
metrics.CHANNELS_INGESTED.inc()
except Exception:
metrics.CHANNEL_INGESTION_ERRORS.inc()
logger.exception(
"[%s] An error occurred while ingesting textable channel history!",
if not channel.permissions_for(channel.guild.me).read_message_history:
logger.warning(
"[%s] Unable to ingest channel history"
" due to lacking permissions. Ignoring.",
channel.id,
)
return
if await db.is_channel_ingested(channel.id):
return
await db.mark_channel_ingested(channel.id)
logger.info("[%s] Starting ingestion of textable channel history.", channel.id)
with metrics.CHANNEL_INGESTION_DURATION.time():
count = 0
async for message in channel.history(limit=MAXIMUM_MESSAGES_PER_CHANNEL):
count += 1
await ingest_message(db, message)
metrics.CHANNEL_INGESTION_MESSAGES.observe(count)
logger.info(
"[%s] Ingestion of channel history complete. %d messages ingested.",
channel.id,
count,
)
metrics.CHANNELS_INGESTED.inc()
+14
View File
@@ -1 +1,15 @@
# Copyright 2026 Logan Fick
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Test suite for Crabstero."""
+2 -2
View File
@@ -21,12 +21,12 @@ import pytest
from crabstero.database import Database
if TYPE_CHECKING:
from collections.abc import AsyncIterator
from collections.abc import AsyncGenerator
from pathlib import Path
@pytest.fixture
async def db(tmp_path: Path) -> AsyncIterator[Database]:
async def db(tmp_path: Path) -> AsyncGenerator[Database]:
"""Yield a Database backed by a temporary SQLite file."""
database = await Database.connect(str(tmp_path / "test.db"))
yield database
+1 -1
View File
@@ -55,7 +55,7 @@ class TestIngestUningestCycle:
assert await db.get_random_image(1) is None
async def test_cache_round_trip(self) -> None:
def test_cache_round_trip(self) -> None:
"""Cache put then pop returns the original entry."""
cache = IngestCache()
entry = CachedMessage(