Merge branch 'master' into clustering

admin_panel/preview/management/commands/preview.py

@@ -0,0 +1,81 @@
+import asyncio
+import os
+import sys
+
+from django.core.management.base import BaseCommand, CommandError, CommandParser
+from tortoise.exceptions import DoesNotExist
+
+from ballsdex.core.image_generator.image_gen import draw_card
+from ballsdex.core.models import Ball, BallInstance, Special
+from ballsdex.settings import settings
+
+from ...utils import refresh_cache
+
+
+class Command(BaseCommand):
+    help = (
+        "Generate a local preview of a card. This will use the system's image viewer "
+        "or print to stdout if the output is being piped."
+    )
+
+    def add_arguments(self, parser: CommandParser):
+        parser.add_argument(
+            "--ball",
+            help=f"The name of the {settings.collectible_name} you want to generate. "
+            "If not provided, the first entry is used.",
+        )
+        parser.add_argument(
+            "--special",
+            help="The special event's background you want to use, otherwise regime is used",
+        )
+
+    async def generate_preview(self, *args, **options):
+        await refresh_cache()
+
+        if ball_name := options.get("ball"):
+            try:
+                ball = await Ball.get(country__iexact=ball_name)
+            except DoesNotExist as e:
+                raise CommandError(
+                    f'No {settings.collectible_name} found with the name "{ball_name}"'
+                ) from e
+        else:
+            ball = await Ball.first()
+            if ball is None:
+                raise CommandError(f"You need at least one {settings.collectible_name} created.")
+
+        special = None
+        if special_name := options.get("special"):
+            try:
+                special = await Special.get(name__iexact=special_name)
+            except DoesNotExist as e:
+                raise CommandError(f'No special found with the name "{special_name}"') from e
+
+        # use stderr to avoid piping
+        self.stderr.write(
+            self.style.SUCCESS(
+                f"Generating card for {ball.country}" + (f" ({special.name})" if special else "")
+            )
+        )
+
+        instance = BallInstance(ball=ball, special=special)
+        image = draw_card(instance, media_path="./media/")
+
+        if sys.platform not in ("win32", "darwin") and not os.environ.get("DISPLAY"):
+            self.stderr.write(
+                self.style.WARNING(
+                    "\nThis command displays the generated card using your system's image viewer, "
+                    "but no display was detected. Are you running this inside Docker?\n"
+                    'You can append "> image.png" at the end of your command to instead write the '
+                    "image to disk, which you can then open manually.\n"
+                )
+            )
+            raise CommandError("No display detected.")
+        if sys.stdout.isatty():
+            image.show(title=ball.country)
+        else:
+            image.save(sys.stdout.buffer, "png")
+
+    def handle(self, *args, **options):
+        loop = asyncio.get_event_loop()
+        loop.run_until_complete(self.generate_preview(*args, **options))

admin_panel/preview/utils.py

@@ -0,0 +1,42 @@
+import os
+
+from tortoise import Tortoise
+
+from ballsdex.__main__ import init_tortoise
+from ballsdex.core.models import (
+    Ball,
+    Economy,
+    Regime,
+    Special,
+    balls,
+    economies,
+    regimes,
+    specials,
+)
+
+
+async def refresh_cache():
+    """
+    Similar to the bot's `load_cache` function without the fancy display. Also handles
+    initializing the connection to Tortoise.
+
+    This must be called on every request, since the image generation relies on cache and we
+    do *not* want caching in the admin panel to happen (since we're actively editing stuff).
+    """
+    if not Tortoise._inited:
+        await init_tortoise(os.environ["BALLSDEXBOT_DB_URL"], skip_migrations=True)
+    balls.clear()
+    for ball in await Ball.all():
+        balls[ball.pk] = ball
+
+    regimes.clear()
+    for regime in await Regime.all():
+        regimes[regime.pk] = regime
+
+    economies.clear()
+    for economy in await Economy.all():
+        economies[economy.pk] = economy
+
+    specials.clear()
+    for special in await Special.all():
+        specials[special.pk] = special

admin_panel/preview/views.py

@@ -1,53 +1,14 @@
-import os
-
 from django.contrib import messages
 from django.http import HttpRequest, HttpResponse
-from tortoise import Tortoise
 
-from ballsdex.__main__ import init_tortoise
 from ballsdex.core.image_generator.image_gen import draw_card
