Improved code quality with more idiomatic Python patterns and safer initialization.
This commit is contained in:
+11
-1
@@ -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
@@ -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()
|
||||
|
||||
@@ -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."""
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
@@ -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))
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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."""
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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
@@ -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
|
||||
|
||||
@@ -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(
|
||||
|
||||
Reference in New Issue
Block a user