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").'
|
' 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(
|
parser.add_argument(
|
||||||
"--ingestion-workers",
|
"--ingestion-workers",
|
||||||
type=int,
|
type=int,
|
||||||
default=int(os.environ.get("INGESTION_WORKERS", "4")),
|
default=workers_default,
|
||||||
help=(
|
help=(
|
||||||
"Number of concurrent ingestion workers"
|
"Number of concurrent ingestion workers"
|
||||||
" (default: INGESTION_WORKERS environment"
|
" (default: INGESTION_WORKERS environment"
|
||||||
|
|||||||
+31
-14
@@ -30,6 +30,7 @@ from crabstero.cache import IngestCache
|
|||||||
from crabstero.database import Database
|
from crabstero.database import Database
|
||||||
from crabstero.listeners import interaction, message, server_events
|
from crabstero.listeners import interaction, message, server_events
|
||||||
from crabstero.metrics import (
|
from crabstero.metrics import (
|
||||||
|
CHANNEL_INGESTION_ERRORS,
|
||||||
DISCORD_EVENTS,
|
DISCORD_EVENTS,
|
||||||
DISCORD_LATENCY,
|
DISCORD_LATENCY,
|
||||||
GUILD_COUNT,
|
GUILD_COUNT,
|
||||||
@@ -83,10 +84,11 @@ class Crabstero(commands.Bot):
|
|||||||
discord.TextChannel | discord.VoiceChannel
|
discord.TextChannel | discord.VoiceChannel
|
||||||
] = asyncio.Queue()
|
] = asyncio.Queue()
|
||||||
self._ingestion_workers: list[asyncio.Task[None]] = []
|
self._ingestion_workers: list[asyncio.Task[None]] = []
|
||||||
self.db: Database
|
self._db: Database | None = None
|
||||||
self.ingest_cache = IngestCache()
|
self.ingest_cache = IngestCache()
|
||||||
|
|
||||||
self._metrics_address = metrics_address
|
self._metrics_address = metrics_address
|
||||||
|
self._metrics_server: MetricsServer | None = None
|
||||||
|
|
||||||
DISCORD_LATENCY.set_function(lambda: self.latency)
|
DISCORD_LATENCY.set_function(lambda: self.latency)
|
||||||
GUILD_COUNT.set_function(lambda: len(self.guilds))
|
GUILD_COUNT.set_function(lambda: len(self.guilds))
|
||||||
@@ -95,15 +97,31 @@ class Crabstero(commands.Bot):
|
|||||||
repo_url = "https://git.logal.dev/LogalDeveloper/Crabstero"
|
repo_url = "https://git.logal.dev/LogalDeveloper/Crabstero"
|
||||||
self.http.user_agent = f"DiscordBot ({repo_url}, {crabstero_version})"
|
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:
|
async def setup_hook(self) -> None:
|
||||||
"""Open the database, start ingestion workers, and load all cogs."""
|
"""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.ingest_cache.start()
|
||||||
self._start_ingestion_workers()
|
self._start_ingestion_workers()
|
||||||
|
|
||||||
if self._metrics_address is not None:
|
if self._metrics_address is not None:
|
||||||
self._metrics_server = MetricsServer(*self._metrics_address)
|
server = MetricsServer(*self._metrics_address)
|
||||||
await self._metrics_server.start()
|
await server.start()
|
||||||
|
self._metrics_server = server
|
||||||
|
|
||||||
if not self._ingest_only:
|
if not self._ingest_only:
|
||||||
await interaction.setup(self)
|
await interaction.setup(self)
|
||||||
@@ -143,13 +161,9 @@ class Crabstero(commands.Bot):
|
|||||||
"""Log that the bot has started successfully."""
|
"""Log that the bot has started successfully."""
|
||||||
logger.info("Crabstero started!")
|
logger.info("Crabstero started!")
|
||||||
|
|
||||||
async def start(self, token: str = "", *, reconnect: bool = True) -> None:
|
async def start(self, **kwargs: object) -> None:
|
||||||
"""Start the bot using the stored token by default.
|
"""Start the bot using the token provided at initialization."""
|
||||||
|
await super().start(self._token, **kwargs)
|
||||||
: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 close(self) -> None:
|
async def close(self) -> None:
|
||||||
"""Cancel ingestion workers, close the database, and then the bot connection."""
|
"""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)
|
await asyncio.gather(*self._ingestion_workers, return_exceptions=True)
|
||||||
self._ingestion_workers.clear()
|
self._ingestion_workers.clear()
|
||||||
self.ingest_cache.stop()
|
self.ingest_cache.stop()
|
||||||
if hasattr(self, "_metrics_server"):
|
if self._metrics_server is not None:
|
||||||
await self._metrics_server.stop()
|
await self._metrics_server.stop()
|
||||||
if hasattr(self, "db"):
|
if self._db is not None:
|
||||||
await self.db.close()
|
await self._db.close()
|
||||||
await super().close()
|
await super().close()
|
||||||
|
|
||||||
def queue_channel_for_ingestion(
|
def queue_channel_for_ingestion(
|
||||||
@@ -188,5 +202,8 @@ class Crabstero(commands.Bot):
|
|||||||
channel = await self._ingestion_queue.get()
|
channel = await self._ingestion_queue.get()
|
||||||
try:
|
try:
|
||||||
await ingest_channel(channel, self.db)
|
await ingest_channel(channel, self.db)
|
||||||
|
except Exception:
|
||||||
|
CHANNEL_INGESTION_ERRORS.inc()
|
||||||
|
logger.exception("Ingestion failed for channel %s", channel.id)
|
||||||
finally:
|
finally:
|
||||||
self._ingestion_queue.task_done()
|
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."""
|
"""Listener cogs for Discord events."""
|
||||||
|
|||||||
@@ -47,7 +47,7 @@ class InteractionCog(commands.Cog):
|
|||||||
name="pingme",
|
name="pingme",
|
||||||
description=(
|
description=(
|
||||||
"Opts into (or back out of) receiving pings"
|
"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:
|
async def pingme(self, interaction: discord.Interaction) -> None:
|
||||||
|
|||||||
@@ -58,7 +58,7 @@ class MessageCog(commands.Cog):
|
|||||||
|
|
||||||
metrics.MESSAGES_PROCESSED.inc()
|
metrics.MESSAGES_PROCESSED.inc()
|
||||||
|
|
||||||
if self.bot._ingest_only:
|
if self.bot.ingest_only:
|
||||||
# In ingest-only mode, never reply — only ingest eligible messages.
|
# In ingest-only mode, never reply — only ingest eligible messages.
|
||||||
if message.type == discord.MessageType.default and not isinstance(
|
if message.type == discord.MessageType.default and not isinstance(
|
||||||
channel, discord.Thread
|
channel, discord.Thread
|
||||||
|
|||||||
@@ -62,13 +62,13 @@ class ServerEventsCog(commands.Cog):
|
|||||||
value=f"{guild.member_count} members",
|
value=f"{guild.member_count} members",
|
||||||
)
|
)
|
||||||
|
|
||||||
if guild.icon:
|
if guild.icon is not None:
|
||||||
embed.set_image(url=guild.icon.url)
|
embed.set_image(url=guild.icon.url)
|
||||||
|
|
||||||
embed.set_footer(text=f"{len(self.bot.guilds)} total servers")
|
embed.set_footer(text=f"{len(self.bot.guilds)} total servers")
|
||||||
|
|
||||||
app_info = await self.bot.application_info()
|
app_info = await self.bot.application_info()
|
||||||
if app_info.owner:
|
if app_info.owner is not None:
|
||||||
with contextlib.suppress(discord.HTTPException):
|
with contextlib.suppress(discord.HTTPException):
|
||||||
await app_info.owner.send(embed=embed)
|
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.
|
LIMIT 1 naturally preserves this weighting.
|
||||||
"""
|
"""
|
||||||
|
|
||||||
|
import itertools
|
||||||
import re
|
import re
|
||||||
from typing import TYPE_CHECKING
|
from typing import TYPE_CHECKING
|
||||||
|
|
||||||
@@ -31,6 +32,19 @@ if TYPE_CHECKING:
|
|||||||
# as is common on Discord.
|
# as is common on Discord.
|
||||||
DEFAULT_SENTENCE_END = "\u00a7"
|
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:
|
def is_complete_sentence(sentence: str) -> bool:
|
||||||
"""Check whether a sentence ends with a valid terminator.
|
"""Check whether a sentence ends with a valid terminator.
|
||||||
@@ -44,7 +58,7 @@ def is_complete_sentence(sentence: str) -> bool:
|
|||||||
if not sentence:
|
if not sentence:
|
||||||
return False
|
return False
|
||||||
|
|
||||||
return sentence[-1] in (DEFAULT_SENTENCE_END, ".", "!", "?")
|
return sentence[-1] in _TERMINATORS
|
||||||
|
|
||||||
|
|
||||||
def _split_sentences(paragraph: str) -> list[str]:
|
def _split_sentences(paragraph: str) -> list[str]:
|
||||||
@@ -55,8 +69,8 @@ def _split_sentences(paragraph: str) -> list[str]:
|
|||||||
"""
|
"""
|
||||||
if not is_complete_sentence(paragraph):
|
if not is_complete_sentence(paragraph):
|
||||||
paragraph += DEFAULT_SENTENCE_END
|
paragraph += DEFAULT_SENTENCE_END
|
||||||
normalized = re.sub(r" +", " ", paragraph.strip().replace("\n", " "))
|
normalized = _normalize_whitespace(paragraph.replace("\n", " "))
|
||||||
return re.split(r"(?<=[.!?]) ", normalized)
|
return _SENTENCE_SPLIT.split(normalized)
|
||||||
|
|
||||||
|
|
||||||
def _tokenize_sentence(
|
def _tokenize_sentence(
|
||||||
@@ -69,14 +83,10 @@ def _tokenize_sentence(
|
|||||||
"""
|
"""
|
||||||
if not is_complete_sentence(sentence):
|
if not is_complete_sentence(sentence):
|
||||||
sentence += DEFAULT_SENTENCE_END
|
sentence += DEFAULT_SENTENCE_END
|
||||||
words = re.sub(r" +", " ", sentence.strip()).split(" ")
|
words = _normalize_whitespace(sentence).split()
|
||||||
|
|
||||||
start_words: list[str] = []
|
start_words: list[str] = [words[0]] if len(words) >= 2 else []
|
||||||
transitions: list[tuple[str, str]] = []
|
transitions: list[tuple[str, str]] = list(itertools.pairwise(words))
|
||||||
for i in range(len(words) - 1):
|
|
||||||
if i == 0:
|
|
||||||
start_words.append(words[i])
|
|
||||||
transitions.append((words[i], words[i + 1]))
|
|
||||||
return start_words, transitions
|
return start_words, transitions
|
||||||
|
|
||||||
|
|
||||||
@@ -162,8 +172,7 @@ async def generate(
|
|||||||
" Chat a bit more so I can learn how this channel talks!"
|
" Chat a bit more so I can learn how this channel talks!"
|
||||||
)
|
)
|
||||||
|
|
||||||
parts: list[str] = []
|
parts: list[str] = [word]
|
||||||
parts.append(word)
|
|
||||||
current_length = len(word)
|
current_length = len(word)
|
||||||
|
|
||||||
# The loop is skipped if the start word already ends a sentence (e.g. "Yes.").
|
# 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.
|
current_length += 1 + len(word) # +1 for the joining space.
|
||||||
|
|
||||||
if current_length >= hard_limit:
|
if current_length >= hard_limit:
|
||||||
result = " ".join(parts)[:hard_limit]
|
break
|
||||||
# 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
|
|
||||||
|
|
||||||
result = " ".join(parts)
|
result = " ".join(parts)
|
||||||
|
if current_length >= hard_limit:
|
||||||
|
result = result[:hard_limit]
|
||||||
|
|
||||||
# Strip the internal sentence-end marker so it never appears in output.
|
# 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]
|
result = result[:-1]
|
||||||
|
|
||||||
metrics.GENERATED_MESSAGE_LENGTH.observe(len(result))
|
metrics.GENERATED_MESSAGE_LENGTH.observe(len(result))
|
||||||
|
|||||||
@@ -66,7 +66,7 @@ async def reply_to_message(db: Database, message: discord.Message) -> None:
|
|||||||
else:
|
else:
|
||||||
channel_id = channel.id
|
channel_id = channel.id
|
||||||
|
|
||||||
body = await markov.generate(db, channel_id, 750, 1000)
|
body = await markov.generate(db, channel_id)
|
||||||
|
|
||||||
embed = None
|
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
|
and channel.permissions_for(guild.me).embed_links
|
||||||
):
|
):
|
||||||
embed = discord.Embed(
|
embed = discord.Embed(
|
||||||
title=await markov.generate(db, channel_id, 200, 300),
|
title=await markov.generate(db, channel_id, soft_limit=200, hard_limit=300),
|
||||||
description=await markov.generate(db, channel_id, 300, 500),
|
description=await markov.generate(
|
||||||
|
db, channel_id, soft_limit=300, hard_limit=500
|
||||||
|
),
|
||||||
)
|
)
|
||||||
|
|
||||||
random_image = await db.get_random_image(channel_id)
|
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."""
|
"""Background tasks for Crabstero."""
|
||||||
|
|||||||
@@ -59,7 +59,6 @@ async def ingest_channel(
|
|||||||
:param channel: The channel to ingest.
|
:param channel: The channel to ingest.
|
||||||
:param db: The database instance.
|
:param db: The database instance.
|
||||||
"""
|
"""
|
||||||
try:
|
|
||||||
if not channel.permissions_for(channel.guild.me).read_message_history:
|
if not channel.permissions_for(channel.guild.me).read_message_history:
|
||||||
logger.warning(
|
logger.warning(
|
||||||
"[%s] Unable to ingest channel history"
|
"[%s] Unable to ingest channel history"
|
||||||
@@ -89,10 +88,3 @@ async def ingest_channel(
|
|||||||
count,
|
count,
|
||||||
)
|
)
|
||||||
metrics.CHANNELS_INGESTED.inc()
|
metrics.CHANNELS_INGESTED.inc()
|
||||||
|
|
||||||
except Exception:
|
|
||||||
metrics.CHANNEL_INGESTION_ERRORS.inc()
|
|
||||||
logger.exception(
|
|
||||||
"[%s] An error occurred while ingesting textable channel history!",
|
|
||||||
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.
|
||||||
|
|
||||||
"""Test suite for Crabstero."""
|
"""Test suite for Crabstero."""
|
||||||
|
|||||||
+2
-2
@@ -21,12 +21,12 @@ import pytest
|
|||||||
from crabstero.database import Database
|
from crabstero.database import Database
|
||||||
|
|
||||||
if TYPE_CHECKING:
|
if TYPE_CHECKING:
|
||||||
from collections.abc import AsyncIterator
|
from collections.abc import AsyncGenerator
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
|
|
||||||
|
|
||||||
@pytest.fixture
|
@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."""
|
"""Yield a Database backed by a temporary SQLite file."""
|
||||||
database = await Database.connect(str(tmp_path / "test.db"))
|
database = await Database.connect(str(tmp_path / "test.db"))
|
||||||
yield database
|
yield database
|
||||||
|
|||||||
@@ -55,7 +55,7 @@ class TestIngestUningestCycle:
|
|||||||
|
|
||||||
assert await db.get_random_image(1) is None
|
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 put then pop returns the original entry."""
|
||||||
cache = IngestCache()
|
cache = IngestCache()
|
||||||
entry = CachedMessage(
|
entry = CachedMessage(
|
||||||
|
|||||||
Reference in New Issue
Block a user