-from ballsdex.core.models import (
-    Ball,
-    BallInstance,
-    Economy,
-    Regime,
-    Special,
-    balls,
-    economies,
-    regimes,
-    specials,
-)
-
-
-async def _refresh_cache():
-    """
-    Similar to the bot's `load_cache` function without the fancy display. Also handles
-    initializing the connection to Tortoise.
-
-    This must be called on every request, since the image generation relies on cache and we
-    do *not* want caching in the admin panel to happen (since we're actively editing stuff).
-    """
-    if not Tortoise._inited:
-        await init_tortoise(os.environ["BALLSDEXBOT_DB_URL"], skip_migrations=True)
-    balls.clear()
-    for ball in await Ball.all():
-        balls[ball.pk] = ball
-
-    regimes.clear()
-    for regime in await Regime.all():
-        regimes[regime.pk] = regime
-
-    economies.clear()
-    for economy in await Economy.all():
-        economies[economy.pk] = economy
+from ballsdex.core.models import Ball, BallInstance, Special
 
-    specials.clear()
-    for special in await Special.all():
-        specials[special.pk] = special
+from .utils import refresh_cache
 
 
 async def render_ballinstance(request: HttpRequest, ball_pk: int) -> HttpResponse:
-    await _refresh_cache()
+    await refresh_cache()
 
     ball = await Ball.get(pk=ball_pk)
     instance = BallInstance(ball=ball)
@@ -59,7 +20,7 @@ async def render_ballinstance(request: HttpRequest, ball_pk: int) -> HttpRespons
 
 
 async def render_special(request: HttpRequest, special_pk: int) -> HttpResponse:
-    await _refresh_cache()
+    await refresh_cache()
 
     ball = await Ball.first()
     if ball is None:

ballsdex/__init__.py

@@ -1 +1 @@
-__version__ = "2.25.1"
+__version__ = "2.25.3"

ballsdex/core/image_generator/image_gen.py

@@ -19,6 +19,18 @@
 CORNERS = ((34, 261), (1393, 992))
 artwork_size = [b - a for a, b in zip(*CORNERS)]
 
+# ===== TIP =====
+#
+# If you want to quickly test the image generation, there is a CLI tool to quickly generate
+# test images locally, without the bot or the admin panel running:
+#
+# With Docker: "docker compose run admin-panel python3 manage.py preview > image.png"
+# Without: "cd admin_panel && poetry run python3 manage.py preview"
+#
+# This will either create a file named "image.png" or directly display it using your system's
+# image viewer. There are options available to specify the ball or the special background,
+# use the "--help" flag to view all options.
+
 title_font = ImageFont.truetype(str(SOURCES_PATH / "ArsenicaTrial-Extrabold.ttf"), 170)
 capacity_name_font = ImageFont.truetype(str(SOURCES_PATH / "Bobby Jones Soft.otf"), 110)
 capacity_description_font = ImageFont.truetype(str(SOURCES_PATH / "OpenSans-Semibold.ttf"), 75)

ballsdex/core/utils/sorting.py

@@ -59,9 +59,16 @@ def sort_balls(
             **{f"{sort.value}_sort": F(f"{sort.value}_bonus") + F(f"ball__{sort.value}")}
         ).order_by(f"-{sort.value}_sort")
     elif sort == SortingChoices.total_stats:
-        return queryset.annotate(
-            stats=F("health_bonus") + F("ball__health") + F("attack_bonus") + F("ball__attack")
-        ).order_by("-stats")
+        return (
+            queryset.select_related("ball")
+            .annotate(
+                stats=RawSQL(
+                    "health_bonus + ballinstance__ball.health + "
+                    "attack_bonus + ballinstance__ball.attack :: BIGINT"
+                )
+            )
+            .order_by("-stats")
+        )
     elif sort == SortingChoices.rarity:
         return queryset.order_by(sort.value, "ball__country")
     else:

ballsdex/core/utils/utils.py

@@ -9,7 +9,9 @@
     from ballsdex.core.bot import BallsDexBot
 
 
-def is_staff(interaction: discord.Interaction) -> bool:
+def is_staff(interaction: discord.Interaction["BallsDexBot"]) -> bool:
+    if interaction.user.id in interaction.client.owner_ids:
+        return True
     if interaction.guild and interaction.guild.id in settings.admin_guild_ids:
         roles = settings.admin_role_ids + settings.root_role_ids
         if any(role.id in roles for role in interaction.user.roles):  # type: ignore
@@ -19,7 +21,7 @@ def is_staff(interaction: discord.Interaction) -> bool:
 
 async def inventory_privacy(
     bot: "BallsDexBot",
-    interaction: discord.Interaction,
+    interaction: discord.Interaction["BallsDexBot"],
     player: Player,
     user_obj: Union[discord.User, discord.Member],
 ):

ballsdex/packages/admin/balls.py

@@ -3,6 +3,7 @@
 import random
 import re
 from pathlib import Path
+from typing import TYPE_CHECKING, cast
 
 import discord
 from discord import app_commands
@@ -19,9 +20,12 @@
     RegimeTransform,
     SpecialTransform,
 )
-from ballsdex.packages.countryballs.countryball import CountryBall
 from ballsdex.settings import settings
 
+if TYPE_CHECKING:
+    from ballsdex.packages.countryballs.cog import CountryBallsSpawner
+    from ballsdex.packages.countryballs.countryball import CountryBall
+
 log = logging.getLogger("ballsdex.packages.admin.balls")
 FILENAME_RE = re.compile(r"^(.+)(\.\S+)$")
 
@@ -47,6 +51,7 @@ class Balls(app_commands.Group):
     async def _spawn_bomb(
         self,
         interaction: discord.Interaction[BallsDexBot],
+        countryball_cls: type["CountryBall"],
         countryball: Ball | None,
         channel: discord.TextChannel,
         n: int,
@@ -77,9 +82,9 @@ async def update_message_loop():
         try:
             for i in range(n):
                 if not countryball:
-                    ball = await CountryBall.get_random()
+                    ball = await countryball_cls.get_random()
                 else:
-                    ball = CountryBall(countryball)
+                    ball = countryball_cls(countryball)
                 ball.special = special
                 ball.atk_bonus = atk_bonus
                 ball.hp_bonus = hp_bonus
@@ -137,10 +142,29 @@ async def spawn(
         # the transformer triggered a response, meaning user tried an incorrect input
         if interaction.response.is_done():
             return
+        cog = cast("CountryBallsSpawner | None", interaction.client.get_cog("CountryBallsSpawner"))
+        if not cog:
+            prefix = (
+                settings.prefix
+                if interaction.client.intents.message_content or not interaction.client.user
+                else f"{interaction.client.user.mention} "
+            )
+            # do not replace `countryballs` with `settings.collectible_name`, it is intended
+            await interaction.response.send_message(
+                "The `countryballs` package is not loaded, this command is unavailable.\n"
+                "Please resolve the errors preventing this package from loading. Use "
+                f'"{prefix}reload countryballs" to try reloading it.',
+                ephemeral=True,
+            )
+            return
 
         if n > 1:
             await self._spawn_bomb(
-                interaction, countryball, channel or interaction.channel, n  # type: ignore
+                interaction,
+                cog.countryball_cls,
+                countryball,
+                channel or interaction.channel,  # type: ignore
+                n,
             )
             await log_action(
                 f"{interaction.user} spawned {settings.collectible_name}"
@@ -152,9 +176,9 @@ async def spawn(
 
         await interaction.response.defer(ephemeral=True, thinking=True)
         if not countryball:
-            ball = await CountryBall.get_random()
+            ball = await cog.countryball_cls.get_random()
         else:
-            ball = CountryBall(countryball)
+            ball = cog.countryball_cls(countryball)
         ball.special = special
         ball.atk_bonus = atk_bonus
         ball.hp_bonus = hp_bonus

ballsdex/packages/balls/cog.py

@@ -376,7 +376,9 @@ async def info(
 
     @app_commands.command()
     @app_commands.checks.cooldown(1, 60, key=lambda i: i.user.id)
-    async def last(self, interaction: discord.Interaction, user: discord.User | None = None):
+    async def last(
+        self, interaction: discord.Interaction["BallsDexBot"], user: discord.User | None = None
+    ):
         """
         Display info of your or another users last caught countryball.
 

ballsdex/packages/countryballs/cog.py

@@ -23,6 +23,7 @@ class CountryBallsSpawner(commands.Cog):
     def __init__(self, bot: "BallsDexBot"):
         self.bot = bot
         self.cache: dict[int, int] = {}
+        self.countryball_cls = CountryBall
 
         module_path, class_name = settings.spawn_manager.rsplit(".", 1)
         module = importlib.import_module(module_path)

ballsdex/packages/countryballs/components.py

@@ -114,11 +114,15 @@ async def catch_ball(
         player, created = await Player.get_or_create(discord_id=user.id)
 
         # stat may vary by +/- 20% of base stat
-        bonus_attack = self.ball.atk_bonus or random.randint(
-            -settings.max_attack_bonus, settings.max_attack_bonus
+        bonus_attack = (
+            self.ball.atk_bonus
+            if self.ball.atk_bonus is not None
+            else random.randint(-settings.max_attack_bonus, settings.max_attack_bonus)
         )
-        bonus_health = self.ball.hp_bonus or random.randint(
-            -settings.max_health_bonus, settings.max_health_bonus
+        bonus_health = (
+            self.ball.hp_bonus
+            if self.ball.hp_bonus is not None
+            else random.randint(-settings.max_health_bonus, settings.max_health_bonus)
         )
 
         # check if we can spawn cards with a special background